#!/usr/bin/env python3
"""Evaluate the auditor model on held-out test examples, fully offline.

Measures, per task:
  - valid JSON         the reply parses as JSON
  - grounded           passes training/check_faithfulness.py (the same check the app uses)
  - fields match       finding: check_id, status and risk level equal the target;
                       summary: score and top gaps equal the target; ai: tools and approval flags equal the target
Run it for the base model and for the fine-tuned adapter to compare.

Usage:
    training/.venv/bin/python training/evaluate.py --count 30                       # base model
    training/.venv/bin/python training/evaluate.py --count 30 --adapter training/models/auditor-lora/final
"""
import argparse
import json
import os
import random
import sys
import time

os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

import torch  # noqa: E402
from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "training"))
from check_faithfulness import check as grounding_problems  # noqa: E402
from build_chat import check_chat  # noqa: E402

CITE = __import__("re").compile(r"\[(\d+)\]")
REFUSAL = __import__("re").compile(r"can only help|don't have reliable information|won't guess|can't see your device", __import__("re").I)


def chat_scores(example, reply):
    """Chat answers: grounded = citations exist and numbers appear in the sources; correct = cites the reference's
    passage, or (for questions the reference declines) declines without citations."""
    graded = {"messages": example["messages"][:-1] + [{"role": "assistant", "content": reply}]}
    grounded = not check_chat(graded)
    target = example["messages"][-1]["content"]
    want, got = set(CITE.findall(target)), set(CITE.findall(reply))
    correct = bool(want & got) if want else (not got and bool(REFUSAL.search(reply)))
    return grounded, correct


def pick(examples, count, seed):
    """Balanced sample: findings across statuses, plus summaries and AI classifications."""
    rng = random.Random(seed)
    groups = {}
    for item in examples:
        meta = item["meta"]
        groups.setdefault((meta["task"], meta.get("status", "")), []).append(item)
    for items in groups.values():
        rng.shuffle(items)
    chosen, keys = [], sorted(groups)
    while len(chosen) < count and any(groups[k] for k in keys):
        for key in keys:
            if groups[key] and len(chosen) < count:
                chosen.append(groups[key].pop())
    return chosen


def fields_match(task, answer, target):
    if task == "finding":
        return (answer.get("check_id") == target["check_id"] and answer.get("status") == target["status"]
                and (answer.get("risk") or {}).get("level") == target["risk"]["level"])
    if task == "summary":
        return (answer.get("score") == target["score"]
                and [g.get("check_id") for g in answer.get("top_gaps", [])] == [g["check_id"] for g in target["top_gaps"]])
    return ([(t.get("tool"), t.get("approved")) for t in answer.get("tools", [])]
            == [(t["tool"], t["approved"]) for t in target["tools"]])


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--base", default=os.path.join(ROOT, "training", "models", "Qwen2.5-0.5B-Instruct"))
    parser.add_argument("--adapter", help="LoRA adapter folder (omit to evaluate the base model)")
    parser.add_argument("--data", default=os.path.join(ROOT, "training", "out", "test.jsonl"))
    parser.add_argument("--count", type=int, default=30)
    parser.add_argument("--max-new-tokens", type=int, default=700)
    parser.add_argument("--seed", type=int, default=5)
    parser.add_argument("--save", help="write every reply to this JSONL file")
    parser.add_argument("--results", help="write the score table to this JSON file (for model cards)")
    args = parser.parse_args()

    torch.set_num_threads(os.cpu_count())
    tokenizer = AutoTokenizer.from_pretrained(args.base, local_files_only=True)
    model = AutoModelForCausalLM.from_pretrained(args.base, local_files_only=True, torch_dtype=torch.float32,
                                                 attn_implementation="eager")
    if args.adapter:
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, args.adapter)
        model = model.merge_and_unload()
    model.eval()

    with open(args.data, encoding="utf-8") as handle:
        examples = pick([json.loads(line) for line in handle], args.count, args.seed)

    totals = {}
    saved = []
    started = time.time()
    for number, example in enumerate(examples, 1):
        task = example["meta"]["task"]
        prompt_ids = tokenizer.apply_chat_template(example["messages"][:-1], add_generation_prompt=True, return_tensors="pt")
        with torch.no_grad():
            output = model.generate(prompt_ids, max_new_tokens=args.max_new_tokens, do_sample=False,
                                    pad_token_id=tokenizer.eos_token_id)
        reply = tokenizer.decode(output[0, prompt_ids.shape[1]:], skip_special_tokens=True).strip()
        row = totals.setdefault(task, {"n": 0, "json": 0, "grounded": 0, "fields": 0})
        row["n"] += 1
        if task == "chat":
            grounded, correct = chat_scores(example, reply)
            row["json"] += 1  # chat is plain text: counted as well-formed
            row["grounded"] += grounded
            row["fields"] += correct
            saved.append({"task": task, "meta": example["meta"], "reply": reply})
            print(f"[{number}/{len(examples)}] {task:<18} {time.time() - started:6.0f}s", flush=True)
            continue
        target = json.loads(example["messages"][-1]["content"])
        try:
            answer = json.loads(reply)
            row["json"] += 1
            graded = {"messages": example["messages"][:-1] + [{"role": "assistant", "content": reply}]}
            if not grounding_problems(graded):
                row["grounded"] += 1
            if fields_match(task, answer, target):
                row["fields"] += 1
        except json.JSONDecodeError:
            pass
        saved.append({"task": task, "meta": example["meta"], "reply": reply})
        print(f"[{number}/{len(examples)}] {task:<18} {time.time() - started:6.0f}s", flush=True)

    label = f"adapter {args.adapter}" if args.adapter else "base model"
    print(f"\nResults for {label} on {len(examples)} test examples:")
    print(f"{'task':<18} {'n':>3} {'valid JSON':>11} {'grounded':>9} {'fields match':>13}")
    for task, row in sorted(totals.items()):
        n = row["n"]
        print(f"{task:<18} {n:>3} {row['json'] / n:>10.0%} {row['grounded'] / n:>9.0%} {row['fields'] / n:>13.0%}")
    if args.results:
        with open(args.results, "w", encoding="utf-8") as handle:
            json.dump({"model": label, "data": os.path.relpath(args.data, ROOT), "count": len(examples), "seed": args.seed,
                       "tasks": {t: {"n": r["n"], "valid_json": r["json"] / r["n"], "grounded": r["grounded"] / r["n"],
                                     "fields_match": r["fields"] / r["n"]} for t, r in sorted(totals.items())},
                       "overall": {k: sum(r[c] for r in totals.values()) / max(1, sum(r["n"] for r in totals.values()))
                                   for k, c in (("valid_json", "json"), ("grounded", "grounded"), ("fields_match", "fields"))}},
                      handle, indent=2)
    if args.save:
        with open(args.save, "w", encoding="utf-8") as handle:
            for item in saved:
                handle.write(json.dumps(item, ensure_ascii=False) + "\n")


if __name__ == "__main__":
    main()