#!/usr/bin/env python3
"""Verify the published LoRA adapter loads and can generate (needs GPU for full test)."""

from __future__ import annotations

import argparse
import json
import os
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",
)


def verify_hub_files(adapter_repo: str, token: str | None) -> dict:
    from huggingface_hub import HfApi

    api = HfApi(token=token)
    info = api.model_info(adapter_repo)
    names = {s.rfilename for s in info.siblings}
    required = {"adapter_config.json", "adapter_model.safetensors"}
    missing = sorted(required - names)
    if missing:
        raise FileNotFoundError(f"Missing on Hub: {missing}")

    cfg = json.loads(
        Path(
            api.hf_hub_download(adapter_repo, "adapter_config.json", token=token)
        ).read_text(encoding="utf-8")
    )
    base = cfg.get("base_model_name_or_path", DEFAULT_BASE)
    return {
        "adapter_repo": adapter_repo,
        "base_model": base,
        "peft_type": cfg.get("peft_type"),
        "target_modules": cfg.get("target_modules"),
        "files": sorted(names),
    }


def smoke_generate(adapter_repo: str, base_model: str, token: str | None) -> str:
    import torch
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA required for smoke generation test")

    tokenizer = AutoTokenizer.from_pretrained(
        adapter_repo, token=token, trust_remote_code=True
    )
    base = AutoModelForCausalLM.from_pretrained(
        base_model,
        token=token,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    model = PeftModel.from_pretrained(base, adapter_repo, token=token)
    model.eval()

    messages = [
        {
            "role": "system",
            "content": "You are Plane Mode Scholar, a concise study coach.",
        },
        {"role": "user", "content": "In one sentence, what is gradient descent?"},
    ]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
    )
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=64, do_sample=False)
    reply = tokenizer.decode(out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True)
    return reply.strip()


def main() -> int:
    parser = argparse.ArgumentParser(description="Verify fine-tuned Nemotron adapter")
    parser.add_argument("--adapter", default=DEFAULT_ADAPTER)
    parser.add_argument("--base", default=None, help="Override base model")
    parser.add_argument("--smoke", action="store_true", help="Run GPU generation test")
    args = parser.parse_args()

    token = resolve_hf_token()
    print(f"Adapter: {args.adapter}")
    meta = verify_hub_files(args.adapter, token)
    base = args.base or meta["base_model"]
    print(f"Base:    {base}")
    print(f"PEFT:    {meta['peft_type']} r={meta.get('target_modules')}")
    print("Hub files OK")

    if not args.smoke:
        print("Skipping generation (pass --smoke on a CUDA machine)")
        return 0

    reply = smoke_generate(args.adapter, base, token)
    print(f"Sample reply: {reply[:300]}")
    return 0 if reply else 1


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