#!/usr/bin/env python3
"""
Convert the published HF PEFT LoRA adapter to GGUF for llama.cpp.

Enables stacking Well-Tuned + Llama Champion + Off the Grid badges:
  base GGUF + --lora plane-mode-study-coach-lora.gguf

Usage:
  python scripts/export_lora_gguf.py
  python scripts/export_lora_gguf.py --dry-run
  PMS_FINETUNED_MODEL=GuusBouwensNL/plane-mode-nemotron-4b-study-coach python scripts/export_lora_gguf.py
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "scripts"))

from hf_auth import resolve_hf_token  # noqa: E402

DEFAULT_ADAPTER = os.environ.get(
    "PMS_FINETUNED_MODEL",
    "GuusBouwensNL/plane-mode-nemotron-4b-study-coach",
)
DEFAULT_BASE = os.environ.get(
    "PMS_FT_BASE_MODEL",
    "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16",
)
DEFAULT_OUT = Path(
    os.environ.get(
        "PMS_LLAMACPP_LORA",
        str(ROOT / "models" / "nemotron" / "plane-mode-study-coach-lora.gguf"),
    )
)
DEFAULT_LLAMA_CPP = Path(os.environ.get("LLAMA_CPP_DIR", ROOT / "vendor" / "llama.cpp"))

# HF config.json for Nemotron 4B dense includes unused MoE keys; llama.cpp then
# tags LoRA as nemotron_h_moe while NVIDIA GGUF base is nemotron_h.
_MOE_KEYS = (
    "num_experts_per_tok",
    "n_routed_experts",
    "n_shared_experts",
    "moe_intermediate_size",
    "moe_shared_expert_intermediate_size",
    "moe_latent_size",
    "moe_shared_expert_overlap",
    "norm_topk_prob",
    "routed_scaling_factor",
    "n_group",
)


def _patch_dense_nemotron_hparams() -> None:
    """Strip spurious MoE keys so LoRA GGUF arch matches dense nemotron_h base."""
    import conversion.base as conv_base

    original = conv_base.ModelBase.load_hparams

    @classmethod
    def load_hparams(cls, dir_model, is_mistral_format):  # type: ignore[no-untyped-def]
        try:
            hparams = original.__func__(cls, dir_model, is_mistral_format)  # type: ignore[attr-defined]
        except AttributeError:
            hparams = original(dir_model, is_mistral_format)
        if hparams.get("num_experts_per_tok"):
            for key in _MOE_KEYS:
                hparams.pop(key, None)
        return hparams

    conv_base.ModelBase.load_hparams = load_hparams  # type: ignore[method-assign]


def _run_convert_lora(
    llama_cpp_dir: Path,
    adapter_dir: Path,
    base_model_id: str,
    outfile: Path,
) -> None:
    sys.path.insert(0, str(llama_cpp_dir))
    _patch_dense_nemotron_hparams()
    argv = [
        "convert_lora_to_gguf.py",
        str(adapter_dir),
        "--base-model-id",
        base_model_id,
        "--trust-remote-code",
        "--outfile",
        str(outfile),
        "--outtype",
        "f16",
    ]
    old_argv = sys.argv
    try:
        sys.argv = argv
        import runpy

        runpy.run_path(str(llama_cpp_dir / "convert_lora_to_gguf.py"), run_name="__main__")
    finally:
        sys.argv = old_argv


def _ensure_llama_cpp(llama_cpp_dir: Path) -> Path:
    convert = llama_cpp_dir / "convert_lora_to_gguf.py"
    if convert.exists():
        return llama_cpp_dir
    llama_cpp_dir.parent.mkdir(parents=True, exist_ok=True)
    print(f"Cloning llama.cpp into {llama_cpp_dir}...")
    subprocess.check_call(
        [
            "git",
            "clone",
            "--depth",
            "1",
            "https://github.com/ggml-org/llama.cpp",
            str(llama_cpp_dir),
        ]
    )
    return llama_cpp_dir


def _download_adapter(adapter_repo: str, dest: Path, token: str | None) -> Path:
    from huggingface_hub import snapshot_download

    dest.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        adapter_repo,
        local_dir=str(dest),
        token=token,
        ignore_patterns=["runs/*", "*.tfevents*"],
    )
    return dest


def export_lora_gguf(
    adapter_repo: str,
    base_model_id: str,
    outfile: Path,
    llama_cpp_dir: Path,
    token: str | None,
    dry_run: bool = False,
) -> Path:
    if dry_run:
        print("=== LoRA → GGUF dry run ===")
        print(f"  adapter:   {adapter_repo}")
        print(f"  base:      {base_model_id}")
        print(f"  output:    {outfile}")
        print(f"  llama.cpp: {llama_cpp_dir}")
        return outfile

    llama_cpp_dir = _ensure_llama_cpp(llama_cpp_dir)
    adapter_dir = outfile.parent / ".hf-lora-cache"
    outfile.parent.mkdir(parents=True, exist_ok=True)

    print(f"Downloading adapter {adapter_repo}...")
    _download_adapter(adapter_repo, adapter_dir, token)

    convert_py = llama_cpp_dir / "convert_lora_to_gguf.py"
    if not convert_py.exists():
        raise FileNotFoundError(f"Missing {convert_py}")
    print("Converting PEFT → GGUF LoRA (dense nemotron_h)...")
    _run_convert_lora(llama_cpp_dir, adapter_dir, base_model_id, outfile)
    print(f"Wrote {outfile} ({outfile.stat().st_size / 1e6:.1f} MB)")
    return outfile


def main() -> int:
    parser = argparse.ArgumentParser(description="Export fine-tuned LoRA to GGUF for llama.cpp")
    parser.add_argument("--adapter", default=DEFAULT_ADAPTER)
    parser.add_argument("--base-model-id", default=DEFAULT_BASE)
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--llama-cpp-dir", type=Path, default=DEFAULT_LLAMA_CPP)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    token = resolve_hf_token()
    export_lora_gguf(
        args.adapter,
        args.base_model_id,
        args.out,
        args.llama_cpp_dir,
        token,
        dry_run=args.dry_run,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())