# ============================================================
# Gemma-3-270M – Analisi Traffico di Rete TCP/IP
# Unsloth + LoRA + Dataset JSONL (CIC-IDS2017 + UNSW-NB15)
# ============================================================
# Struttura basata sul tuo script, adattata per il dominio
# di analisi del traffico di rete con mappatura MITRE ATT&CK.
#
# PREREQUISITI:
#   Google Colab con runtime GPU (T4 basta)
#   !pip install --no-deps unsloth
#   !pip install transformers datasets trl peft accelerate sentencepiece
#
# FILE NECESSARI (nella stessa cartella dello script):
#   dataset.jsonl   ← generato dalla script apposita
# ============================================================

# ---------- INSTALL (Colab) ----------
# !pip install --no-deps unsloth
# !pip install transformers datasets trl peft accelerate sentencepiece

# ---------- IMPORT ----------
from unsloth import FastModel
from unsloth.chat_templates import get_chat_template, train_on_responses_only
import torch
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig

# ---------- CONFIG ----------
MODEL_NAME   = "unsloth/gemma-3-270m-it"
DATASET_PATH = "dataset_traffico.jsonl"          # <== il JSONL che abbiamo generato
OUTPUT_DIR   = "outputs"
MAX_SEQ_LENGTH = 512                    # 512 basta per questi prompt, risparmia memoria

# ---------- LOAD MODEL ----------
model, tokenizer = FastModel.from_pretrained(
    model_name     = MODEL_NAME,
    max_seq_length = MAX_SEQ_LENGTH,
    load_in_4bit   = False,
    load_in_8bit   = False,
    full_finetuning = False,
)

