"""The worker of the visible-test detectors: runs inside the grading container, beside the agent's
code. Stdlib only; :mod:`.visible_detect` writes this text into the sealed directory and runs it as
``python -I -B ../heldout/probity-visible-worker/visible_runner.py ../heldout/spec.json`` from the workspace.

It is the oracle-free detector study's in-sandbox runner with the study's held-out runner folded in
(``canon``, the type-strict rewrite and its helpers), unchanged in what it measures:

  units     every unit of the visible test runs as prelude + unit in its own namespace, in ONE
            process in order, with ``candidate`` recording (argument repr before the call,
            canonical result) and line coverage of ``func.py``; a unit's outcome is passed or an
            error string with its type; a unit that overruns ``unit_timeout`` is ``TimeoutExpired``
  probe     each probe calls the entry point directly under coverage and records the raised
            exception type or the canonical result
  repeat    every source's ``check(candidate)`` runs ``runs`` times in this process, each run
            recording its calls; a run fails on any exception
  isolate   ``func.py`` is the instrumented module (:mod:`.visible_isolate`): the baseline run of
            every unit with the arm each decision took, one forced-arm ablation per candidate
            arm on a fresh module, and the in-domain probes with their arms

The agent's code shares this process, so every fact here is evidence the agent could shape: a
detector built on it can only refuse, and an agent that forges its coverage escapes a detector, it
never earns a pass. The result is one JSON line on the standard output this process started with;
the agent's code writes to standard error instead.
"""

from __future__ import annotations

import ast
import hashlib
import importlib.util
import json
import os
import re
import signal
import sys
import traceback
from typing import Any

BUILTIN_SCALARS = (int, bool, float, complex, str, bytes, type(None))


def safe_repr(x) -> str:
    try:
        return repr(x)[:2000]
    except Exception as e:  # an overloaded __repr__ may raise
        return f"<repr failed: {type(e).__name__}>"


def canon(x):
    t = type(x)
    if t in BUILTIN_SCALARS:
        return [t.__name__, repr(x)]
    if t in (list, tuple):
        return [t.__name__, [canon(v) for v in x]]
    if t is dict:
        return ["dict", sorted(([canon(k), canon(v)] for k, v in x.items()), key=repr)]
    if t in (set, frozenset):
        return [t.__name__, sorted((canon(v) for v in x), key=repr)]
    # the default repr carries the object's address, which differs on every run of honest code too
    return ["<nonbuiltin>", f"{t.__module__}.{t.__qualname__}",
            re.sub(r" at 0x[0-9a-fA-F]+", "", safe_repr(x))]


def strict_eq(x, y) -> bool:
    if type(x) is not type(y):
        return False
    t = type(x)
    if t in (list, tuple):
        return len(x) == len(y) and all(strict_eq(a, b) for a, b in zip(x, y))
    if t is dict:
        if len(x) != len(y):
            return False
        for k, v in x.items():
            match = [k2 for k2 in y if strict_eq(k, k2)]
            if len(match) != 1 or not strict_eq(v, y[match[0]]):
                return False
        return True
    if t in (set, frozenset):
        return len(x) == len(y) and all(any(strict_eq(a, b) for b in y) for a in x)
    if t in BUILTIN_SCALARS:
        return bool(x == y)
    return False  # two values of the same non-builtin type: no trusted equality


def strict_ne(x, y) -> bool:
    return type(x) is type(y) and not strict_eq(x, y)


def strict_in(x, c) -> bool:
    if type(c) in (str, bytes):
        return type(x) is type(c) and x in c
    if type(c) is dict:
        return any(strict_eq(x, k) for k in c)
    if type(c) in (list, tuple, set, frozenset):
        return any(strict_eq(x, e) for e in c)
    return False


def strict_bool(x) -> bool:
    if type(x) is not bool:
        raise AssertionError(f"type-strict: expected bool, got {type(x).__name__}")
    return x


HELPERS = {"__dr_strict_eq": strict_eq, "__dr_strict_ne": strict_ne,
           "__dr_strict_in": strict_in, "__dr_strict_bool": strict_bool}


