#!/usr/bin/env python3
"""
Export un checkpoint LoRA en GGUF pour Ollama.

Stratégie : charge le modèle de base via Unsloth (pour bénéficier de
save_pretrained_gguf), puis applique le LoRA depuis le checkpoint
via PeftModel.from_pretrained en copiant la méthode GGUF.

Usage: python scripts/export_checkpoint_gguf.py [--checkpoint models/pmsi-lora-checkpoints/checkpoint-7000]
"""

import argparse
from pathlib import Path

BASE = Path(__file__).resolve().parent.parent
OUTPUT = BASE / "models"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", default=str(OUTPUT / "pmsi-lora-checkpoints" / "checkpoint-7000"))
    parser.add_argument("--model", default="unsloth/gemma-3-12b-it-bnb-4bit")
    parser.add_argument("--max-seq-length", type=int, default=512)
    parser.add_argument("--quant", default="q4_k_m")
    args = parser.parse_args()

    from unsloth import FastLanguageModel
    from peft import PeftModel

    checkpoint_dir = Path(args.checkpoint)
    gguf_dir = OUTPUT / "pmsi-gguf-v2"
    gguf_dir.mkdir(parents=True, exist_ok=True)

    # Étape 1 : Charger base via Unsloth (installe save_pretrained_gguf)
    print(f"[1/3] Chargement du modèle de base via Unsloth...")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.model,
        max_seq_length=args.max_seq_length,
        dtype=None,
        load_in_4bit=True,
    )

    # Sauver la méthode GGUF avant que PeftModel ne l'écrase
    save_gguf_fn = model.save_pretrained_gguf

    # Étape 2 : Appliquer le LoRA checkpoint
    print(f"[2/3] Application LoRA depuis {checkpoint_dir.name}...")
    model = PeftModel.from_pretrained(model, str(checkpoint_dir))

    # Réattacher la méthode GGUF d'Unsloth au modèle PeftModel
    model.save_pretrained_gguf = save_gguf_fn.__func__.__get__(model, type(model))

    # Vérifier que les poids LoRA sont bien chargés
    lora_params = [n for n, p in model.named_parameters() if "lora" in n and p.requires_grad]
    print(f"  Paramètres LoRA actifs : {len(lora_params)}")

    # Étape 3 : Export GGUF
    print(f"[3/3] Export GGUF ({args.quant})...")
    model.save_pretrained_gguf(
        str(gguf_dir),
        tokenizer,
        quantization_method=args.quant,
    )

    # Résultat
    gguf_files = sorted(gguf_dir.glob("*.gguf"), key=lambda f: f.stat().st_size)
    if not gguf_files:
        print("Aucun GGUF produit !")
        return

    final_gguf = gguf_files[0]
    for g in gguf_files:
        size_gb = g.stat().st_size / 1024**3
        print(f"  {g.name} ({size_gb:.1f} Go)")

    # Modelfile
    modelfile_path = gguf_dir / "Modelfile"
    with open(modelfile_path, "w") as f:
        f.write(f"FROM {final_gguf.name}\n\n")
        f.write("PARAMETER temperature 0.3\n")
        f.write("PARAMETER top_p 0.9\n")
        f.write("PARAMETER num_ctx 8192\n")

    print(f"\nTerminé !")
    print(f"  GGUF : {final_gguf}")
    print(f"\nPour importer dans Ollama :")
    print(f"  cd {gguf_dir}")
    print(f"  ollama create pmsi-coder -f Modelfile")


if __name__ == "__main__":
    main()