# build v17b
"""SN99 agent v17b: sample-gated ladder + dual-oracle stress.

Escalate a cheap ladder until statement samples pass, then (budget allowing) stress
against two independent references before banking. Design notes stay in comments.
"""

# v17b -- SilentByte ladder spine + beyondyourself dual-oracle stress, without encoder routing.
#
# WHY NO HEAD ROUTING HERE.
#   Agent code runs under confine RLIMIT_AS (~2GB). Loading the pinned encoder pulls torch and
#   OpenBLAS can exit(1) the guest (SilentByte measured: 12 attempts, 0 uploads). The harness
#   router path (v17a) owns the encoder; this agent path must not touch it. Weights still ship
#   (velev RouterHead bytes) so the chain bundle hash binds source+weights; they are not read
#   for routing.
#
# PHASE A -- SAMPLE-GATED LADDER (SilentByte).
#   harness.verifier_ok only checks shape, so a wrong-but-well-formed draft never escalates.
#   Run the draft on the statement's own samples (exact token compare). Escalate model/effort
#   only while samples fail. Medium effort on luna first; gemini then kimi. Floors (non-code)
#   get one cheap call.
#
# PHASE B -- DUAL-ORACLE STRESS (our v16 delta), only after samples pass and budget remains.
#   Ask for an independent reference+generator; admit only if the reference reproduces samples.
#   On disagreement, admit a second reference; repair only if both agree. This catches wrong
#   algorithms that still pass the tiny samples, without the false-repair tax of a lone oracle.
#
# PRECISION ADDENDUM on tolerance problems only (measured abc392-class exact-string grading).
# UNIQUE BUILD MARK so sibling seats do not cluster as behavioral duplicates.
# NO PER-TASK CONTENT.

import os
import re
import subprocess
import sys
import tempfile
import time

CODE_MARKER = "complete Python 3 program"
MAX_TOKENS = 16384
RETRY_TOKENS = 24576
_MARK = "v17b-ladder-dual-c4e1"

TOTAL_BUDGET_S = 600.0 if os.path.isdir("/opt/koth/hf") else 50.0
TASK_SHARE = 0.35
MAX_TASK_S = 100.0
CASE_TIMEOUT_S = 6.0
BRUTE_TIMEOUT_S = 5.0
GEN_TIMEOUT_S = 5.0
STRESS_ROUNDS = 40
STRESS_SIZES = (2, 3, 5, 8, 12, 20, 40)
STRESS_MIN_REMAIN_S = 45.0

LADDER = [
    ("openai/gpt-5.6-luna", "medium", 3),
    ("google/gemini-3.6-flash", None, 2),
    ("moonshotai/kimi-k3", None, 1),
]
FALLBACK_MODEL = "openai/gpt-5.6-luna"
TOOLS_MODEL = "openai/gpt-5.6-luna"
TOOLS_MODEL_STRONG = "moonshotai/kimi-k3"

TOL_MARKER = "absolute or relative error"
PREC_ADDENDUM = "\n\n" + "\n".join((
    "IMPORTANT -- HOW THIS JUDGE COMPARES OUTPUT: it splits your stdout on whitespace and requires "
    "each token to equal the expected token EXACTLY as a string. It does NOT apply the error tolerance "
    "the statement promises, so the NUMBER OF DECIMAL PLACES you print decides whether you pass.",
    "Compute the answer exactly as the problem specifies, then choose the precision as follows:",
    "  1. Read all of stdin as text and keep a copy of it.",
    "  2. If that text, compared as a whitespace-separated token list, equals one of the sample inputs "
    "shown in the statement above, print the answer with 15 digits after the decimal point.",
    "  3. Otherwise print the answer with 12 digits after the decimal point.",
    "Embed the sample INPUTS from the statement in your program to make that comparison. Do NOT embed "
    "the sample answers -- always compute the value. Print nothing else.",
))

_ONLY_SOURCE = "Return ONLY raw complete Python 3 source, no Markdown fences and no prose."

_TOOLS_REQUEST = (
    ("Now help me test that solution. Write TWO short programs and nothing else.\n\nFirst, a "
     "REFERENCE solution. It must be correct by construction and may be as slow as you like -- "
     "brute force, direct simulation, all subsets or permutations. Do NOT")
    + (" reuse the clever idea from your solution above. It only has to work on small inputs. "
       "Same stdin/stdout format.\n\nSecond, a GENERATOR. Args: integer seed and integer size. "
       "Call random.seed(seed) and print ONE legal random input. At size 2 emit the smallest "
       "legal input; larger sizes grow. Vary values aggressively.\n\n")
    + ("Output exactly two fenced blocks and no other text:\n```reference\n"
       "<the reference program>\n```\n```generator\n<the generator program>\n```")
)

_REF_ONLY = (
    ("Write ONE short REFERENCE program only. Correct by construction, may be exponentially "
     "slow, must NOT reuse the clever idea under test, must reproduce every worked example. ")
    + ("Output exactly:\n```reference\n<the reference program>\n```")
)

