"""Merge a LoRA adapter into its base model, export to GGUF, quantise, and try it once.

Purpose: turn the adapter directory that train-sft.py produced into a single model
         that every engine in Level 2 can serve: merge, save, convert to GGUF with
         llama.cpp's convert_hf_to_gguf.py, quantise with llama-quantize, and run one
         prompt through llama-cli so the export is proved rather than assumed.
Platform: spark, strix, nvidia (and mac for the merge and conversion steps; Track M's
          adapters come from mlx_lm.fuse instead, and the page says so)
Minimum memory: 8 GB
Assumes: torch, transformers and peft installed; a llama.cpp checkout for
         convert_hf_to_gguf.py and built binaries for llama-quantize and llama-cli,
         as built in Part 6; enough disk for the merged model plus two GGUF files.

Usage: python merge-and-export.py --adapter runs/sft-qwen3-0.6b --llama-cpp ~/llama.cpp
       python merge-and-export.py --adapter runs/sft-qwen3-0.6b --skip-gguf   # merge only
"""
from __future__ import annotations

import argparse
import shutil
import subprocess
import sys
from pathlib import Path

import torch
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer


def run(command: list[str]) -> None:
    """Run one external command, showing it first, and stop the script if it fails."""
    print("+ " + " ".join(str(part) for part in command))
    subprocess.run(command, check=True)


def find_tool(explicit: str | None, name: str) -> str:
    """Prefer the path the reader gave, then the one on PATH; fail with a useful message."""
    if explicit:
        return explicit
    found = shutil.which(name)
    if found is None:
        raise SystemExit(f"{name} is not on PATH; pass its path explicitly (built in Part 6)")
    return found


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--adapter", required=True, help="directory saved by train-sft.py")
    parser.add_argument("--merged-dir", default=None, help="where to write the merged model")
    parser.add_argument("--llama-cpp", default=None, help="llama.cpp checkout holding convert_hf_to_gguf.py")
    parser.add_argument("--llama-quantize", default=None, help="path to the llama-quantize binary")
    parser.add_argument("--llama-cli", default=None, help="path to the llama-cli binary")
    parser.add_argument("--quant", default="Q4_K_M", help="quantisation type accepted by llama-quantize")
    parser.add_argument("--outtype", default="bf16", choices=["f32", "f16", "bf16", "q8_0", "auto"],
                        help="conversion precision passed to convert_hf_to_gguf.py")
    parser.add_argument("--prompt", default="How much memory do the weights of a 5 billion parameter model need at BF16?")
    parser.add_argument("--predict", type=int, default=96, help="tokens to generate in the smoke test")
    parser.add_argument("--skip-gguf", action="store_true", help="merge and save only")
    parser.add_argument("--skip-run", action="store_true", help="do not run llama-cli at the end")
    args = parser.parse_args()

    adapter = Path(args.adapter)
    if not (adapter / "adapter_config.json").is_file():
        raise SystemExit(f"{adapter} does not look like a PEFT adapter directory (no adapter_config.json)")
    merged = Path(args.merged_dir) if args.merged_dir else adapter.with_name(adapter.name + "-merged")

    base_id = PeftConfig.from_pretrained(str(adapter)).base_model_name_or_path
    print(f"base model: {base_id}")
    print(f"adapter:    {adapter}")
    print(f"merged to:  {merged}")

    # The merge happens on the CPU in bfloat16: it is a weight arithmetic step, not a
    # forward pass, so it needs no accelerator and only one copy of the weights.
    base = AutoModelForCausalLM.from_pretrained(base_id, dtype=torch.bfloat16, device_map="cpu")
    model = PeftModel.from_pretrained(base, str(adapter))
    model = model.merge_and_unload()
    model.save_pretrained(str(merged))

    tokenizer = AutoTokenizer.from_pretrained(str(adapter) if (adapter / "tokenizer_config.json").is_file() else base_id)
    tokenizer.save_pretrained(str(merged))
    print(f"merged model written to {merged}")

    if args.skip_gguf:
        return

    if args.llama_cpp is None:
        raise SystemExit("pass --llama-cpp <path to your llama.cpp checkout>, or --skip-gguf")
    converter = Path(args.llama_cpp) / "convert_hf_to_gguf.py"
    if not converter.is_file():
        raise SystemExit(f"{converter} not found; --llama-cpp must point at a llama.cpp checkout")

    gguf_full = merged.with_suffix(".gguf")
    run([sys.executable, str(converter), str(merged), "--outfile", str(gguf_full), "--outtype", args.outtype])

    gguf_quant = merged.with_name(f"{merged.name}-{args.quant}.gguf")
    run([find_tool(args.llama_quantize, "llama-quantize"), str(gguf_full), str(gguf_quant), args.quant])

    for path in (gguf_full, gguf_quant):
        print(f"{path.name}: {path.stat().st_size:,} bytes")

    if args.skip_run:
        return
    run([
        find_tool(args.llama_cli, "llama-cli"),
        "--model", str(gguf_quant),
        "--prompt", args.prompt,
        "--predict", str(args.predict),
        "--temp", "0",
        "--seed", "0",
    ])


if __name__ == "__main__":
    main()