#!/usr/bin/env python3
"""Run the self-eval task set against models and write ``proficiency``.

Why this exists: ``proficiency_score`` is the ONLY category-dependent term in
the composite, so until this table has data, ``task_category`` cannot change a
routing decision at all — the classifier's category output is computed, paid
for, and then discarded. This is what makes it matter.

Scoring is objective wherever the category admits it (see evals/tasks.yaml):
code is executed against checks, exact answers are compared, tool calls are
inspected structurally. Only the four prose categories fall back to a judge.

Scores accumulate rather than replace (``proficiency_store.add_self_eval``),
so running this repeatedly tightens estimates and pushes categories past
``self_eval_min_samples`` into a proper blend with the leaderboard prior.

    python eval_proficiency.py --dry-run          # plan only, no calls
    python eval_proficiency.py                    # every identity, every task
    python eval_proficiency.py --models kimi-k3 --categories coding_general

SAFETY: `code` tasks execute model-generated Python. Isolation is a
subprocess with a wall-clock timeout, running in a temp directory — not a
container and not a real sandbox. The task prompts ask for small pure
functions, so nothing invites filesystem or network use, but treat this as
"bounded", not "safe against a hostile model".
"""

from __future__ import annotations

import argparse
import json
import os
import re
import secrets
import sqlite3
import subprocess
import sys
import tempfile
from collections import defaultdict
from pathlib import Path
from typing import Optional

import requests
import yaml

from config import RouterConfig, load_config
from proficiency_store import add_self_eval, propagate_to_variants

TASKS_PATH = "evals/tasks.yaml"
CODE_TIMEOUT_SECONDS = 15
CALL_TIMEOUT_SECONDS = 300

# Generous on purpose. These are reasoning models and the trace is billed
# against the same budget as the answer: at 1200, qwen3.6-35b spent ~4,200
# characters thinking and returned an EMPTY content field, scoring 0.00 on
# tasks it can plainly do. The cap must clear the trace, not just the answer,
# and the harder task set makes traces longer still.
#
# Clamped per model at request time: this exceeds gemma-4-31b's advertised
# 16384 output limit, and asking for more than a model allows is an error
# rather than a silent truncation.
EVAL_MAX_TOKENS = 24000

# Preamble for the sandbox: common imports so ordinary solutions run, plus the
# `raises` helper the checks use. Defined AFTER the model's code so a model
# that shadows one of these does not break the harness.
HARNESS_PREAMBLE = "import collections, itertools, json, math, re, string, time\n"
HARNESS_HELPERS = '''
def raises(exc, fn, *a, **kw):
    try:
        fn(*a, **kw)
    except exc:
        return True
    except Exception:
        return False
    return False
'''

FENCE_RE = re.compile(r"^\s*```[a-zA-Z]*\n(.*?)```", re.DOTALL | re.MULTILINE)


# --- extraction and scoring ----------------------------------------------

def strip_fences(text: str) -> str:
    """Pull code out of a markdown block if the model added one anyway.

    Every code prompt says "no markdown fences" and models add them regularly,
    so this is normalization rather than leniency — a model that solved the
    task shouldn't score zero for formatting.
    """
    match = FENCE_RE.search(text or "")
    code = match.group(1) if match else (text or "")
    # Strip the WHOLE reply's leading/trailing whitespace: a single leading
    # space turns an otherwise perfect function into IndentationError, which
    # is how kimi-k2.7-code scored 0.00 on every coding task in an early run.
    # Internal indentation is untouched.
    return code.strip()


def parse_check_results(stdout: str, marker: str) -> dict[int, bool]:
    """Read the harness's own verdicts out of a script's stdout.

    Keyed by check INDEX rather than counted, which bounds the total at the
    number of checks even if a line somehow repeats, and drops anything that
    is not exactly one of the harness's own three-token lines.
    """
    results: dict[int, bool] = {}
    for line in stdout.splitlines():
        parts = line.split()
        if len(parts) != 3 or parts[0] != marker or parts[2] not in ("PASS", "FAIL"):
            continue
        try:
            index = int(parts[1])
        except ValueError:
            continue
        results[index] = parts[2] == "PASS"
    return results


