#!/usr/bin/env python3
"""
Training script for LLM-S (Show-Command Selector)
Maps Wazuh alerts → required diagnostic show commands.

Model: Hermes-3-Llama-3.1-8B
Training: LoRA
"""

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

# -----------------------------
# PATHS
# -----------------------------
BASE_MODEL = r"D:\dKorpesio\git_llm_wazuh\hermes\Hermes-3-Llama-3.1-8B"
DATASET = "llm_s_dataset_v5.jsonl"
OUTPUT = "./lora_llm_s"

# -----------------------------
# LOAD MODEL + TOKENIZER
# -----------------------------
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
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.05,
    bias="none",
    target_modules=["q_proj", "v_proj"],
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, lora_cfg)

# -----------------------------
# DATASET PREPROCESSING
# -----------------------------
ds = load_dataset("json", data_files=DATASET)["train"].train_test_split(test_size=0.1)
train_data, eval_data = ds["train"], ds["test"]


def format_sample(example):
    """
    Convert dataset entry into a full prompt for training.
    """

    instruction = example["instruction"]
    wazuh_alert = json.dumps(example["wazuh_alert"], indent=2)
    requested_show = "\n".join(example["requested_show"])

    output_block = f"Requested show commands:\n{requested_show}"

    prompt = (
        f"### Instruction:\n"
        f"{instruction}\n\n"
        f"### Wazuh alert:\n"
        f"{wazuh_alert}\n\n"
        f"### Response:\n"
        f"{output_block}"
    )

    tokens = tokenizer(
        prompt,
        truncation=True,
        max_length=1024,
        padding="max_length"
    )

    tokens["labels"] = tokens["input_ids"].copy()
    return tokens


train_dataset = train_data.map(format_sample)
eval_dataset = eval_data.map(format_sample)

# -----------------------------
# TRAINING ARGS
# -----------------------------
training_args = TrainingArguments(
    output_dir=OUTPUT,
    num_train_epochs=2,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    fp16=True,
    logging_steps=20,
    save_total_limit=2,
    save_strategy="epoch",
    report_to="none"
)

# -----------------------------
# TRAINER
# -----------------------------
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset
)

# -----------------------------
# RUN TRAINING
# -----------------------------
trainer.train()

model.save_pretrained(OUTPUT)
tokenizer.save_pretrained(OUTPUT)

print("✅ Training complete — LoRA adapter saved to:", OUTPUT)