#!/usr/bin/env python3
"""
Train SEC1-LLM (ACL security incidents)
Incidents:
- ACL blocking legitimate traffic
- ACL misconfiguration
- Excessive deny entries
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, json

BASE_MODEL = r"D:\dKorpesio\git_llm_wazuh\hermes\Hermes-3-Llama-3.1-8B"
DATASET    = "datasets/sec3_dataset_v2_900.jsonl"
OUT_DIR    = "./sec_llm/lora_llm_sec3"

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

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

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

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"]}

You are a network security automation model.
Based on the incident context and decision constraints,
generate appropriate Cisco IOS configuration commands.

Output ONLY CLI FIX COMMANDS.
Do not explain.


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

### Wazuh alert:
{json.dumps(ex["wazuh_alert"], 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)

training_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=training_args,
    train_dataset=train_ds,
    eval_dataset=eval_ds
)

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