"""
Export the Qwen2.5-0.5B email-triage LoRA adapter to llama.cpp-compatible GGUF.

VARIANT EXPERIMENT -- Qwen2.5 0.5B.
Companion to train/train_qwen2_5_0_5b_lora.py. Loads the base
Qwen2.5-0.5B-Instruct model with that variant's trained LoRA adapter, merges
the weights, and quantizes to the requested GGUF format(s). Outputs go into
the same clearly-labeled directory tree as the training run.

Outputs:
    outputs/qwen2.5-0.5b/gguf/grimoire-qwen2.5-0.5b-triage-q4_k_m.gguf
    outputs/qwen2.5-0.5b/gguf/grimoire-qwen2.5-0.5b-triage-q3_k_m.gguf  (optional)
    outputs/qwen2.5-0.5b/merged/                                       (optional)

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

import argparse
from pathlib import Path


def parse_args():
    parser = argparse.ArgumentParser(description="Export Qwen2.5-0.5B fine-tuned LoRA to GGUF")
    parser.add_argument("--model_name", default="grimoire-qwen2.5-0.5b-triage", help="Base name for GGUF/Ollama model")
    parser.add_argument("--base_model", default="Qwen/Qwen2.5-0.5B-Instruct", help="Base HF model name/path")
    parser.add_argument("--lora_dir", default="outputs/qwen2.5-0.5b/lora", help="Directory with LoRA adapter")
    parser.add_argument("--output_dir", default="outputs/qwen2.5-0.5b/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)
    parser.add_argument("--merged_dir", default="outputs/qwen2.5-0.5b/merged", help="Optional merged HF model output")
    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)

    # This Unsloth build's from_pretrained() doesn't accept adapter_name_or_path
    # (TypeError: ...__init__() got an unexpected keyword argument
    # 'adapter_name_or_path' -- confirmed against the 1.5B pipeline, see
    # DEPLOYMENT.md). Point model_name directly at the LoRA directory instead
    # (it has adapter_config.json with base_model_name_or_path set) --
    # Unsloth's own documented pattern, loads base+adapter as one call and is
    # correctly tagged as PEFT for save_pretrained_gguf().
    print(f"Loading base model + LoRA adapter from {args.lora_dir} ...")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.lora_dir,
        max_seq_length=args.max_seq_length,
        dtype=None,
        load_in_4bit=True,
    )

    # Export GGUF(s)
    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)")

    # Save merged HF model (useful for non-GGUF inference / debugging)
    if args.merged_dir:
        merged_dir = Path(args.merged_dir)
        merged_dir.mkdir(parents=True, exist_ok=True)
        print(f"Saving merged HF model to {merged_dir}")
        merged = model.merge_and_unload()
        merged.save_pretrained(merged_dir)
        tokenizer.save_pretrained(merged_dir)


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