#!/usr/bin/env python3
"""
Training script for BGP-1 Incident Solver LLM

Incidents:
 - bgp_session_flap
 - bgp_hold_timer_expiration
 - bgp_neighborship_reset

Output: ONLY CLI FIX COMMANDS (no explanation).
"""

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/bgp1_dataset_v3_900.jsonl"
OUTPUT_DIR = "./bgp_llm/lora_llm_bgp1"

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):
    """
    Prompt format aligned with dataset v3
    """

    devices_section = "\n".join([
        (
            f"- {d['name']} | AS {d['local_as']} "
            f"| neighbor {d['neighbor_ip']} remote-as {d['remote_as']} "
            f"| iface {d['interface']} "
            f"| peer KA {d['peer_keepalive']} hold {d['peer_hold']}"
        )
        for d in example["devices"]
    ])

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

    prompt = f"""
### Instruction:
{example['instruction']}
Rules:
- NEVER set BGP timers to 10 30
- For neighborship reset, do NOT modify configuration
- For hold timer expiration, timers MUST match the peer
- Prefer minimal-impact actions
Do NOT provide explanation.

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

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

### Devices:
{devices_section}

### Response (CLI FIX COMMANDS ONLY):
{cli_fix}
""".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=8,
    learning_rate=2e-4,
    fp16=True,
    logging_steps=20,
    save_strategy="epoch",
    save_total_limit=2,
    report_to="none",
    max_grad_norm=0.3
)

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

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