"""Evaluate persona and route contracts after Qwen3 SFT.

Reports syntax/shape validity, persona mood accuracy, route intent sequence
exact match, per-action intent accuracy, and slot exact match.
"""
from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

os.environ.setdefault("HF_HOME", "/mnt/D/.cache/huggingface")
os.environ.setdefault("HF_DATASETS_CACHE", "/mnt/D/.cache/huggingface/datasets")

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

HERE = Path(__file__).resolve().parent


def rows(path: str, limit: int) -> list[dict]:
    out = [json.loads(line) for line in Path(path).read_text(encoding="utf-8").splitlines()
           if line.strip()]
    return out[:limit] if limit else out


def normalized_route(value):
    if isinstance(value, dict):
        value = [value]
    return value if isinstance(value, list) else None


@torch.inference_mode()
def complete(tok, model, messages, max_new_tokens):
    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
                                     enable_thinking=False)
    ids = tok(prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(model.device)
    out = model.generate(input_ids=ids, max_new_tokens=max_new_tokens, do_sample=False,
                         pad_token_id=tok.eos_token_id)
    return tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True).strip()


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default=str(HERE / "Qwen3-1.7B-ru-cpt"))
    ap.add_argument("--adapter", default=str(HERE / "Qwen3-1.7B-maven-sft"))
    ap.add_argument("--persona-eval", default=str(HERE / "data/persona_eval.jsonl"))
    ap.add_argument("--route-eval", default=str(HERE / "data/route_eval.jsonl"))
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--output")
    args = ap.parse_args()
    tok = AutoTokenizer.from_pretrained(args.adapter)
    base = AutoModelForCausalLM.from_pretrained(args.base, torch_dtype=torch.bfloat16,
                                                attn_implementation="eager", device_map={"":0})
    model = PeftModel.from_pretrained(base, args.adapter).eval()

    persona_valid = persona_mood = 0
    persona_rows = rows(args.persona_eval, args.limit)
    for row in persona_rows:
        expected = json.loads(row["messages"][-1]["content"])
        raw = complete(tok, model, row["messages"][:-1], 160)
        try:
            got = json.loads(raw)
            valid = set(got) == {"response", "mood"} and isinstance(got["response"], str)
        except Exception:
            valid, got = False, {}
        persona_valid += valid
        persona_mood += valid and got["mood"] == expected["mood"]

    route_valid = route_seq = action_total = intent_ok = slots_ok = 0
    route_rows = rows(args.route_eval, args.limit)
    for row in route_rows:
        expected = normalized_route(json.loads(row["messages"][-1]["content"]))
        raw = complete(tok, model, row["messages"][:-1], 160)
        try:
            got = normalized_route(json.loads(raw))
            valid = bool(got) and all(isinstance(x, dict) and "intent" in x for x in got)
        except Exception:
            valid, got = False, None
        route_valid += valid
        if not valid:
            continue
        route_seq += [x["intent"] for x in got] == [x["intent"] for x in expected]
        for exp, actual in zip(expected, got):
            action_total += 1
            intent_ok += actual.get("intent") == exp.get("intent")
            slots_ok += all(actual.get(k) == v for k, v in exp.items() if k != "intent")

    report = {
        "persona": {"cases":len(persona_rows),
                    "json_valid_pct":100*persona_valid/len(persona_rows),
                    "mood_accuracy_pct":100*persona_mood/len(persona_rows)},
        "route": {"cases":len(route_rows), "json_valid_pct":100*route_valid/len(route_rows),
                  "intent_sequence_exact_pct":100*route_seq/len(route_rows),
                  "action_intent_accuracy_pct":100*intent_ok/action_total if action_total else 0,
                  "slot_exact_pct":100*slots_ok/action_total if action_total else 0},
    }
    rendered = json.dumps(report, indent=2, ensure_ascii=False) + "\n"
    if args.output:
        Path(args.output).write_text(rendered, encoding="utf-8")
    print(rendered, end="")


if __name__ == "__main__":
    main()