"""Numeric-correctness evaluation for the civil reactor model.

Why this exists
---------------
The repository's existing benchmarks score assistant-token negative
log-likelihood / perplexity. The project's own independent eval showed that
metric rewards format memorisation that does not transfer: an adapter can score
"53% better" on a template-sharing benchmark while making the model *worse* than
the base on open-ended content. Perplexity is the wrong yardstick for a model
whose whole point is producing *correct numbers*.

This harness measures the thing that actually matters: does the stated number
match what the audited engine computes, within a declared tolerance? Ground
truth is not hand-written; it is computed at run time from the same engines the
answer is supposed to agree with (via ``civil_tools``), so the suite cannot
drift from the physics and there is nothing to keep in sync by hand.

Three modes
-----------
* ``--self-check`` (default): recompute every item's ground truth and confirm
  the suite's own ``reference_answer`` contains it within tolerance. Needs no
  model and runs in CI; it validates both the suite and the scorer.
* ``--answers FILE``: score a JSONL of model outputs (``{"id", "answer"}``) and
  report per-category and overall accuracy.
* ``--endpoint-url URL``: query an OpenAI-compatible local endpoint for each
  prompt, then score. Uses only stdlib ``urllib`` — no extra dependency.

A number is scored, not prose. Extraction masks isotope labels (``Xe-135``,
``I-135``, ``U-235``) so a mass number is never mistaken for an answer, then
compares every remaining number in the text against the engine value with
``math.isclose``. A relative tolerance travels with each item, because "within
2%" means different things for a period (~50 s) and an inventory (~1e15 /cm^3).
"""

from __future__ import annotations

import argparse
import json
import math
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from civil_safety import read_jsonl
from civil_tools import call_tool

ROOT = Path(__file__).resolve().parent
DEFAULT_SUITE = ROOT / "eval_numeric.jsonl"

# A signed integer/decimal/scientific literal. Kept deliberately strict so it
# does not swallow stray punctuation.
NUMBER_RE = re.compile(r"[-+]?(?:\d+\.?\d*|\.\d+)(?:[eE][-+]?\d+)?")

# Element-mass isotope labels (Xe-135, I-135, U-235, Np-239). Masked before
# number extraction so "-135" in "Xe-135" is never read as the answer -135.
ISOTOPE_RE = re.compile(r"[A-Za-z]{1,2}-\d{1,3}\b")


def mask_isotopes(text: str) -> str:
    """Blank out element-mass labels so their mass numbers are not extracted."""
    return ISOTOPE_RE.sub(" ", text)


def extract_numbers(text: str) -> list[float]:
    """Every numeric literal in ``text`` (isotope labels masked out first)."""
    cleaned = mask_isotopes(text)
    values: list[float] = []
    for match in NUMBER_RE.finditer(cleaned):
        token = match.group()
        try:
            values.append(float(token))
        except ValueError:
            continue
    return values


def answer_matches(
    expected: float, text: str, *, rel_tol: float, abs_tol: float
) -> tuple[bool, float | None]:
    """True iff some number in ``text`` is close to ``expected``.

    Returns the matched value too, for reporting. ``abs_tol`` covers values at
    or near zero, where a relative tolerance is meaningless.
    """
    for value in extract_numbers(text):
        if math.isclose(value, expected, rel_tol=rel_tol, abs_tol=abs_tol):
            return True, value
    return False, None


def ground_truth(item: dict[str, Any]) -> float:
    """Engine-computed reference value for an item's ``check`` spec."""
    check = item.get("check")
    if not isinstance(check, dict):
        raise ValueError(f"{item.get('id', '<no id>')}: missing 'check' object")
    tool = check.get("tool")
    field_name = check.get("field")
    args = check.get("args", {})
    result = call_tool(str(tool), args)
    if not result.ok:
        raise ValueError(
            f"{item.get('id')}: ground-truth tool '{tool}' failed: {result.error}"
        )
    if field_name not in result.values:
        raise ValueError(
            f"{item.get('id')}: tool '{tool}' has no field '{field_name}' "
            f"(has: {', '.join(result.values)})"
        )
    return result.values[str(field_name)]


@dataclass
class ScoreRow:
    id: str
    category: str
    expected: float
    passed: bool
    matched: float | None
    detail: str = ""


@dataclass
class Report:
    rows: list[ScoreRow] = field(default_factory=list)

    @property
    def total(self) -> int:
        return len(self.rows)

    @property
    def passed(self) -> int:
        return sum(1 for r in self.rows if r.passed)

    @property
    def accuracy(self) -> float:
        return self.passed / self.total if self.total else 0.0

    def by_category(self) -> dict[str, tuple[int, int]]:
        acc: dict[str, list[int]] = {}
        for row in self.rows:
            bucket = acc.setdefault(row.category, [0, 0])
            bucket[0] += int(row.passed)
            bucket[1] += 1
        return {k: (v[0], v[1]) for k, v in sorted(acc.items())}

    def to_dict(self) -> dict[str, Any]:
        return {
            "total": self.total,
            "passed": self.passed,
            "accuracy": self.accuracy,
            "by_category": {
                k: {"passed": p, "total": t} for k, (p, t) in self.by_category().items()
            },
            "failures": [
                {
                    "id": r.id,
                    "category": r.category,
                    "expected": r.expected,
                    "matched": r.matched,
                    "detail": r.detail,
                }
                for r in self.rows
                if not r.passed
            ],
        }

    def render(self) -> str:
        lines = ["Numeric-correctness report", "=" * 34]
        for category, (p, t) in self.by_category().items():
            lines.append(f"  {category:34s} {p:3d}/{t:<3d}  {(p / t if t else 0):6.1%}")
        lines.append("-" * 34)
        lines.append(
            f"  {'OVERALL':34s} {self.passed:3d}/{self.total:<3d}  {self.accuracy:6.1%}"
        )
        failures = [r for r in self.rows if not r.passed]
        if failures:
            lines.append("\nFailures:")
            for r in failures:
                got = "no numeric match" if r.matched is None else f"got {r.matched:g}"
                lines.append(
                    f"  - {r.id} ({r.category}): expected {r.expected:g}, {got}. {r.detail}"
                )
        return "\n".join(lines)


