#!/usr/bin/env python3
"""
Train SWITCH2-LLM
Incidents:
- VLAN mismatch
- Trunk negotiation issues
- Native VLAN mismatch
"""

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

BASE_MODEL = r"D:\dKorpesio\git_llm_wazuh\hermes\Hermes-3-Llama-3.1-8B"
DATASET    = "datasets/switch3_dataset_900.jsonl"
OUT_DIR    = "./switch_llm/lora_llm_switch3"

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

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

lora = 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(model, lora)

dataset = load_dataset("json", data_files=DATASET)["train"].train_test_split(
    test_size=0.1, seed=42
)

def format_sample(ex):
    cli = "\n".join(ex["cli_fix"])
    prompt = f"""
### Instruction:
{ex['instruction']}

Rules:
- Output ONLY valid Cisco IOS/IOS-XE CLI commands
- Do NOT include show/debug commands
- Do NOT provide explanation

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

### Wazuh alert:
{json.dumps(ex['wazuh_alert'], indent=2)}

### Devices:
{json.dumps(ex['devices'], indent=2)}

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

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

train_ds = dataset["train"].map(format_sample)
eval_ds  = dataset["test"].map(format_sample)

args = TrainingArguments(
    output_dir=OUT_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=args,
    train_dataset=train_ds,
    eval_dataset=eval_ds
)

if __name__ == "__main__":
    trainer.train()
    model.save_pretrained(OUT_DIR)
    print("✅ SWITCH3-LLM training finished")