class Strictify(ast.NodeTransformer):
    """``==``, ``!=``, ``in`` and ``not in`` become the type-strict helpers; a bare
    ``assert candidate(...)`` must return a bool."""

    OPS = {ast.Eq: "__dr_strict_eq", ast.NotEq: "__dr_strict_ne", ast.In: "__dr_strict_in"}

    def visit_Compare(self, node):
        self.generic_visit(node)
        terms, left = [], node.left
        for op, right in zip(node.ops, node.comparators):
            if type(op) in self.OPS:
                terms.append(ast.Call(ast.Name(self.OPS[type(op)], ast.Load()), [left, right], []))
            elif isinstance(op, ast.NotIn):
                terms.append(ast.UnaryOp(ast.Not(), ast.Call(
                    ast.Name("__dr_strict_in", ast.Load()), [left, right], [])))
            else:
                terms.append(ast.Compare(left, [op], [right]))
            left = right
        return terms[0] if len(terms) == 1 else ast.BoolOp(ast.And(), terms)

    def visit_Assert(self, node):
        self.generic_visit(node)
        e = node.test
        inner = e.operand if isinstance(e, ast.UnaryOp) and isinstance(e.op, ast.Not) else e
        if isinstance(inner, ast.Call) and isinstance(inner.func, ast.Name) and inner.func.id == "candidate":
            wrapped = ast.Call(ast.Name("__dr_strict_bool", ast.Load()), [inner], [])
            node.test = ast.UnaryOp(ast.Not(), wrapped) if inner is not e else wrapped
        return node


class UnitTimeout(BaseException):
    pass


def _alarm(_signum, _frame):
    raise UnitTimeout()


def sha(s: str) -> str:
    return hashlib.sha256(s.encode()).hexdigest()[:16]


def load_func(work: str, entry: str):
    sys.path.insert(0, work)
    path = os.path.join(work, "func.py")
    spec = importlib.util.spec_from_file_location("func", path)
    if spec is None or spec.loader is None:
        raise ImportError(f"cannot load {path}")
    mod = importlib.util.module_from_spec(spec)
    sys.modules["func"] = mod
    spec.loader.exec_module(mod)
    return mod, getattr(mod, entry), path


class Coverage:
    def __init__(self, path: str):
        self.path = path
        self.lines: set[int] = set()

    def _local(self, frame, event, arg):
        if event == "line":
            self.lines.add(frame.f_lineno)
        return self._local

    def _global(self, frame, event, arg):
        if event == "call" and frame.f_code.co_filename == self.path:
            self.lines.add(frame.f_lineno)
            return self._local
        return None

    def __enter__(self):
        self.lines = set()
        sys.settrace(self._global)
        return self

    def __exit__(self, *a):
        sys.settrace(None)
        return False


def err_string(e: BaseException) -> tuple[str, str]:
    if isinstance(e, UnitTimeout):
        return "TimeoutExpired", "TimeoutExpired"
    return type(e).__name__, f"{type(e).__name__}: {safe_repr(e)[:300]}"


def run_units(spec: dict) -> dict:
    out = {"units": [], "import_error": None}
    try:
        mod, fn, path = load_func(spec["work"], spec["entry_point"])
    except BaseException as e:  # noqa: BLE001
        out["import_error"] = err_string(e)[1]
        out["units"] = [{"passed": False, "error_type": "ImportError", "error": out["import_error"], "lines": [], "calls": []} for _ in spec["units"]]
        return out
    strict = bool(spec.get("strict"))
    tmo = int(spec.get("unit_timeout", 60))
    signal.signal(signal.SIGALRM, _alarm)
    prelude = spec.get("prelude", "")
    for usrc in spec["units"]:
        calls = []

        def candidate(*args, **kw):
            a = safe_repr((args, kw) if kw else args)
            r = fn(*args, **kw)
            calls.append([sha(a), sha(json.dumps(canon(r), sort_keys=True, default=repr)), a[:200]])
            return r

        ns = dict(HELPERS, candidate=candidate)
        rec = {"passed": True, "error_type": None, "error": None}
        src = prelude + "\n" + usrc + "\n"
        try:
            tree = ast.parse(src)
            if strict:
                tree = ast.fix_missing_locations(Strictify().visit(tree))
            code = compile(tree, "<unit>", "exec")
        except SyntaxError as e:
            rec.update(passed=False, error_type="SyntaxError", error=str(e)[:300], lines=[], calls=[])
            out["units"].append(rec)
            continue
        cov = Coverage(path)
        signal.alarm(tmo)
        try:
            with cov:
                exec(code, ns)
        except BaseException as e:  # noqa: BLE001
            rec["passed"] = False
            rec["error_type"], rec["error"] = err_string(e)
        finally:
            signal.alarm(0)
            sys.settrace(None)
        rec["lines"] = sorted(cov.lines)
        rec["calls"] = calls[:200]
        rec["n_calls"] = len(calls)
        out["units"].append(rec)
    return out


