"""
Export a trained LoRA adapter to llama.cpp-compatible GGUF.

This loads the base Qwen2.5-1.5B-Instruct model with the trained LoRA adapter,
merges the weights, and quantizes to the requested GGUF format(s).

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

Usage:
    python train/export_gguf.py
    python train/export_gguf.py --methods q4_k_m q3_k_m
    python train/export_gguf.py --lora_dir outputs/lora --base_model Qwen/Qwen2.5-1.5B-Instruct
"""

import argparse
from pathlib import Path


def parse_args():
    parser = argparse.ArgumentParser(description="Export fine-tuned LoRA to GGUF")
    parser.add_argument("--base_model", default="Qwen/Qwen2.5-1.5B-Instruct", help="Base HF model name/path")
    parser.add_argument("--lora_dir", default="outputs/lora", help="Directory with LoRA adapter")
    parser.add_argument("--output_dir", default="outputs/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/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: Qwen2ForCausalLM.__init__() got an unexpected keyword argument
    # 'adapter_name_or_path'). Loading the base model separately and attaching
    # via plain peft.PeftModel.from_pretrained *works* for inference, but
    # save_pretrained_gguf() doesn't recognize a plain PeftModel as PEFT
    # ("Model is not a PEFT model. Saving directly without LoRA merge...") and
    # then fails on an unrelated weight-conversion bug trying to save it as if
    # it were a full model. Pointing model_name directly at the LoRA directory
    # (which has adapter_config.json with base_model_name_or_path set) is
    # Unsloth's own documented pattern for this and loads base+adapter as a
    # single call, correctly tagged as PEFT.
    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)
    model_name = "grimoire-qwen2.5-1.5b-triage"
    for method in args.methods:
        print(f"Exporting GGUF with quantization={method} ...")
        model.save_pretrained_gguf(
            str(out_dir / 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)