"""Eval Qwen3-4B + q3b_adapter on test/holdout."""
import json
import sys
from pathlib import Path
from mlx_lm import load, generate

BASE = "mlx-community/Qwen3-4B-Instruct-2507-4bit"
ADAPTER = sys.argv[1] if len(sys.argv) > 1 else "poc/llm-finetune/training/q3b_adapter_h1"
TEST = Path(sys.argv[2] if len(sys.argv) > 2 else "/tmp/test_100.jsonl")


import re


def parse(text):
    text = text.strip()
    # Strip Qwen3 hybrid-reasoning <think>...</think> block
    text = re.sub(r'<think>.*?</think>\s*', '', text, flags=re.DOTALL).strip()
    if text.startswith("```"):
        text = text.split("```")[1]
        if text.startswith("json"):
            text = text[4:]
    try:
        return json.loads(text.strip())
    except Exception:
        return None


def main():
    print(f"loading {BASE} + {ADAPTER}")
    model, tok = load(BASE, adapter_path=ADAPTER)
    rows = [json.loads(l) for l in TEST.read_text().splitlines() if l.strip()]
    print(f"test rows: {len(rows)}")

    schema_ok = intent_em = slots_em = 0
    for i, row in enumerate(rows):
        msgs = row["messages"]
        gold = json.loads(msgs[2]["content"])
        prompt = tok.apply_chat_template(msgs[:2], add_generation_prompt=True, tokenize=False)
        out = generate(model, tok, prompt=prompt, max_tokens=200, verbose=False)
        parsed = parse(out)
        is_schema = bool(parsed and "intent" in parsed and "slots" in parsed)
        is_intent = bool(parsed and parsed.get("intent") == gold.get("intent"))
        is_slots = bool(parsed and parsed.get("slots") == gold.get("slots"))
        schema_ok += int(is_schema); intent_em += int(is_intent); slots_em += int(is_slots)

    n = len(rows)
    print(f"\n=== Q3B RESULTS ===")
    print(f"adapter: {ADAPTER}")
    print(f"schema_valid: {schema_ok}/{n} = {schema_ok/n:.1%}")
    print(f"intent_em:    {intent_em}/{n} = {intent_em/n:.1%}")
    print(f"slots_em:     {slots_em}/{n} = {slots_em/n:.1%}")


if __name__ == "__main__":
    main()