"""
Quality scoring for benchmark outputs.

Two independent layers, deliberately separate:

  format_score  — deterministic checks declared per-prompt in prompts.json.
                  Free, reproducible, and the only score used by default.
  judge_score   — optional LLM-as-judge rating 1-5. Costs tokens, is not
                  reproducible run to run, and is reported alongside the format
                  score rather than blended into it.

Blending a stochastic judge into a deterministic score produces a number nobody
can defend. Keep them in separate columns.
"""

from __future__ import annotations

import ast
import json
import re
from typing import Any, Optional

FENCE_RE = re.compile(r"```(?:[a-zA-Z0-9_+-]*)\n(.*?)```", re.S)


# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #

def strip_fence(text: str) -> str:
    m = FENCE_RE.search(text or "")
    return m.group(1).strip() if m else (text or "").strip()


def strip_thinking(text: str) -> str:
    """Remove <think>...</think> blocks emitted by hybrid reasoning models."""
    return re.sub(r"<think>.*?</think>", "", text or "", flags=re.S).strip()


def _first_json(text: str) -> Optional[Any]:
    """Best-effort: parse the first JSON value in the text."""
    s = strip_fence(strip_thinking(text))
    try:
        return json.loads(s)
    except Exception:  # noqa: BLE001
        pass
    for opener, closer in (("{", "}"), ("[", "]")):
        i, j = s.find(opener), s.rfind(closer)
        if i != -1 and j > i:
            try:
                return json.loads(s[i : j + 1])
            except Exception:  # noqa: BLE001
                continue
    return None


def _sentences(text: str) -> list[str]:
    return [p for p in re.split(r"(?<=[.!?])\s+", (text or "").strip()) if p]


_TYPE_OK = {
    "str": lambda v: isinstance(v, str),
    "num": lambda v: isinstance(v, (int, float)) and not isinstance(v, bool),
    "bool": lambda v: isinstance(v, bool),
}


# --------------------------------------------------------------------------- #
# format checks
# --------------------------------------------------------------------------- #

def format_score(text: str, check: dict) -> tuple[float, str]:
    """Return (score in 0.0-1.0, short reason). Unknown check types score 0.5."""
    body = strip_thinking(text or "")
    if not body.strip():
        return 0.0, "empty output"

    kind = check.get("type")

    if kind == "contains_any":
        vals = check.get("values", [])
        hit = any(v.lower() in body.lower() for v in vals)
        if not hit:
            return 0.0, f"missing any of {vals}"
        mw = check.get("max_words")
        if mw and len(body.split()) > mw:
            return 0.6, f"correct but verbose ({len(body.split())}w > {mw})"
        return 1.0, "ok"

    if kind == "max_words":
        n = len(body.split())
        return (1.0, "ok") if n <= check["value"] else (0.0, f"{n}w > {check['value']}")

    if kind == "sentence_count":
        n = len(_sentences(body))
        want, tol = check["value"], check.get("tolerance", 0)
        if n == want:
            return 1.0, "ok"
        if abs(n - want) <= tol:
            return 0.7, f"{n} sentences, wanted {want}"
        return 0.0, f"{n} sentences, wanted {want}"

    if kind == "line_count_min":
        n = len([l for l in body.splitlines() if l.strip()])
        return (1.0, "ok") if n >= check["value"] else (0.0, f"{n} lines < {check['value']}")

    if kind == "python_compiles":
        code = strip_fence(body)
        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            return 0.0, f"syntax error line {e.lineno}"
        want = check.get("must_define")
        if want:
            names = {n.name for n in ast.walk(tree)
                     if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))}
            if want not in names:
                return 0.4, f"parses but no def {want}()"
        return 1.0, "ok"

    if kind == "json_schema":
        obj = _first_json(body)
        if obj is None:
            return 0.0, "not parseable as JSON"
        if not isinstance(obj, dict):
            return 0.2, f"JSON is {type(obj).__name__}, wanted object"
        missing = [k for k in check.get("required_keys", []) if k not in obj]
        if missing:
            return 0.3, f"missing keys {missing}"
        for k, t in (check.get("types") or {}).items():
            if k in obj and not _TYPE_OK[t](obj[k]):
                return 0.6, f"{k} is {type(obj[k]).__name__}, wanted {t}"
        ak = check.get("array_key")
        if ak:
            arr = obj.get(ak)
            if not isinstance(arr, list):
                return 0.5, f"{ak} is not an array"
            if check.get("array_len") and len(arr) != check["array_len"]:
                return 0.7, f"{ak} has {len(arr)}, wanted {check['array_len']}"
        # Strictness bonus condition: prose wrapped around the JSON is a real failure
        # mode for structured-output use, so flag it rather than pass silently.
        if not strip_fence(body).lstrip().startswith(("{", "[")):
            return 0.8, "valid JSON but wrapped in prose"
        return 1.0, "ok"

    if kind == "json_array":
        obj = _first_json(body)
        if not isinstance(obj, list):
            return 0.0, "not a JSON array"
        if check.get("length") and len(obj) != check["length"]:
            return 0.5, f"{len(obj)} items, wanted {check['length']}"
        if check.get("lowercase") and any(
            isinstance(x, str) and x != x.lower() for x in obj
        ):
            return 0.7, "items not lowercase"
        return 1.0, "ok"

    return 0.5, f"unknown check type {kind!r}"


# --------------------------------------------------------------------------- #
# optional LLM-as-judge
# --------------------------------------------------------------------------- #

JUDGE_SYSTEM = (
    "You grade model outputs for a benchmark. Judge only whether the output does what the "
    "instruction asked: correctness, completeness, and obedience to stated format and length "
    "constraints. Ignore style and tone. Reply with ONLY a single integer 1-5.\n"
    "5 = fully correct and follows every constraint\n"
    "4 = correct, one minor constraint missed\n"
    "3 = partially correct or ignores a stated format\n"
    "2 = mostly wrong but on topic\n"
    "1 = wrong, empty, or ignores the instruction"
)


def judge_score(client, judge_model: str, prompt: str, output: str) -> Optional[int]:
    """Rate 1-5 with a judge model. Returns None if the judge fails or misbehaves.

    The judge is a cost. Only call it when --judge-model is passed explicitly.
    """
    if not (output or "").strip():
        return 1
    msg = (
        f"INSTRUCTION GIVEN TO THE MODEL:\n{prompt}\n\n"
        f"MODEL OUTPUT:\n{strip_thinking(output)[:4000]}\n\n"
        "Rating (1-5), integer only:"
    )
    try:
        r = client.chat.completions.create(
            model=judge_model,
            messages=[
                {"role": "system", "content": JUDGE_SYSTEM},
                {"role": "user", "content": msg},
            ],
            max_tokens=8,
            temperature=0.0,
        )
        m = re.search(r"[1-5]", r.choices[0].message.content or "")
        return int(m.group()) if m else None
    except Exception:  # noqa: BLE001
        return None