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

# Lokálny model
model_name = r"D:\dKorpesio\git_llm_wazuh\hermes\Hermes-3-Llama-3.1-8B"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16
)

# LoRA konfigurácia
lora_config = 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_config)

# Dataset
dataset = load_dataset("json", data_files="network_context_4000.jsonl")["train"].train_test_split(test_size=0.1)
train_data, eval_data = dataset["train"], dataset["test"]

# Funkcia na formátovanie vstupov
def format_sample(example):
    # Extrahuj príkazy, verifikáciu a vysvetlenie z gold_actions
    try:
        actions = example["gold_actions"][0]
        commands = "\n".join(actions.get("commands", []))
        explain = actions.get("explain", "")
        verify = "\n".join(actions.get("verify", []))
        rollback = "\n".join(actions.get("rollback", []))
    except Exception:
        commands, explain, verify, rollback = "", "", "", ""
    
    # Výstupný text, ktorý sa model má naučiť generovať
    output = (
        f"Proposed configuration:\n{commands}\n\n"
        f"Verification steps:\n{verify}\n\n"
        f"Rollback plan:\n{rollback}\n\n"
        f"Explanation:\n{explain}"
    )
    
    # Prompt pre model
    prompt = (
        f"### Instruction:\n{example['instruction']}\n\n"
        f"### Device configuration context:\n{example['config_raw']}\n\n"
        f"### Wazuh alert:\n{json.dumps(example['wazuh_alert'], indent=2)}\n\n"
        f"### Response:\n{output}"
    )
    
    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)

# Tréningové argumenty
training_args = TrainingArguments(
    output_dir="./lora_hermes_cisco",
    num_train_epochs=3,
    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(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset
)

trainer.train()
model.save_pretrained("./lora_hermes_cisco")