def _tolerances(item: dict[str, Any]) -> tuple[float, float]:
    check = item.get("check", {})
    return float(check.get("rel_tol", 0.02)), float(check.get("abs_tol", 1e-9))


def score_answers(items: list[dict[str, Any]], answers: dict[str, str]) -> Report:
    """Score model answers (id -> answer text) against engine ground truth."""
    report = Report()
    for item in items:
        item_id = str(item.get("id", "<no id>"))
        category = str(item.get("category", "uncategorized"))
        expected = ground_truth(item)
        rel_tol, abs_tol = _tolerances(item)
        text = answers.get(item_id)
        if text is None:
            report.rows.append(
                ScoreRow(item_id, category, expected, False, None, "no answer supplied")
            )
            continue
        passed, matched = answer_matches(
            expected, text, rel_tol=rel_tol, abs_tol=abs_tol
        )
        report.rows.append(ScoreRow(item_id, category, expected, passed, matched))
    return report


def self_check(items: list[dict[str, Any]]) -> Report:
    """Confirm each item's reference_answer embeds its engine ground truth."""
    answers = {}
    for item in items:
        ref = item.get("reference_answer")
        if not isinstance(ref, str) or not ref.strip():
            raise ValueError(
                f"{item.get('id')}: self-check needs a 'reference_answer' string"
            )
        answers[str(item.get("id"))] = ref
    return score_answers(items, answers)


def _query_endpoint(url: str, prompt: str, *, system: str, timeout: float) -> str:
    """POST one chat completion to an OpenAI-compatible endpoint (stdlib only)."""
    import urllib.error
    import urllib.request

    payload = json.dumps(
        {
            "model": "vers3dynamics-civil-reactor-expert",
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": prompt},
            ],
            "temperature": 0.0,
            "max_tokens": 512,
        }
    ).encode("utf-8")
    request = urllib.request.Request(
        url, data=payload, headers={"Content-Type": "application/json"}
    )
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            body = json.loads(response.read().decode("utf-8"))
    except (urllib.error.URLError, TimeoutError) as exc:  # pragma: no cover - network
        raise RuntimeError(f"endpoint request failed: {exc}") from exc
    return str(body["choices"][0]["message"]["content"])


_ENDPOINT_SYSTEM = (
    "You are Vers3Dynamics Civil Reactor Expert. Answer civil-nuclear questions "
    "with explicit numbers and units. Stay civil-only."
)


def collect_endpoint_answers(
    items: list[dict[str, Any]], url: str, *, timeout: float = 60.0
) -> dict[str, str]:
    answers: dict[str, str] = {}
    for item in items:
        answers[str(item.get("id"))] = _query_endpoint(
            url, str(item.get("prompt", "")), system=_ENDPOINT_SYSTEM, timeout=timeout
        )
    return answers


def _load_answers_file(path: Path) -> dict[str, str]:
    answers: dict[str, str] = {}
    for row in read_jsonl(path):
        answers[str(row.get("id"))] = str(row.get("answer", ""))
    return answers


def _main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Numeric-correctness eval against the physics engines."
    )
    parser.add_argument(
        "--suite", type=Path, default=DEFAULT_SUITE, help="Eval suite JSONL."
    )
    parser.add_argument(
        "--self-check", action="store_true", help="Validate suite + scorer (no model)."
    )
    parser.add_argument(
        "--answers", type=Path, help="JSONL of model answers ({id, answer})."
    )
    parser.add_argument(
        "--endpoint-url", help="OpenAI-compatible endpoint to query per prompt."
    )
    parser.add_argument("--report", type=Path, help="Write the JSON report here.")
    parser.add_argument(
        "--min-accuracy",
        type=float,
        default=None,
        help="Fail (exit 1) below this accuracy.",
    )
    args = parser.parse_args(argv)

    items = read_jsonl(args.suite)
    if not items:
        print(f"empty suite: {args.suite}", file=sys.stderr)
        return 2

    if args.answers:
        report = score_answers(items, _load_answers_file(args.answers))
    elif args.endpoint_url:
        report = score_answers(
            items, collect_endpoint_answers(items, args.endpoint_url)
        )
    else:
        report = self_check(items)  # default mode

    print(report.render())
    if args.report:
        args.report.write_text(json.dumps(report.to_dict(), indent=2), encoding="utf-8")
        print(f"\nwrote {args.report}")

    if args.min_accuracy is not None and report.accuracy < args.min_accuracy:
        print(
            f"\nFAIL: accuracy {report.accuracy:.1%} < required {args.min_accuracy:.1%}",
            file=sys.stderr,
        )
        return 1
    # In self-check mode, any failure is a broken suite and must be non-zero.
    if not args.answers and not args.endpoint_url and report.passed != report.total:
        return 1
    return 0


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