def score_code(model_output: str, checks: list[str]) -> tuple[float, str]:
    """Execute the model's code and eval each check against it.

    Returns (fraction of checks passing, detail). Partial credit is
    deliberate: a function correct on three of four cases is genuinely better
    than one that fails everything, and a binary score throws that away.

    The verdict lines carry a per-run nonce because this harness runs the
    model's code as ``__main__``, and models routinely append a demo block. One
    that prints the word PASS -- `print("self-test:", f(2) == 4 and "PASS")` --
    was counted as a passing check by the old substring count, and a two-check
    task scored 1.50. The model cannot predict the nonce, so it cannot vote on
    its own work. That was the fourth harness bug here to score the rig rather
    than the model; this one inflated instead of deflating, which is why it
    went unnoticed.
    """
    if not checks:
        return 0.0, "no checks"

    code = strip_fences(model_output)
    marker = f"CHECK-{secrets.token_hex(8)}"
    harness = (
        HARNESS_PREAMBLE
        + code
        + "\n"
        + HARNESS_HELPERS
        + f"\nCHECKS = {checks!r}\n"
        + "for _i, _c in enumerate(CHECKS):\n"
        + "    try:\n"
        + "        _ok = bool(eval(_c))\n"
        + "    except Exception:\n"
        + "        _ok = False\n"
        + f"    print({marker!r}, _i, 'PASS' if _ok else 'FAIL')\n"
    )
    with tempfile.TemporaryDirectory() as tmp:
        script = Path(tmp) / "harness.py"
        script.write_text(harness)
        try:
            proc = subprocess.run(
                [sys.executable, str(script)],
                capture_output=True,
                text=True,
                timeout=CODE_TIMEOUT_SECONDS,
                cwd=tmp,
            )
        except subprocess.TimeoutExpired:
            return 0.0, "timeout"

    results = parse_check_results(proc.stdout, marker)
    passed = sum(1 for i in range(len(checks)) if results.get(i))
    if passed == 0 and proc.returncode != 0:
        # Distinguish "wrote broken code" from "wrote code that fails cases"
        first_line = (proc.stderr or "").strip().splitlines()[-1:] or [""]
        return 0.0, f"did not run: {first_line[0][:80]}"
    return passed / len(checks), f"{passed}/{len(checks)} checks"


def normalize_answer(text: str) -> str:
    """Reduce a free-text reply to a comparable token.

    Models wrap a number in prose, punctuation, or thousands separators even
    when told not to; none of that is the thing being measured.
    """
    cleaned = (text or "").strip().replace(",", "")
    numbers = re.findall(r"-?\d+(?:\.\d+)?", cleaned)
    if numbers:
        value = numbers[-1]  # the conclusion, if it reasoned out loud
        return value.rstrip("0").rstrip(".") if "." in value else value
    return cleaned.lower().strip(" .!\"'")


def score_exact(model_output: str, expected: str) -> tuple[float, str]:
    got = normalize_answer(model_output)
    want = normalize_answer(expected)
    return (1.0, f"{got!r}") if got == want else (0.0, f"got {got!r} want {want!r}")


def score_tool(tool_calls: list, task: dict) -> tuple[float, str]:
    """Score a tool-use task structurally — no judge required.

    Three things are worth distinguishing and each is a task in the set:
    calling the right tool, choosing correctly between several, and NOT
    calling a tool when none applies. The last is scored because a model that
    reaches for a tool on every prompt is a real failure mode in an agent
    loop.
    """
    expected_tool = task.get("expect_tool")

    if expected_tool is None:
        return (1.0, "correctly abstained") if not tool_calls else (
            0.0,
            f"called {tool_calls[0]['function']['name']} when none applied",
        )

    if not tool_calls:
        return 0.0, "no tool call"

    called = tool_calls[0]["function"]["name"]
    if called != expected_tool:
        return 0.0, f"called {called}, wanted {expected_tool}"

    # Right tool. Give partial credit and check the arguments that matter.
    expected_args = task.get("expect_args") or {}
    if not expected_args:
        return 1.0, "correct tool"
    try:
        got_args = json.loads(tool_calls[0]["function"].get("arguments") or "{}")
    except json.JSONDecodeError:
        return 0.5, "correct tool, unparseable arguments"

    matched = 0
    for key, want in expected_args.items():
        got = got_args.get(key)
        if isinstance(want, str) and isinstance(got, str):
            ok = want.lower() in got.lower()
        else:
            ok = str(got) == str(want)
        matched += bool(ok)
    # Floor at 0.5 for picking the right tool; arguments carry the rest.
    return 0.5 + 0.5 * (matched / len(expected_args)), f"correct tool, {matched}/{len(expected_args)} args"