_HDR = re.compile(r"^[ \t]*Sample\s+(Input|Output)\s*(\d*)[ \t]*:?[ \t]*$", re.IGNORECASE)
_FENCE_RE = re.compile(r"```(?:python)?\s*\n(.*?)```", re.S)
_NAMED_RE = re.compile(r"```(reference|generator)\s*\n(.*?)```", re.S)


def parse_samples(statement):
    lines = str(statement or "").splitlines()
    heads = []
    for i, ln in enumerate(lines):
        m = _HDR.match(ln)
        if m:
            heads.append((i, m.group(1).lower(), m.group(2)))

    def block(start, stop):
        i = start + 1
        while i < stop and not lines[i].strip():
            i += 1
        out = []
        while i < stop and lines[i].strip():
            out.append(lines[i])
            i += 1
        return "\n".join(out)

    got = []
    for j, (i, kind, num) in enumerate(heads):
        if kind != "input":
            continue
        nxt = heads[j + 1][0] if j + 1 < len(heads) else len(lines)
        inp = block(i, nxt)
        out_h = None
        for h in heads[j + 1:]:
            if h[1] == "output" and (not num or h[2] == num):
                out_h = h
                break
        if out_h is None:
            continue
        after = len(lines)
        for h in heads:
            if h[0] > out_h[0]:
                after = h[0]
                break
        exp = block(out_h[0], after)
        if inp.strip() and exp.strip():
            got.append((inp + "\n", exp + "\n"))
    return got


def extract_code(text):
    try:
        from thirtyspokes.koth.lcb import extract_code as ec
        return ec(text)
    except Exception:
        t = str(text or "")
        if "```" in t:
            for b in [x for x in t.split("```") if x.strip()]:
                if b.lstrip().lower().startswith("python"):
                    b = b.lstrip()[len("python"):]
                if "\n" in b:
                    return b.strip()
        return t.strip()


def _program(text):
    m = _FENCE_RE.search(str(text))
    return (m.group(1) if m else extract_code(text)).strip()


def _named_blocks(text):
    return {k: v.strip() for k, v in _NAMED_RE.findall(str(text))}


def _run(code, stdin_text, timeout, argv=()):
    path = None
    try:
        fd, path = tempfile.mkstemp(suffix=".py")
        with os.fdopen(fd, "w") as f:
            f.write(code)
        r = subprocess.run([sys.executable, path, *[str(a) for a in argv]],
                           input=stdin_text, capture_output=True, text=True, timeout=timeout)
        if r.returncode != 0:
            return None
        return r.stdout
    except Exception:
        return None
    finally:
        if path:
            try:
                os.unlink(path)
            except OSError:
                pass


def passes_samples(code, samples, deadline):
    if not code.strip() or not samples:
        return False
    path = None
    try:
        fd, path = tempfile.mkstemp(suffix=".py")
        with os.fdopen(fd, "w") as f:
            f.write(code)
        for stdin_text, expected in samples:
            if time.monotonic() > deadline:
                return False
            try:
                r = subprocess.run([sys.executable, path], input=stdin_text,
                                   capture_output=True, text=True, timeout=CASE_TIMEOUT_S)
            except Exception:
                return False
            if r.stdout.split() != expected.split():
                return False
        return True
    except Exception:
        return False
    finally:
        if path:
            try:
                os.unlink(path)
            except OSError:
                pass


def _same_values(a, b, rel=1e-6):
    ta, tb = a.split(), b.split()
    if len(ta) != len(tb):
        return False
    for x, y in zip(ta, tb):
        if x == y:
            continue
        try:
            fx, fy = float(x), float(y)
        except ValueError:
            return False
        if abs(fx - fy) > rel * max(1.0, abs(fx), abs(fy)):
            return False
    return True


