#!/usr/bin/env python3
"""Score a model on the held-out eval sets via in-process HF generate (no vLLM).

Useful when vLLM can't (yet) serve a model but the HF forward path works — e.g. the
Qwen3.6-27B base baseline. Emits the SAME report format as eval_endpoint.py
(reports/eval/<label>_eval.{json,md}) so numbers stay comparable across the base
model, in-training checkpoints, and any later vLLM-served runs.

Inference uses only the forward pass, so it is unaffected by the FLA gated-delta
*backward* kernel issue that blocks training on Hopper.

Example (on the GPU host, base model):
  python training/scripts/eval_hf_model.py --model Qwen/Qwen3.6-27B --label base \
    --eval data/eval/vuln_detection_test.jsonl --eval data/eval/knowledge_mcq.jsonl \
    --report-dir reports/eval --sample 100
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from intraining_eval import run_eval_sets  # noqa: E402


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--model", default="Qwen/Qwen3.6-27B")
    p.add_argument("--eval", action="append", required=True, help="Eval JSONL (repeatable).")
    p.add_argument("--label", required=True)
    p.add_argument("--report-dir", default="reports/eval")
    p.add_argument("--sample", type=int, default=100, help="Rows per set (default 100).")
    p.add_argument("--max-new-tokens", type=int, default=256)
    p.add_argument("--enable-thinking", action="store_true",
                   help="Let the model emit a <think> block first (slow; off by default for a fast tracking eval).")
    p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32", "auto"])
    p.add_argument("--device-map", default="auto")
    return p.parse_args()


def load_model(model_name: str, dtype: str, device_map: str):
    import torch
    import transformers

    dt = {"bfloat16": torch.bfloat16, "float16": torch.float16,
          "float32": torch.float32, "auto": "auto"}[dtype]
    errors = []
    for cls_name in ["AutoModelForMultimodalLM", "AutoModelForImageTextToText",
                     "AutoModelForVision2Seq", "AutoModelForCausalLM"]:
        cls = getattr(transformers, cls_name, None)
        if cls is None:
            continue
        try:
            model = cls.from_pretrained(model_name, dtype=dt, device_map=device_map, trust_remote_code=True)
            print(f"loaded with {cls_name}")
            return model
        except Exception as exc:  # noqa: BLE001
            errors.append(f"{cls_name}: {exc!r}")
    raise RuntimeError("could not load model:\n" + "\n".join(errors))


def main() -> int:
    args = parse_args()
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = load_model(args.model, args.dtype, args.device_map)

    sets = run_eval_sets(model, tokenizer, args.eval, sample_per_set=args.sample,
                         max_new_tokens=args.max_new_tokens, enable_thinking=args.enable_thinking)
    results = list(sets.values())

    report_dir = Path(args.report_dir)
    report_dir.mkdir(parents=True, exist_ok=True)
    payload = {"label": args.label, "model": args.model, "mock": False, "via": "hf_generate", "results": results}
    (report_dir / f"{args.label}_eval.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")

    lines = [f"# Endpoint Eval — {args.label} (HF generate)", "", f"- Model: `{args.model}`", ""]
    for r in results:
        if "error" in r:
            lines.append(f"## ERROR: {r}")
            continue
        lines.append(f"## {r.get('kind')}")
        lines.append(f"- n: {r.get('n')}  accuracy: **{r.get('accuracy', 0):.2%}**")
        if r.get("kind") == "vuln_detection":
            lines.append(f"- precision(vuln): {r.get('precision_vuln', 0):.2%}  "
                         f"recall(vuln): {r.get('recall_vuln', 0):.2%}  f1: {r.get('f1_vuln', 0):.2%}  "
                         f"unparsed: {r.get('unparsed')}")
        else:
            lines.append(f"- unparsed: {r.get('unparsed')}")
        lines.append("")
    (report_dir / f"{args.label}_eval.md").write_text("\n".join(lines) + "\n", encoding="utf-8")

    print(json.dumps(payload, indent=2))
    return 0


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