"""Verified-retry routing agent. No memorization: nothing here is keyed to a known task.

The whole edge is a loop any unseen problem also gets: ask a model, RUN its program against the
sample cases the statement itself publishes, and if a sample fails, hand the model the concrete
counter-example and ask again -- escalating the pool model as attempts go. Measured on 445 real
enclave responses: 63% of wrong answers already fail a published sample, so they are detectable
before grading; the other 37% pass the samples and no loop can see them.

Deliberately absent (this is the part the held-out audit ejects, and it is why it is absent):
  * no prompt->answer table, no per-task fingerprint routing, no hand-written algorithm contracts.
Every decision below is computed from the prompt in front of it, so held-out tasks get exactly the
same treatment as pool ones.

`weights` is a tiny JSON knob file, not a lookup table: entry rung, escalation order, deadlines.
"""

import json
import re
import subprocess
import sys
import time

_MODELS = (
    "qwen/qwen3.7-flash",
    "deepseek/deepseek-v4-flash",
    "deepseek/deepseek-v4-pro",
    "z-ai/glm-5.2",
    "openai/gpt-5.6-luna",
    "google/gemini-3.6-flash",
    "moonshotai/kimi-k3",
)
_PARAMS = {"max_tokens": 16384, "reasoning": {"effort": "low"}}

# Time. The operator's attempt deadline is 900 s for the whole epoch (6 tasks), so the agent keeps
# its own budget well inside that: a watchdog-abandoned run is graded as a WRONG ANSWER, which
# would cost more than any retry can win back.
_RUN_DEADLINE_S = 560.0        # whole-epoch ceiling this agent will not cross
_TASK_LADDER_S = 70.0          # per-task ceiling for the retry ladder
_EXEC_BUDGET_S = 20.0          # per-task wall clock spent RUNNING candidate programs
_CASE_TIMEOUT_S = 3.0          # one sample case
_MAX_CASES = 4                 # sample cases checked per attempt

_RETRY = (
    "Your previous program was run on a sample case published in the statement above and it was "
    "wrong. On the input\n%s\nit printed\n%s\nbut the statement's own expected output is\n%s\n"
    "Work out where the reasoning breaks and write a corrected complete program. Match the "
    "expected output exactly, including the number of digits and the number of lines."
)
_ONLY_SOURCE = ("Return ONLY raw complete Python 3 source, no Markdown fences, no prose, "
                "no explanation before or after the code.")


def _is_code(prompt):
    t = str(prompt)
    return ("Write a complete Python 3 program" in t
            and "standard input" in t and "standard output" in t)


def _samples(prompt):
    """(stdin, expected) pairs the STATEMENT publishes. Generic parse, no task knowledge.

    The answer is the first paragraph after each marker: the blocks that follow it are prose
    explaining the case, and including them was what made an early version of this check reject
    correct programs.
    """
    t = str(prompt).replace("\r\n", "\n").replace("\r", "\n")
    parts = re.split(r"\n\s*Sample (Input|Output) \d+\s*\n", t)
    ins, outs = [], []
    for i in range(1, len(parts) - 1, 2):
        first = parts[i + 1].split("\n\n")[0].strip("\n")
        (ins if parts[i] == "Input" else outs).append(first)
    return list(zip(ins, outs))[:_MAX_CASES]


def _extract(answer):
    t = str(answer).strip()
    if t.startswith("```"):
        t = re.sub(r"^```[a-zA-Z0-9]*\n", "", t)
        t = re.sub(r"\n```\s*$", "", t)
    return t


def _run_one(code, stdin, timeout):
    """Execute a candidate program on one sample input. Returns stdout, or None if it did not run.

    `-I` isolates the interpreter (no site, no env-driven imports). Failure to run is NOT evidence
    the answer is wrong -- it is treated as unknown, so a sandbox quirk can never manufacture a
    retry that replaces a good answer with a worse one.
    """
    try:
        r = subprocess.run([sys.executable, "-I", "-c", code],
                           input=stdin if stdin.endswith("\n") else stdin + "\n",
                           capture_output=True, text=True, timeout=timeout)
        return r.stdout
    except Exception:
        return None


def _check(code, cases, clock):
    """(passes, fails, first_failure). `clock` is a one-element list of remaining exec seconds."""
    if not code.strip():
        return 0, 1, (cases[0][0], "", cases[0][1]) if cases else (0, 0, None)
    passes = fails = 0
    bad = None
    for stdin, want in cases:
        if clock[0] <= 0.0:
            break
        t0 = time.monotonic()
        got = _run_one(code, stdin, _CASE_TIMEOUT_S)
        clock[0] -= time.monotonic() - t0
        if got is None:                      # crashed or timed out -> unknown, not a failure
            continue
        if got.split() == want.split():
            passes += 1
        elif bad is None:
            fails += 1
            bad = (stdin, got, want)
        else:
            fails += 1
    return passes, fails, bad


def build_agent(weights):
    cfg = json.loads(bytes(weights).decode("utf-8"))
    if cfg.get("kind") != "verified-retry-1":
        raise ValueError("weights do not match the verified-retry-1 contract")
    entry = int(cfg["entry"])                       # rung for the first attempt
    ladder = [int(x) for x in cfg["ladder"]]        # rungs for successive retries
    floor_rung = int(cfg.get("floor", entry))       # non-code benchmarks (weight-0 floors)
    if not 0 <= entry < len(_MODELS) or any(not 0 <= r < len(_MODELS) for r in ladder):
        raise ValueError("rung out of range")
    started = [None]

    def agent(prompt, call_model):
        if started[0] is None:
            started[0] = time.monotonic()
        text = str(prompt)
        params = {"max_tokens": _PARAMS["max_tokens"],
                  "reasoning": dict(_PARAMS["reasoning"])}

        if not _is_code(text):                      # MMLU / GSM8K: one cheap call, no loop
            return call_model(_MODELS[floor_rung], [{"role": "user", "content": text}], params)

        first = call_model(_MODELS[entry], [{"role": "user", "content": text}], params)
        try:
            cases = _samples(text)
            if not cases:
                return first
            t_task = time.monotonic()
            clock = [_EXEC_BUDGET_S]
            best = first
            passes, fails, bad = _check(_extract(first), cases, clock)
            if fails == 0 or bad is None:
                return first                        # samples agree -> ship it
            for rung in ladder:
                if clock[0] <= 0.0:
                    break
                if time.monotonic() - t_task > _TASK_LADDER_S:
                    break
                if time.monotonic() - started[0] > _RUN_DEADLINE_S:
                    break
                retry = text + "\n\n" + (_RETRY % bad) + "\n" + _ONLY_SOURCE
                nxt = call_model(_MODELS[rung], [{"role": "user", "content": retry}], params)
                npass, nfail, nbad = _check(_extract(nxt), cases, clock)
                if npass > passes:                  # keep the strictly better candidate only
                    best, passes, fails = nxt, npass, nfail
                    bad = nbad if nbad is not None else bad
                if nfail == 0 and npass > 0:
                    return nxt
                elif nbad is not None:
                    bad = nbad
            return best
        except Exception:
            return first                            # any loop fault -> the plain answer stands

    return agent