#!/usr/bin/env python3
"""
Training script for OSPF-1 Incident Solver LLM
Incidents:
 - ospf_neighbor_down
 - ospf_stuck_init_2way
 - ospf_stuck_exstart_exchange

Model learns to output ONLY CLI FIX COMMANDS.
"""

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/ospf2_dataset_fixed_900.jsonl"
OUTPUT_DIR = "./ospf_llm/lora_llm_ospf2"

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) 

# -----------------------------
# DATASET LOADING
# -----------------------------
dataset = load_dataset("json", data_files=DATA_FILE)["train"]

# -----------------------------
# FORMAT PROMPT FUNCTION
# -----------------------------
def format_sample(example):

    devices_section = "\n".join([
        f"- {d['name']} | iface {d['interface']} | mtu {d['mtu']} | state {d['ospf_state']} | rid {d['router_id']}"
        for d in example["devices"]
    ])

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

    prompt = f"""
### Instruction:
{example['instruction']}

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

### Involved devices:
{devices_section}

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

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

train_dataset = dataset.map(format_sample)

# -----------------------------
# TRAINING ARGUMENTS
# -----------------------------
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 = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset
)

# -----------------------------
# START TRAINING
# -----------------------------
trainer.train()
model.save_pretrained(OUTPUT_DIR)

print("\n✅ Training complete. Model saved to:", OUTPUT_DIR)