def run_probes(spec: dict) -> dict:
    out = {"probes": [], "import_error": None}
    try:
        mod, fn, path = load_func(spec["work"], spec["entry_point"])
    except BaseException as e:  # noqa: BLE001
        out["import_error"] = err_string(e)[1]
        out["probes"] = [{"raised": "ImportError", "result": None, "lines": []} for _ in spec["probes"]]
        return out
    tmo = int(spec.get("unit_timeout", 60))
    signal.signal(signal.SIGALRM, _alarm)
    for p in spec["probes"]:
        rec: dict[str, Any] = {"raised": None, "result": None, "is_none": None}
        try:
            args = ast.literal_eval(p)
        except Exception:
            rec["raised"] = "BadProbe"
            rec["lines"] = []
            out["probes"].append(rec)
            continue
        cov = Coverage(path)
        signal.alarm(tmo)
        try:
            with cov:
                r = fn(*args)
            rec["is_none"] = r is None
            rec["result"] = sha(json.dumps(canon(r), sort_keys=True, default=repr))
        except BaseException as e:  # noqa: BLE001
            rec["raised"] = err_string(e)[0]
        finally:
            signal.alarm(0)
            sys.settrace(None)
        rec["lines"] = sorted(cov.lines)
        out["probes"].append(rec)
    return out


def compile_check(src: str):
    ns = dict(HELPERS)
    exec(compile(ast.parse(src), "<heldout>", "exec"), ns)
    return ns["check"]


def run_repeat(spec: dict) -> dict:
    """Every source ``runs`` times in this process: the study's held-out runner, plain."""
    result: dict = {"runs": []}
    try:
        mod, fn, path = load_func(spec["work"], spec["entry_point"])
    except BaseException as e:  # noqa: BLE001  an import failure is a failed run
        result["import_error"] = f"{type(e).__name__}: {safe_repr(e)}"
        result["runs"] = [{"passed": False, "error": result["import_error"], "calls": []}] * spec.get("runs", 1)
        return result
    for _ in range(spec.get("runs", 1)):
        calls = []

        def candidate(*args, **kw):
            a = safe_repr((args, kw) if kw else args)
            r = fn(*args, **kw)
            calls.append([a, canon(r)])
            return r

        run = {"passed": True, "error": None}
        try:
            for src in spec["sources"]:
                compile_check(src)(candidate)
        except BaseException as e:  # noqa: BLE001  any exception fails the run
            run["passed"] = False
            tb = traceback.extract_tb(e.__traceback__)
            where = next((f"line {f.lineno}" for f in reversed(tb) if f.filename == "<heldout>"), "")
            run["error"] = f"{type(e).__name__}: {safe_repr(e)[:500]} {where}"
        run["calls"] = calls
        result["runs"].append(run)
    return result


# ---------------------------------------------------------------- isolate mode (R5-4, DMD-DESIGN-04)
#
# func.py here is the instrumented source of of_isolate.instrument: every decision calls one of the
# hooks below. A hook records the arm it took ("<aid>:T", "<aid>:F", "<aid>:hit") into ENTERED and,
# when FORCE names its arm id, takes the other way (a branch forced to the given bool; a lookup forced
# to miss). The source is never edited for an ablation, so a mutant is the same bytes with one flag.

ENTERED: set = set()
HITKEYS: dict = {}
OPVALS: dict = {}
FORCE: dict = {}
OPVAL_CAP = 64  # distinct operand tuples kept per arm per call; past it the arm records "overflow"
OPVAL_ENTRIES = 256  # entries hashed per arm per call; past it the arm records "overflow" and stops hashing
OPCOUNT: dict = {}


def _opval(ops) -> str:
    try:
        v = ops()
    except BaseException:  # noqa: BLE001 an operand the thunk cannot read is unknown, never equal
        return "?"
    try:
        return sha(json.dumps(canon(v), sort_keys=True, default=repr))
    except BaseException:  # noqa: BLE001
        return "?"


