"""
exec_checks.py — single source of truth for exec-gold hygiene + match canonicalization.

WHY THIS EXISTS
---------------
The data generator, the cleaner, the gold census, and the runtime verification ladder
must agree on (a) what counts as a malformed exec gold and (b) how two outputs are
compared for equality. When each reimplements its own rules they drift, and you get
the failure we already saw: bash checked as python, structured tag-output flagged
"unstructured". One module, imported everywhere:

    build_hyper_dataset.py  -> repair() + degeneracy() at author time
    remove_junk.py          -> repair() + degeneracy() + python_for_compiled()
    audit_gold.py           -> (recommended) import well-formedness helpers here
    runtime verifier         -> canonicalize() for match-against-gold

PRECISION vs RECALL (important design split)
--------------------------------------------
audit_gold is RECALL-oriented: flag anything suspicious so a human reads it.
remove_junk is PRECISION-oriented: drop only what is junk under ALL plausible task
intents. The functions here are labeled so each caller picks the right ones:
  * degeneracy()          -> precision-safe drop (junk under any intent)
  * python_for_compiled() -> precision-safe drop (a Python def is never the answer to
                             a C++/Rust/Go/Java task in anything we author)
  * repair()              -> lossless-intent normalization (strip wrapping fence /
                             leading "Sure, here is:" preamble) — safe to always apply
  * canonicalize()        -> equality normalization for match-against-gold

Dependency-free (stdlib only) so every consumer can import it without pulling torch.
"""

from __future__ import annotations

import json
import re
from collections import Counter

# Compiled / typed languages where Qwen tends to fall back to Python.
COMPILED_LANGS = {"c++", "cpp", "rust", "go", "golang", "java", "c", "c#", "csharp"}

# Leading conversational preamble that occasionally survives the open/exec prompt.
_PREAMBLE_RE = re.compile(
    r"^\s*(sure[,!.]?\s*|certainly[,!.]?\s*|of course[,!.]?\s*|"
    r"here(?:'s| is| are)\b[^\n:]*:?\s*|below is\b[^\n:]*:?\s*|"
    r"the (?:answer|result|output) is[:]?\s*)",
    re.I)

_PY_DEF = re.compile(r"(?:^|\n)\s*def\s+\w+\s*\([^)]*\)\s*:", re.M)
_PY_IMPORT = re.compile(r"(?:^|\n)\s*(?:import\s+\w+|from\s+\w+\s+import)\b", re.M)


def strip_fences(s):
    """Strip a SINGLE wrapping ```lang ... ``` fence enclosing the WHOLE string.
    Conservative: only when fence-count == 2 and it wraps everything; multi-block,
    inline-fenced, and non-fenced strings are left untouched."""
    if not isinstance(s, str):
        return s
    t = s.strip()
    if not (t.startswith("```") and t.endswith("```") and len(t) >= 6):
        return s
    if t.count("```") != 2:
        return s
    nl = t.find("\n")
    if nl == -1:
        return s
    inner = t[nl + 1:].rstrip()
    if inner.endswith("```"):
        inner = inner[:-3]
    return inner.strip()


def strip_preamble(s):
    """Remove leading conversational clauses ('Sure, here is the code:'). Iterates to
    a fixed point (capped) so stacked clauses are handled; leaves the artifact intact."""
    if not isinstance(s, str):
        return s
    for _ in range(3):
        new = _PREAMBLE_RE.sub("", s, count=1)
        if new == s:
            break
        s = new
    return s


def repair(s):
    """Lossless-intent normalization applied to any OUTPUT (never inputs): drop a
    wrapping fence, then a leading preamble. Safe to always apply."""
    if not isinstance(s, str):
        return s
    return strip_preamble(strip_fences(s)).strip()


def degeneracy(s):
    """Return a reason string if the output is catastrophically degenerate (empty or
    a repetition loop), else None. Bad under ALL task intents -> precision-safe drop."""
    if not isinstance(s, str):
        return "non-string"
    t = s.strip()
    if not t:
        return "empty"
    lines = [ln for ln in t.splitlines() if ln.strip()]
    if len(lines) >= 6:
        top = Counter(lines).most_common(1)[0][1]
        if top / len(lines) > 0.6:
            return "repeated-line loop"
    toks = t.split()
    if len(toks) >= 20:
        top = Counter(toks).most_common(1)[0][1]
        if top / len(toks) > 0.7:
            return "repeated-token loop"
    return None


def python_for_compiled(domain, out):
    """High-precision: OUTPUT is clearly a Python function/script authored for a
    COMPILED-language task. Returns a reason or None. Gates on domain so it self-skips
    everywhere else."""
    if (domain or "").strip().lower() not in COMPILED_LANGS:
        return None
    if not isinstance(out, str):
        return None
    head = out[:400]
    c_family_syntax = (";" in head) or ("{" in head) or ("#include" in head)
    looks_python = bool(_PY_DEF.search(out) or _PY_IMPORT.search(head))
    if looks_python and not c_family_syntax:
        return "python authored for compiled-language task"
    return None


def canonicalize(s):
    """Equality normalization for match-against-gold. The runtime verifier and any
    eval-time match MUST use this same function so 'equal' means the same thing
    everywhere. JSON is canonicalized (sorted keys, no whitespace); everything else
    is whitespace-collapsed and stripped. Extend per-type as the ladder grows
    (e.g. numeric rounding tolerance), but keep it HERE so it stays shared."""
    if not isinstance(s, str):
        return s
    t = s.strip()
    try:
        return json.dumps(json.loads(t), sort_keys=True, separators=(",", ":"))
    except Exception:
        return re.sub(r"\s+", " ", t)


def matches(pred, gold):
    """Deterministic equality after canonicalization. The strongest exec rung when a
    gold exists (eval-time / consensus). NOT available at runtime for a NOVEL request
    with no gold — there you fall to parse/compile (floor) or RAG-match (knowledge)."""
    return canonicalize(pred) == canonicalize(gold)