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

Incidents:
 - ospf_auth_mismatch
 - ospf_hello_dead_mismatch
 - ospf_network_type_mismatch

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

# ==========================
# KONFIGURÁCIA MODELU A DÁT
# ==========================

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

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

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

# ==========================
# LoRA NASTAVENIE
# ==========================

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)

# ==========================
# NAČÍTANIE DATASETU
# ==========================

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

def format_sample(example):
    """
    Pripraví prompt v štýle:

    ### Instruction:
    ...

    ### Incident type:
    ospf_auth_mismatch

    ### Involved devices:
    - CORE-R1 | iface Gi0/1 | auth md5 | hello 10 | dead 40 | net broadcast
    - ACCESS-SW1 | ...

    ### Response (ONLY CLI FIX COMMANDS):
    interface Gi0/1
    ip ospf authentication md5
    clear ip ospf process
    """

    devices_section = "\n".join([
        (
            f"- {d['name']} | iface {d['interface']} | "
            f"auth {d['authentication']} | hello {d['hello']} | "
            f"dead {d['dead']} | net {d['network_type']}"
        )
        for d in example["devices"]
    ])

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

    prompt = f"""
### Instruction:
{example['instruction']}
You MUST always include "clear ip ospf process" as the final command.

### 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)

# ==========================
# TRÉNINGOVÉ PARAMETRE
# ==========================

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
)

# ==========================
# ŠTART TRÉNINGU
# ==========================

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