# --- provider calls -------------------------------------------------------

def call_model(
    base_url: str,
    api_key: str,
    model_id: str,
    task: dict,
    max_output_tokens: Optional[int] = None,
) -> tuple[str, list, bool]:
    """One completion. Returns (text, tool_calls, truncated)."""
    budget = EVAL_MAX_TOKENS
    if max_output_tokens:
        budget = min(budget, max_output_tokens)
    body = {
        "model": model_id,
        "messages": [{"role": "user", "content": task["prompt"]}],
        "temperature": 0,
        "max_tokens": budget,
    }
    if task.get("tools"):
        body["tools"] = task["tools"]

    resp = requests.post(
        f"{base_url}/chat/completions",
        headers={"authorization": f"Bearer {api_key}"},
        json=body,
        timeout=CALL_TIMEOUT_SECONDS,
    )
    resp.raise_for_status()
    choice = (resp.json().get("choices") or [{}])[0]
    message = choice.get("message") or {}
    # finish_reason 'length' means the budget ran out mid-answer. Whatever
    # came back is an artifact of the cap, so the caller skips rather than
    # scoring it — the same rule as an unusable judge reply.
    truncated = choice.get("finish_reason") == "length"
    return message.get("content") or "", message.get("tool_calls") or [], truncated


JUDGE_SYSTEM = (
    "You are scoring one model's answer against a rubric. Reply with ONLY a "
    'JSON object: {"score": <number 0 to 1>, "reason": "<at most 12 words>"}. '
    "Do not explain your reasoning outside the JSON. "
    "Be strict: award a high score only if every rubric requirement is met."
)

JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)


def extract_judge_json(raw: str) -> Optional[dict]:
    """Pull the score object out of a judge reply.

    Judges are themselves reasoning models and leak their thinking into the
    content despite response_format, so the object usually arrives wrapped in
    prose ("Let me evaluate the answer against the rubric: ... {...}"). Before
    this existed, 44% of judge calls were unparseable.
    """
    if not raw:
        return None
    candidates = [raw]
    match = JSON_OBJECT_RE.search(raw)
    if match:
        candidates.append(match.group(0))
    for candidate in candidates:
        try:
            parsed = json.loads(candidate)
        except json.JSONDecodeError:
            continue
        if isinstance(parsed, dict) and "score" in parsed:
            return parsed
    return None


def score_judge(
    base_url: str, api_key: str, judge_model: str, task: dict, answer: str
) -> Optional[tuple[float, str]]:
    """Score a prose task against its rubric.

    Returns None when the judge itself failed, so the caller can SKIP the
    task rather than record it. Scoring a judge malfunction as 0.0 would
    charge the model for the judge's formatting — it put genuine 0.0 entries
    in translation and general_chat on the first run, for answers that were
    in fact fine.
    """
    body = {
        "model": judge_model,
        "messages": [
            {"role": "system", "content": JUDGE_SYSTEM},
            {
                "role": "user",
                "content": (
                    f"TASK GIVEN TO THE MODEL:\n{task['prompt']}\n\n"
                    f"RUBRIC:\n{task['rubric']}\n\n"
                    f"THE MODEL'S ANSWER:\n{answer}"
                ),
            },
        ],
        "temperature": 0,
        # Generous, because a reasoning judge spends tokens thinking before
        # the JSON and a truncated object is unparseable.
        "max_tokens": 1500,
        "response_format": {"type": "json_object"},
    }
    resp = requests.post(
        f"{base_url}/chat/completions",
        headers={"authorization": f"Bearer {api_key}"},
        json=body,
        timeout=CALL_TIMEOUT_SECONDS,
    )
    resp.raise_for_status()
    raw = ((resp.json().get("choices") or [{}])[0].get("message") or {}).get("content") or ""
    parsed = extract_judge_json(raw)
    if parsed is None:
        return None
    try:
        score = float(parsed["score"])
    except (KeyError, TypeError, ValueError):
        return None
    return max(0.0, min(1.0, score)), str(parsed.get("reason", ""))[:70]


