#!/usr/bin/env python3
"""
Training script for OSPF-4 Incident Solver LLM

Incidents:
 - ospf_lsa_flood
 - ospf_lsdb_inconsistency
 - ospf_redistribution_issue

Model learns to output ONLY CLI FIX COMMANDS.
AREA has been removed from devices in OSPF4 dataset.
"""

from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
import torch
import json

# ==========================
# MODEL & DATA CONFIG
# ==========================

BASE_MODEL = r"D:\dKorpesio\git_llm_wazuh\hermes\Hermes-3-Llama-3.1-8B"
DATA_FILE = "datasets/ospf4_dataset_fixed_1500.jsonl"
OUTPUT_DIR = "./ospf_llm/lora_llm_ospf4"

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="auto"
)

# ==========================
# LoRA CONFIG
# ==========================

lora_cfg = LoraConfig(
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    target_modules=["q_proj", "v_proj"],
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(base_model, lora_cfg)

# ==========================
# LOAD DATASET
# ==========================

dataset = load_dataset("json", data_files=DATA_FILE)["train"]

def format_sample(example):
    """
    Trénovací prompt pre OSPF4

    ### Instruction:
    ...

    ### Incident type:
    ospf_lsa_flood

    ### Involved devices:
    - CORE-R1 | iface Gi0/3 | lsa_rate 350 | checksum False seq False | proto bgp | missing_subnets True leak False
    - ACCESS-SW1 | ...

    ### Response (ONLY CLI FIX COMMANDS):
    router ospf 1
    timers throttle lsa all 20 200 5000
    clear ip ospf process
    """

    devices_section = "\n".join([
        (
            f"- {d['name']} | iface {d['interface']} "
            f"| lsa_rate {d['lsa_rate']} "
            f"| checksum {d['checksum_mismatch']} seq {d['seq_mismatch']} "
            f"| proto {d['redistribution_protocol']} "
            f"| missing_subnets {d['missing_subnets']} "
            f"| route_leak {d['route_leak']}"
        )
        for d in example["devices"]
    ])

    cli_fixes = "\n".join(example["cli_fix"])

    prompt = f"""
### Instruction:
{example['instruction']}
Do NOT provide explanation.
Always modify the device specified in the Wazuh alert.
You MUST always include "clear ip ospf process" as the final command when applicable.

### Incident type:
{example['incident_type']}

### Involved devices:
{devices_section}

### Response (ONLY CLI FIX COMMANDS):
{cli_fixes}
""".strip()

    tokens = tokenizer(
        prompt,
        truncation=True,
        max_length=768,
        padding="max_length"
    )
    tokens["labels"] = tokens["input_ids"].copy()
    return tokens

train_dataset = dataset.map(format_sample)

# ==========================
# TRAINING ARGS
# ==========================

training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=3,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    fp16=True,
    logging_steps=20,
    save_strategy="epoch",
    save_total_limit=2,
    report_to="none"
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset
)

# ==========================
# TRAIN
# ==========================

if __name__ == "__main__":
    trainer.train()
    model.save_pretrained(OUTPUT_DIR)
    print("\n✅ Training complete. Model saved to:", OUTPUT_DIR)