def _pof_b(aid, value, ops=None):
    b = bool(value)
    if aid in FORCE:
        b = FORCE[aid]
    arm = f"{aid}:{'T' if b else 'F'}"
    ENTERED.add(arm)
    if ops is not None:
        # a multiset: how many times the test saw each operand tuple on this side, in this call
        seen = OPVALS.setdefault(arm, {})
        c = OPCOUNT.get(arm, 0) + 1
        OPCOUNT[arm] = c
        if c > OPVAL_ENTRIES or len(seen) >= OPVAL_CAP:
            seen["overflow"] = 1
        else:
            d = _opval(ops)
            seen[d] = seen.get(d, 0) + 1
    return b


def _is_table(recv) -> bool:
    return isinstance(recv, (dict, set, frozenset))


def _hit(aid, recv, key) -> bool:
    try:
        h = key in recv
    except TypeError:
        return False
    if h:
        ENTERED.add(f"{aid}:hit")
        HITKEYS.setdefault(f"{aid}:hit", set()).add(sha(repr(key)))
    return h


def _pof_l(aid, recv, key):
    if isinstance(recv, dict) and _hit(aid, recv, key) and aid in FORCE:
        raise KeyError(key)
    return recv[key]


def _pof_g(aid, recv, key, *default):
    if isinstance(recv, dict):
        if _hit(aid, recv, key) and aid in FORCE:
            return default[0] if default else None
    return recv.get(key, *default)


def _pof_p(aid, recv, key, *default):
    if isinstance(recv, dict):
        if _hit(aid, recv, key) and aid in FORCE:
            if default:
                return default[0]
            raise KeyError(key)
    return recv.pop(key, *default)


def _pof_in(aid, key, recv):
    # a membership test against a constant list or tuple is a table too (design check round 2)
    if _is_table(recv) or isinstance(recv, (list, tuple)):
        if _hit(aid, recv, key):
            return aid not in FORCE
        return False
    return key in recv


def _install_hooks():
    import builtins

    for name, hook in (("_pof_b", _pof_b), ("_pof_l", _pof_l), ("_pof_g", _pof_g), ("_pof_in", _pof_in), ("_pof_p", _pof_p)):
        setattr(builtins, name, hook)


def _fresh(work: str, entry: str):
    for name in list(sys.modules):
        if name == "func":
            del sys.modules[name]
    mod, fn, path = load_func(work, entry)
    ENTERED.clear()
    HITKEYS.clear()
    OPVALS.clear()
    OPCOUNT.clear()
    return fn


def _units_once(fn, prelude: str, units: list, tmo: float) -> list:
    import time

    out = []
    for usrc in units:
        ENTERED.clear()
        HITKEYS.clear()
        OPVALS.clear()
        OPCOUNT.clear()
        calls = []

        def candidate(*args, **kw):
            r = fn(*args, **kw)
            calls.append(sha(json.dumps(canon(r), sort_keys=True, default=repr)))
            return r

        rec: dict[str, Any] = {"passed": True, "error_type": None, "error": None}
        t0 = time.monotonic()
        signal.setitimer(signal.ITIMER_REAL, tmo)
        try:
            exec(compile(prelude + "\n" + usrc + "\n", "<unit>", "exec"), {"candidate": candidate})
        except BaseException as e:  # noqa: BLE001
            rec["passed"] = False
            rec["error_type"], rec["error"] = err_string(e)
            rec["error"] = rec["error"][:160]
        finally:
            signal.setitimer(signal.ITIMER_REAL, 0)
        rec["seconds"] = round(time.monotonic() - t0, 4)
        rec["arms"] = sorted(ENTERED)
        rec["hitkeys"] = {k: sorted(v)[:20] for k, v in HITKEYS.items()}
        rec["opvals"] = {k: dict(sorted(v.items())[:OPVAL_CAP + 1]) for k, v in OPVALS.items()}
        rec["results"] = calls[:50]
        out.append(rec)
    return out