# --- selection ------------------------------------------------------------

def eval_identities(conn: sqlite3.Connection, cfg: RouterConfig) -> list[dict]:
    """Model rows worth measuring directly.

    Flex rows are normally excluded: same weights, same reasoning setting,
    different queue — their quality is the standard row's, inherited via
    ``propagate_to_variants``. ``-fast`` rows ARE measured, because reasoning
    being off genuinely changes answers.

    But that exclusion assumes an evaluated standard equivalent exists, and
    it does not always: glm-5.2-flex reasons by default, while the only
    routable standard glm row is glm-5.2-fast (reasoning reduced) — the
    matching glm-5.2 is canary and therefore never evaluated. Such orphans
    are measured directly rather than left with no proficiency at all.
    """
    placeholders = ",".join("?" * len(cfg.routing.allowed_access_levels))
    rows = conn.execute(
        f"""
        SELECT model_id, base_model_id, reasoning_mode, max_output_tokens FROM models m
         WHERE availability = 'active'
           AND access_level IN ({placeholders})
           AND (
                latency_class = 'standard'
                -- ...or a flex row with no routable standard equivalent to
                -- inherit from, which would otherwise score nothing at all
                OR NOT EXISTS (
                    SELECT 1 FROM models s
                     WHERE s.base_model_id = m.base_model_id
                       AND s.provider = m.provider
                       AND s.reasoning_mode = m.reasoning_mode
                       AND s.context_variant = m.context_variant
                       AND s.latency_class = 'standard'
                       AND s.availability = 'active'
                       AND s.access_level IN ({placeholders})
                )
           )
         ORDER BY model_id
        """,
        tuple(cfg.routing.allowed_access_levels) * 2,
    ).fetchall()
    return [
        {
            "model_id": r[0],
            "base_model_id": r[1],
            "reasoning_mode": r[2],
            "max_output_tokens": r[3],
        }
        for r in rows
    ]


# Tried in order when the configured judge belongs to the model under test.
# A LIST rather than one name, because a single alternate can itself be that
# model: with `--judge-model qwen3.6-35b`, the old guard fired on qwen3.6-35b
# and handed back qwen3.6-35b -- the exact conflict it exists to prevent, and
# silent, because self-judged scores look like any other judge output.
ALTERNATE_JUDGES = ("kimi-k3", "qwen3.6-35b", "gemma-4-31b")