# ---------- LoRA ----------
# Configurazione più aggressiva sul rank (r=64) per un dominio specifico come questo.
# Target modules: tutti i proiettori del transformer.
model = FastModel.get_peft_model(
    model,
    r              = 64,
    target_modules = [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_alpha     = 64,
    lora_dropout   = 0,
    bias           = "none",
    use_gradient_checkpointing = "unsloth",
    random_state   = 3407,
)

# ---------- CHAT TEMPLATE (Gemma-3) ----------
tokenizer = get_chat_template(
    tokenizer,
    chat_template = "gemma3",
)

# ---------- LOAD DATASET ----------
dataset = load_dataset(
    "json",
    data_files = DATASET_PATH,
    split      = "train",
)

print(f"Dataset caricato: {len(dataset)} righe")
print(f"Campi presenti: {dataset.column_names}")
print(f"\nEsempio riga 0:")
print(dataset[0])

# ---------- CONVERT TO CHATML ----------
# Il JSONL ha campi: instruction, input, output
# Li convertiamo nel formato conversations [system, user, assistant]
# che Gemma-3 si aspetta.
def convert_to_chatml(example):
    system_prompt = example["instruction"]

    # Se c'è un campo 'context' lo aggiungiamo al system prompt
    if "context" in example and example["context"]:
        system_prompt += f"\nContesto: {example['context']}."

    return {
        "conversations": [
            {"role": "system",    "content": system_prompt},
            {"role": "user",      "content": example["input"]},
            {"role": "assistant", "content": example["output"]},
        ]
    }

dataset = dataset.map(convert_to_chatml)

# ---------- APPLY GEMMA-3 TEMPLATE ----------
# Applica il template di chat di Gemma-3 a ogni esempio.
# Questo produce la stringa finale che il modello vedrà durante il training.
def formatting_prompts_func(examples):
    convos = examples["conversations"]
    texts = [
        tokenizer.apply_chat_template(
            convo,
            tokenize            = False,
            add_generation_prompt = False,
        ).removeprefix("<bos>")
        for convo in convos
    ]
    return {"text": texts}

dataset = dataset.map(formatting_prompts_func, batched=True)

# Verifica come appare un prompt formattato
print("\n" + "=" * 60)
print(" PROMPT FORMATTATO (esempio)")
print("=" * 60)
print(dataset[0]["text"])
print("=" * 60)

# ---------- TRAINER ----------
trainer = SFTTrainer(
    model          = model,
    tokenizer      = tokenizer,
    train_dataset  = dataset,
    eval_dataset   = None,
    args = SFTConfig(
        dataset_text_field          = "text",
        per_device_train_batch_size = 4,
        gradient_accumulation_steps = 4,       # batch effettivo = 4 * 4 = 16
        warmup_steps                = 10,
        max_steps                   = 500,     # ~500 step su 10k righe con batch 16
        learning_rate               = 2e-5,
        logging_steps               = 25,
        optim                       = "adamw_8bit",
        weight_decay                = 0.001,
        lr_scheduler_type           = "linear",
        seed                        = 3407,
        output_dir                  = OUTPUT_DIR,
        report_to                   = "none",
    ),
)

# ---------- TRAIN ONLY ON ASSISTANT ----------
# Fondamentale: il modello calcola il loss SOLO sulla risposta dell'assistant,
# non sul prompt. Così non "impara" a ripetere la domanda.
trainer = train_on_responses_only(
    trainer,
    instruction_part = "<start_of_turn>user\n",
    response_part    = "<start_of_turn>model\n",
)

# ---------- TRAIN ----------
trainer.train()

# ---------- SAVE LoRA ----------
model.save_pretrained("gemma3-traffico-rete-lora")
tokenizer.save_pretrained("gemma3-traffico-rete-lora")
print("\n✓ Modello LoRA salvato in: gemma3-traffico-rete-lora/")
model.save_pretrained_merged(
    "gemma3-traffico-rete-lora",  # cartella output
    tokenizer,
    save_method="merged_16bit"       # Float16 per GGUF
)
model.save_pretrained_gguf(
        "gemma3-traffico-rete-lora",
        tokenizer,
        quantization_method = "BF16", # For now only Q8_0, BF16, F16 supported
    )
# ---------- INFERENCE: TEST ----------
# Dopo il training, prova il modello con alcuni flussi di esempio.
from transformers import TextStreamer

test_cases = [
    # Caso 1: profilo tipico DoS (masse enormi di byte src, pochissimi dst, durata minima)
    "Protocollo: tcp | Porta dst: 80 | Byte src: 480000 | Byte dst: 40 | Pacchetti: 5200 | Durata: 0.015s",
    # Caso 2: traffico normale HTTPS
    "Protocollo: tcp | Porta dst: 443 | Byte src: 1500 | Byte dst: 6200 | Pacchetti: 9 | Durata: 3.200s",
    # Caso 3: profilo PortScan (tanti dst diversi, pochi byte, durata quasi zero)
    "Protocollo: tcp | Porta dst: 22 | Byte src: 60 | Byte dst: 0 | Pacchetti: 1 | Durata: 0.002s",
    # Caso 4: profilo Brute Force su SSH
    "Protocollo: tcp | Porta dst: 22 | Byte src: 3200 | Byte dst: 8500 | Pacchetti: 45 | Durata: 1.800s",
    # Caso 5: profilo Infiltration / esfiltrazioni dati
    "Protocollo: tcp | Porta dst: 443 | Byte src: 8000 | Byte dst: 120000 | Pacchetti: 200 | Durata: 25.500s",
]

streamer = TextStreamer(tokenizer, skip_prompt=True)

for i, test_input in enumerate(test_cases, 1):
    messages = [
        {
            "role": "system",
            "content": (
                "Analizza il seguente flusso di traffico di rete TCP/IP. "
                "Classifica se è traffico normale o un attacco. "
                "Se è un attacco, indica la categoria e la tecnica MITRE ATT&CK corrispondente."
            ),
        },
        {"role": "user", "content": test_input},
    ]

    text = tokenizer.apply_chat_template(
        messages,
        tokenize            = False,
        add_generation_prompt = True,
    ).removeprefix("<bos>")

    print(f"\n{'─' * 60}")
    print(f" TEST {i}: {test_input[:80]}...")
    print(f"{'─' * 60}")
    print(" Risposta: ", end="")

    _ = model.generate(
        **tokenizer(text, return_tensors="pt").to("cuda"),
        max_new_tokens = 128,
        temperature    = 0.3,       # bassa temperatura = risposte più deterministe
        top_p          = 0.9,
        top_k          = 40,
        streamer       = streamer,
    )

# ---------- SAVE MERGED (opzionale) ----------
# Unisce i pesi LoRA al modello base e salva come modello completo.
# Utile per deployare senza dipendenza da PEFT.
#
model.save_pretrained_merged(
    "gemma3-traffico-rete-merged",
    tokenizer,
    save_method = "merged_16bit",
)
#
# ---------- SAVE GGUF (opzionale) ----------
# Formato GGUF per inferenza locale con llama.cpp / Ollama.
#
model.save_pretrained_gguf(
    "gemma3-traffico-rete-gguf",
    tokenizer,
    quantization_method = "Q8_0",   # Q8_0 = buon equilibrio qualità/dimensione
)