"""
finetune.py
===========
Fine-tuning con LoRA/QLoRA di un LLM su dataset.jsonl.
Usa TRL SFTTrainer — funziona su GPU consumer (4GB+) o su Google Colab gratuito.

Modello base consigliato (leggero, multilingue, italiano):
    Qwen/Qwen2.5-0.5B-Instruct   ← 0.5B param, ottimo su CPU/GPU piccola
    Qwen/Qwen2.5-1.5B-Instruct   ← 1.5B param, qualità migliore
    google/gemma-2-2b-it          ← 2B param, ottimo in italiano

Uso:
    # Training completo
    python finetune.py --dataset dataset.jsonl --hf-repo TUO_USERNAME/generai-model

    # Solo dry-run (verifica che funzioni)
    python finetune.py --dataset dataset.jsonl --dry-run
"""

import argparse
import json
import os
import sys
from pathlib import Path

# ── Dipendenze opzionali — errore chiaro se mancano ───────────────────────────
def _check_deps():
    missing = []
    for pkg in ["transformers", "trl", "peft", "datasets", "torch"]:
        try:
            __import__(pkg)
        except ImportError:
            missing.append(pkg)
    if missing:
        print("❌ Dipendenze mancanti per il training:")
        print(f"   pip install {' '.join(missing)}")
        print("   (oppure: pip install -r requirements-train.txt)")
        sys.exit(1)

_check_deps()

import torch
from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, TaskType
from trl import SFTTrainer, SFTConfig


DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
OUTPUT_DIR    = "./generai-finetuned"


def load_dataset_jsonl(path: str) -> Dataset:
    rows = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    if not rows:
        print(f"❌ Dataset vuoto: {path}")
        sys.exit(1)
    print(f"[dataset] {len(rows)} esempi caricati da {path}")
    return Dataset.from_list(rows)


def apply_chat_template(examples, tokenizer):
    """Applica il chat template del modello ai messaggi."""
    texts = []
    for msgs in examples["messages"]:
        text = tokenizer.apply_chat_template(
            msgs,
            tokenize=False,
            add_generation_prompt=False,
        )
        texts.append(text)
    return {"text": texts}


def build_model_and_tokenizer(model_name: str, use_4bit: bool):
    print(f"[model] Caricamento: {model_name}")

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    bnb_config = None
    if use_4bit:
        try:
            import bitsandbytes  # noqa
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
                bnb_4bit_compute_dtype=torch.bfloat16,
            )
            print("[model] Quantizzazione 4-bit attivata (QLoRA)")
        except ImportError:
            print("[model] bitsandbytes non disponibile — uso fp16 standard")
            use_4bit = False

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        quantization_config=bnb_config,
        torch_dtype=torch.bfloat16 if not use_4bit else None,
        device_map="auto",
        trust_remote_code=True,
    )
    return model, tokenizer


def build_lora_config() -> LoraConfig:
    return LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=16,               # rank — più alto = più parametri addestrati
        lora_alpha=32,
        lora_dropout=0.05,
        bias="none",
        target_modules="all-linear",
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset",   required=True,          help="Path al file dataset.jsonl")
    parser.add_argument("--model",     default=DEFAULT_MODEL,  help=f"Modello base HuggingFace (default: {DEFAULT_MODEL})")
    parser.add_argument("--hf-repo",   default="",             help="Username/nome-repo su HuggingFace per upload automatico")
    parser.add_argument("--epochs",    type=int,   default=3,  help="Numero di epoche (default: 3)")
    parser.add_argument("--batch",     type=int,   default=2,  help="Batch size (default: 2)")
    parser.add_argument("--lr",        type=float, default=2e-4, help="Learning rate (default: 2e-4)")
    parser.add_argument("--max-len",   type=int,   default=512, help="Lunghezza massima token (default: 512)")
    parser.add_argument("--4bit",      dest="use_4bit", action="store_true", help="Usa QLoRA 4-bit (serve bitsandbytes)")
    parser.add_argument("--dry-run",   action="store_true",    help="Verifica configurazione senza addestrare")
    args = parser.parse_args()

    if not Path(args.dataset).exists():
        print(f"❌ File non trovato: {args.dataset}")
        print("   Esegui prima: python export_dataset.py")
        sys.exit(1)

    # ── Carica dataset ─────────────────────────────────────────────────────────
    raw_ds = load_dataset_jsonl(args.dataset)

    # ── Carica modello ─────────────────────────────────────────────────────────
    model, tokenizer = build_model_and_tokenizer(args.model, args.use_4bit)

    # ── Applica chat template ──────────────────────────────────────────────────
    ds = raw_ds.map(
        lambda x: apply_chat_template(x, tokenizer),
        batched=True,
        remove_columns=raw_ds.column_names,
    )
    print(f"[dataset] Esempio tokenizzato:\n{ds[0]['text'][:300]}...\n")

    if args.dry_run:
        print("✅ Dry-run completato. Configurazione corretta.")
        print(f"   Modello   : {args.model}")
        print(f"   Dataset   : {len(ds)} esempi")
        print(f"   Output    : {OUTPUT_DIR}")
        return

    # ── LoRA ───────────────────────────────────────────────────────────────────
    lora_cfg = build_lora_config()
    model = get_peft_model(model, lora_cfg)
    model.print_trainable_parameters()

    # ── Training config ────────────────────────────────────────────────────────
    sft_cfg = SFTConfig(
        output_dir=OUTPUT_DIR,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch,
        gradient_accumulation_steps=4,
        learning_rate=args.lr,
        lr_scheduler_type="cosine",
        warmup_ratio=0.1,
        logging_steps=10,
        save_steps=50,
        save_total_limit=2,
        bf16=torch.cuda.is_available(),
        fp16=False,
        max_seq_length=args.max_len,
        dataset_text_field="text",
        report_to="none",
        push_to_hub=bool(args.hf_repo),
        hub_model_id=args.hf_repo or None,
    )

    trainer = SFTTrainer(
        model=model,
        train_dataset=ds,
        args=sft_cfg,
        tokenizer=tokenizer,
    )

    print("\n🚀 Avvio training...\n")
    trainer.train()

    # ── Salva ──────────────────────────────────────────────────────────────────
    trainer.save_model(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)
    print(f"\n✅ Modello salvato in: {OUTPUT_DIR}")

    if args.hf_repo:
        print(f"⬆️  Upload su HuggingFace: https://huggingface.co/{args.hf_repo}")
        trainer.push_to_hub()
    else:
        print("\nPer caricare su HuggingFace:")
        print(f"  huggingface-cli upload TUO_USERNAME/generai-model {OUTPUT_DIR}")
        print("\nPer usarlo in GenerAI:")
        print(f"  Imposta HF_MODEL={args.hf_repo or 'TUO_USERNAME/generai-model'} in .env")


if __name__ == "__main__":
    main()