"""
Export the fully fine-tuned Danube3-500M email-triage model to GGUF.

Companion to train/train_danube3_500m_full.py. No adapter to merge -- the
model is loaded directly and quantized.

Usage:
    python train/export_gguf_danube3_500m_full.py
    python train/export_gguf_danube3_500m_full.py --methods q4_k_m q3_k_m
"""

import argparse
from pathlib import Path


def parse_args():
    parser = argparse.ArgumentParser(description="Export full fine-tuned Danube3-500M to GGUF")
    parser.add_argument("--model_name", default="grimoire-danube3-500m-triage-full", help="Base name for GGUF/Ollama model")
    parser.add_argument("--model_dir", default="outputs/danube3-500m-full/model", help="Directory with fine-tuned model")
    parser.add_argument("--output_dir", default="outputs/danube3-500m-full/gguf", help="Where to write .gguf files")
    parser.add_argument(
        "--methods",
        nargs="+",
        default=["q4_k_m"],
        help="Quantization methods to produce (e.g. q4_k_m q3_k_m q2_k)",
    )
    parser.add_argument("--max_seq_length", type=int, default=2048)
    return parser.parse_args()


def main(args):
    from unsloth import FastLanguageModel

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"Loading full fine-tuned model from {args.model_dir} ...")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.model_dir,
        max_seq_length=args.max_seq_length,
        dtype=None,
        load_in_4bit=False,
    )

    for method in args.methods:
        print(f"Exporting GGUF with quantization={method} ...")
        model.save_pretrained_gguf(
            str(out_dir / args.model_name),
            tokenizer,
            quantization_method=method,
        )

    print("Done. Files:")
    for f in sorted(out_dir.glob("*.gguf")):
        print(f"  {f} ({f.stat().st_size / 1e6:.1f} MB)")


if __name__ == "__main__":
    args = parse_args()
    main(args)