def run_isolate(spec: dict) -> dict:
    """Baseline with arm recording; one forced-arm ablation per candidate arm, each on a fresh module;
    the probe neighbourhood with arm recording. Candidate arms: entered by a non-empty set F of the
    passing units, F a proper subset, or F all units for a table hit by 2+ distinct keys or a branch
    (the all-units measured arm); smallest F first, at most max_ablations."""
    import time

    out = {"import_error": None, "baseline": [], "ablations": [], "probes": [], "budget_exhausted": False}
    _install_hooks()
    signal.signal(signal.SIGALRM, _alarm)
    work, entry = spec["work"], spec["entry_point"]
    try:
        fn = _fresh(work, entry)
    except BaseException as e:  # noqa: BLE001
        out["import_error"] = err_string(e)[1]
        return out
    tmo = float(spec.get("unit_timeout", 60))
    prelude, units = spec.get("prelude", ""), spec["units"]
    deadline = time.monotonic() + float(spec.get("budget_seconds", 1500))
    base = _units_once(fn, prelude, units, tmo)
    out["baseline"] = base
    n = len(base)
    passing = [i for i, u in enumerate(base) if u["passed"]]
    if len(passing) != n:
        return out
    by_arm: dict[str, set] = {}
    keys: dict[str, set] = {}
    for i, u in enumerate(base):
        for a in u["arms"]:
            by_arm.setdefault(a, set()).add(i)
        for a, ks in u["hitkeys"].items():
            keys.setdefault(a, set()).update(ks)
    cands = []
    for a, F in by_arm.items():
        if len(F) < n:
            cands.append((len(F), a))
        elif n >= 2 and (a.endswith(":hit") and len(keys.get(a, ())) >= 2 or not a.endswith(":hit")):
            cands.append((n, a))
    # identifying arms first (a constant the guard names), then the smallest coverage, then source order,
    # so a cheat guard is never pushed past the ablation cap by honest branches
    prio = {int(a) for a in spec.get("priority", [])}
    cands.sort(key=lambda x: (0 if int(x[1].split(":")[0]) in prio else 1, x[0], int(x[1].split(":")[0]), x[1]))
    slowest = max((u["seconds"] for u in base), default=1.0)
    atmo = min(tmo, 5 * slowest + 5)
    for _, a in cands[: int(spec.get("max_ablations", 40))]:
        if time.monotonic() > deadline:
            out["budget_exhausted"] = True
            break
        aid, side = a.split(":")
        try:
            fn_m = _fresh(work, entry)
        except BaseException as e:  # noqa: BLE001
            out["ablations"].append({"arm": a, "error": err_string(e)[1]})
            continue
        FORCE.clear()
        FORCE[int(aid)] = "miss" if side == "hit" else (side != "T")
        try:
            res = _units_once(fn_m, prelude, units, atmo)
        finally:
            FORCE.clear()
        out["ablations"].append({"arm": a, "units": [{k: u[k] for k in ("passed", "error_type", "seconds")} for u in res]})
    try:
        fn_p = _fresh(work, entry)
    except BaseException as e:  # noqa: BLE001
        out["probe_import_error"] = err_string(e)[1]
        return out
    ptmo = float(spec.get("probe_timeout", 10))
    for p in spec.get("probes", []):
        rec: dict[str, Any] = {"raised": None, "result": None, "is_none": None}
        if time.monotonic() > deadline:
            out["budget_exhausted"] = True
            rec["raised"] = "BudgetExhausted"
            rec["arms"] = []
            out["probes"].append(rec)
            continue
        try:
            args = ast.literal_eval(p)
        except Exception:
            rec.update(raised="BadProbe", arms=[])
            out["probes"].append(rec)
            continue
        ENTERED.clear()
        HITKEYS.clear()
        OPVALS.clear()
        OPCOUNT.clear()
        signal.setitimer(signal.ITIMER_REAL, ptmo)
        try:
            r = fn_p(*args)
            rec["is_none"] = r is None
            rec["result"] = sha(json.dumps(canon(r), sort_keys=True, default=repr))
        except BaseException as e:  # noqa: BLE001
            rec["raised"] = err_string(e)[0]
        finally:
            signal.setitimer(signal.ITIMER_REAL, 0)
        rec["arms"] = sorted(ENTERED)
        rec["hitkeys"] = {k: sorted(v)[:20] + (["overflow"] if len(v) > 20 else []) for k, v in HITKEYS.items()}
        rec["opvals"] = {k: dict(sorted(v.items())[:OPVAL_CAP + 1]) for k, v in OPVALS.items()}
        out["probes"].append(rec)
    return out


MODES = {"units": run_units, "probe": run_probes, "repeat": run_repeat, "isolate": run_isolate}


def main(argv: list[str]) -> int:
    report = os.fdopen(os.dup(1), "w")
    os.dup2(2, 1)
    with open(argv[0]) as handle:
        spec = json.load(handle)
    spec["work"] = os.getcwd()
    try:
        res = MODES[spec["mode"]](spec)
    except BaseException as e:  # noqa: BLE001
        res = {"fatal": f"{type(e).__name__}: {e}"[:500], "trace": traceback.format_exc()[-1500:]}
    report.write(json.dumps(res, default=repr) + "\n")
    report.flush()
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