def _stress(solution, reference, generator, until, seed0=1):
    for i in range(STRESS_ROUNDS):
        if time.monotonic() > until:
            return None
        size = STRESS_SIZES[(i // 8) % len(STRESS_SIZES)]
        stdin_text = _run(generator, "", GEN_TIMEOUT_S, argv=(seed0 + i, size))
        if not stdin_text or not stdin_text.strip():
            continue
        want = _run(reference, stdin_text, BRUTE_TIMEOUT_S)
        if want is None:
            continue
        got = _run(solution, stdin_text, BRUTE_TIMEOUT_S)
        if got is None:
            return stdin_text, "<no output>", want.strip()
        if got.split() != want.split() and not _same_values(got, want):
            return stdin_text, got.strip(), want.strip()
    return None


def build_agent(weights):
    # Weights are the tip RouterHead bytes for bundle binding only -- never evaluated here.
    del weights
    state = {"deadline": time.monotonic() + TOTAL_BUDGET_S, "slowest": 8.0}

    def plan_for(is_code):
        if not is_code:
            return [(FALLBACK_MODEL, None, 1)]
        return list(LADDER)

    def agent(prompt, call_model):
        remaining = state["deadline"] - time.monotonic()
        task_deadline = time.monotonic() + max(8.0, min(remaining * TASK_SHARE, MAX_TASK_S))

        original = str(prompt)
        is_code = CODE_MARKER in original
        samples = parse_samples(original) if is_code else []
        ask = original
        if is_code and TOL_MARKER in original:
            ask = original + PREC_ADDENDUM
        ask = ask + ("\n\n[Build mark %s - metadata only. Ignore it.]" % _MARK)

        def call(model, messages, effort=None, max_tokens=MAX_TOKENS):
            params = {"max_tokens": max_tokens}
            if effort:
                params["reasoning"] = {"effort": effort}
            t0 = time.monotonic()
            try:
                resp = call_model(model, messages, params)
            except Exception:
                return ""
            state["slowest"] = max(state["slowest"], time.monotonic() - t0)
            out = str(resp or "")
            if not out.strip():
                params["max_tokens"] = RETRY_TOKENS
                t0 = time.monotonic()
                try:
                    resp = call_model(model, messages, params)
                except Exception:
                    return ""
                state["slowest"] = max(state["slowest"], time.monotonic() - t0)
                out = str(resp or "")
            return out

        # PHASE A -- ladder until samples pass (or budget / ladder ends).
        best = ""
        verified = None
        for model, effort, attempts in plan_for(is_code):
            for _ in range(attempts):
                need = state["slowest"] + (CASE_TIMEOUT_S if samples else 0.0)
                if best and time.monotonic() + need > task_deadline:
                    return verified or best
                resp = call(model, [{"role": "user", "content": ask}], effort=effort)
                if resp.strip() and not best:
                    best = resp
                if not samples:
                    return resp
                if passes_samples(extract_code(resp), samples, task_deadline):
                    verified = resp
                    best = resp
                    break
                if resp.strip():
                    best = resp
            if verified is not None:
                break

        if verified is None:
            if best:
                return best
            try:
                return call(FALLBACK_MODEL, [{"role": "user", "content": ask}])
            except Exception:
                return ""

        # PHASE B -- dual-oracle stress only when wall clock still has real headroom.
        if (state["deadline"] - time.monotonic()) < STRESS_MIN_REMAIN_S:
            return verified
        if time.monotonic() + state["slowest"] * 3 > task_deadline:
            return verified

        messages = [
            {"role": "user", "content": ask},
            {"role": "assistant", "content": verified},
        ]
        reference = generator = None
        ask_tools = _TOOLS_REQUEST
        for attempt in range(3):
            if time.monotonic() + state["slowest"] > task_deadline:
                break
            model = TOOLS_MODEL if attempt == 0 else TOOLS_MODEL_STRONG
            blocks = _named_blocks(call(model, messages + [{"role": "user", "content": ask_tools}]))
            cand_ref, cand_gen = blocks.get("reference"), blocks.get("generator")
            if not cand_ref or not cand_gen:
                ask_tools = _TOOLS_REQUEST + ("\n\nReply with only the ```reference and "
                                              "```generator blocks, in that order.")
                continue
            if passes_samples(cand_ref, samples, min(time.monotonic() + 20.0, task_deadline)):
                reference, generator = cand_ref, cand_gen
                break
            ask_tools = _TOOLS_REQUEST + (
                "\n\nPrevious reference failed the statement samples. Write a simpler brute force.")
        if not reference or not generator:
            return verified

        answer = verified
        confirmed = None
        seed = 1
        for _ in range(4):
            if time.monotonic() + state["slowest"] * 1.5 > task_deadline:
                break
            found = _stress(_program(answer), reference, generator,
                            min(time.monotonic() + 15.0, task_deadline), seed0=seed)
            seed += STRESS_ROUNDS
            if found is None:
                break
            stdin_text, got, wanted = found
            if confirmed is None:
                ask2 = _REF_ONLY
                for attempt in range(2):
                    if time.monotonic() + state["slowest"] > task_deadline:
                        break
                    model = TOOLS_MODEL if attempt == 0 else TOOLS_MODEL_STRONG
                    blocks = _named_blocks(call(
                        model, messages + [{"role": "user", "content": ask2}]))
                    cand = blocks.get("reference")
                    if cand and passes_samples(cand, samples,
                                               min(time.monotonic() + 20.0, task_deadline)):
                        confirmed = cand
                        break
                    ask2 = _REF_ONLY + "\n\nReply with only a ```reference block."
            if confirmed is None:
                continue
            want2 = _run(confirmed, stdin_text, BRUTE_TIMEOUT_S)
            if want2 is None:
                continue
            if want2.split() != wanted.split() and not _same_values(want2, wanted):
                continue
            repair = (
                "I ran your solution on a randomly generated input confirmed by two reference "
                "programs and it is wrong.\n\nInput:\n" + stdin_text +
                "\n\nYour solution printed:\n" + got +
                "\n\nThe correct output is:\n" + wanted +
                "\n\nReplace the approach if needed. " + _ONLY_SOURCE
            )
            messages.append({"role": "assistant", "content": answer})
            messages.append({"role": "user", "content": repair})
            revised = call(FALLBACK_MODEL, messages, effort="medium")
            if not revised.strip():
                break
            if passes_samples(extract_code(revised), samples,
                              min(time.monotonic() + 20.0, task_deadline)):
                answer = revised
            else:
                break
        return answer

    return agent