def judge_for(model_id: str, default_judge: str) -> Optional[str]:
    """Pick a judge that is not the model being judged, or None.

    A model scoring its own prose is a known bias, and the default judge is
    itself in the evaluated set. Matching is on the FAMILY, so a -fast row does
    not get judged by its reasoning-on sibling either -- same weights.

    Returns None when every candidate shares the model's family, so the caller
    skips the task. No sample beats a self-graded one, which is the same rule
    score_judge already follows for a judge that returned nothing usable.
    """
    from poller import parse_base_model_id

    family = parse_base_model_id(model_id)
    for candidate in (default_judge, *ALTERNATE_JUDGES):
        if parse_base_model_id(candidate) != family:
            return candidate
    return None


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--tasks", default=TASKS_PATH)
    ap.add_argument("--models", help="comma-separated model_ids")
    ap.add_argument("--categories", help="comma-separated categories")
    ap.add_argument("--judge-model", default="kimi-k3")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    cfg = load_config("config.yaml")
    tasks = (yaml.safe_load(Path(args.tasks).read_text()) or {}).get("tasks") or []

    if args.categories:
        wanted = {c.strip() for c in args.categories.split(",")}
        tasks = [t for t in tasks if t["category"] in wanted]

    unknown = {t["category"] for t in tasks} - set(cfg.proficiency.categories)
    if unknown:
        print(f"tasks reference unknown categories: {sorted(unknown)}", file=sys.stderr)
        return 1

    conn = sqlite3.connect(cfg.database.path)
    identities = eval_identities(conn, cfg)
    if args.models:
        wanted = {m.strip() for m in args.models.split(",")}
        identities = [i for i in identities if i["model_id"] in wanted]

    if not identities or not tasks:
        print("nothing to run", file=sys.stderr)
        return 1

    judge_tasks = sum(1 for t in tasks if t["kind"] == "judge")
    print(
        f"{len(identities)} models x {len(tasks)} tasks = "
        f"{len(identities) * len(tasks)} calls"
        + (f" (+{len(identities) * judge_tasks} judge calls)" if judge_tasks else "")
    )
    if args.dry_run:
        for i in identities:
            budget = min(EVAL_MAX_TOKENS, i["max_output_tokens"] or EVAL_MAX_TOKENS)
            print(
                f"  {i['model_id']:24s} reasoning={i['reasoning_mode']:8s} "
                f"budget={budget}"
            )
        print(f"  judge: {args.judge_model}")
        return 0

    settings = cfg.dispatch_providers["neuralwatt"]
    api_key = os.environ.get(settings.api_key_env)
    if not api_key:
        print(f"{settings.api_key_env} is not set", file=sys.stderr)
        return 1

    for identity in identities:
        model_id = identity["model_id"]
        by_category: dict[str, list[float]] = defaultdict(list)
        print(f"\n{model_id}")

        for task in tasks:
            try:
                text, tool_calls, truncated = call_model(
                    settings.base_url,
                    api_key,
                    model_id,
                    task,
                    identity["max_output_tokens"],
                )
            except requests.RequestException as e:
                print(f"  {task['id']:30s} CALL FAILED {type(e).__name__}")
                continue

            if truncated and not tool_calls:
                print(
                    f"  {task['id']:30s}  ---  truncated at "
                    f"{EVAL_MAX_TOKENS} tokens, skipped"
                )
                continue

            kind = task["kind"]
            try:
                if kind == "code":
                    score, detail = score_code(text, task["checks"])
                elif kind == "exact":
                    score, detail = score_exact(text, str(task["answer"]))
                elif kind == "tool":
                    score, detail = score_tool(tool_calls, task)
                elif kind == "judge":
                    judge_model = judge_for(model_id, args.judge_model)
                    if judge_model is None:
                        print(
                            f"  {task['id']:30s}  ---  no judge outside "
                            f"{model_id}'s family, skipped"
                        )
                        continue
                    judged = score_judge(
                        settings.base_url,
                        api_key,
                        judge_model,
                        task,
                        text,
                    )
                    if judged is None:
                        # No sample rather than a zero: a judge that failed to
                        # emit parseable JSON says nothing about the model.
                        print(f"  {task['id']:30s}  ---  judge unusable, skipped")
                        continue
                    score, detail = judged
                else:
                    print(f"  {task['id']:30s} unknown kind {kind!r}")
                    continue
            except requests.RequestException as e:
                print(f"  {task['id']:30s} SCORING FAILED {type(e).__name__}")
                continue

            by_category[task["category"]].append(score)
            print(f"  {task['id']:30s} {score:4.2f}  {detail}")

        for category, scores in by_category.items():
            add_self_eval(conn, cfg, model_id, "neuralwatt", category, scores)
        conn.commit()

    # Equivalent serving variants inherit from the row actually measured —
    # not from the family id, which may name a row that was never evaluated
    # (glm-5.2 is canary, so glm-5.2-flex inherited nothing and scored blank).
    propagated = 0
    for identity in identities:
        propagated += propagate_to_variants(
            conn, cfg, identity["model_id"], "neuralwatt"
        )
    conn.commit()
    print(f"\npropagated {propagated} inherited rows to serving variants")
    conn.close()
    return 0


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