import marimo as mo

import asyncio
import json
import os
import random
import re
import time

import httpx

BASE_URL = os.environ.get("OPENAI_BASE_URL", "http://localhost:4000/v1")
#MODEL = "openrouter/openai/gpt-oss-120b"
#MODEL = "openai-codex/gpt-5.6-luna"
MODEL = "anthropic/claude-opus-5"

_http = httpx.AsyncClient(
    base_url=BASE_URL,
    headers={"Authorization": "Bearer " + os.environ.get("OPENAI_API_KEY", "not-needed-for-local-openai-compatible-servers")},
    timeout=120.0,
)

async def chat(messages, stats, tools=None):
    """POST one /chat/completions request and return the assistant message dict.

    Ticks the live `stats` model-call/token counters and strips the `<|...|>`
    junk some providers append to tool-call function names.
    """
    payload = {"model": MODEL, "messages": messages}
    if tools:
        payload["tools"] = tools
        payload["parallel_tool_calls"] = True
    response = await _http.post("/chat/completions", json=payload)
    response.raise_for_status()
    body = response.json()
    stats["model_calls"] += 1
    usage = body.get("usage") or {}
    stats["input_tokens"] += usage.get("prompt_tokens") or 0
    stats["output_tokens"] += usage.get("completion_tokens") or 0
    message = body["choices"][0]["message"]
    for tool_call in message.get("tool_calls") or []:
        function = tool_call["function"]
        function["name"] = function["name"].split("<|", 1)[0]
    return message

def compare_door_digit(secret_digit, digit):
    return -1 if digit < secret_digit else (0 if digit == secret_digit else 1)

_POSITION = {"type": "integer", "minimum": 1, "maximum": ROOM_DIGITS, "description": f"keypad position, 1 through {ROOM_DIGITS}"}
_DIGIT = {"type": "integer", "minimum": 0, "maximum": 9, "description": "digit guess, 0 through 9"}
_PACKED = {
    "type": "integer",
    "minimum": 10,
    "maximum": ROOM_DIGITS * 10 + 9,
    "description": f"packed probe PD: tens digit = position 1 through {ROOM_DIGITS}, units digit = digit guess 0 through 9; 24 probes digit 4 at position 2",
}
_CODE_TEXT = {"type": "string", "description": f"the full {ROOM_DIGITS}-digit code"}
_CODE_FINAL = {"type": "string", "description": f'the full {ROOM_DIGITS}-digit code as a string, e.g. "13579"; omit until every position has returned an exact 0'}

_RESULT_LEGEND = "Each probe result compares your digit to the secret digit at that position: -1 means your digit is too low (try higher), +1 means your digit is too high (try lower), 0 means exact match."
_GUESS_LEGEND = "Two wrong guesses lose the game."

def _spec(name, description, properties, required=None):
    """Chat-completions function-tool spec with an object parameter schema."""
    parameters = {"type": "object", "properties": properties}
    if required:
        parameters["required"] = required
    return {"type": "function", "function": {"name": name, "description": description, "parameters": parameters}}

_PROBE_SPEC = _spec(
    "probe",
    f"Compare one digit guess against the secret door keypad digit at a 1-indexed position. {_RESULT_LEGEND}",
    {"position": _POSITION, "digit": _DIGIT},
    required=["position", "digit"],
)
_PROBE_MANY_SPEC = _spec(
    "probe_many",
    f"Compare many digit guesses at once, one comparison per item. {_RESULT_LEGEND}",
    {"probes": {"type": "array", "items": {"type": "object", "properties": {"position": _POSITION, "digit": _DIGIT}, "required": ["position", "digit"]}}},
    required=["probes"],
)
_GUESS_SPEC = _spec(
    "guess",
    f"Submit the full door keypad code to escape the room. {_GUESS_LEGEND}",
    {"code": _CODE_TEXT},
    required=["code"],
)
_EXEC_SPEC = _spec(
    "exec",
    f"Run a batch of probes and/or submit a guess in one call. {_RESULT_LEGEND} {_GUESS_LEGEND} "
    "Omit the guess field entirely while probing; any non-empty guess consumes one of your two attempts.",
    {"probes": {"type": "array", "items": _PACKED}, "guess": _CODE_FINAL},
)
_EXEC_OPS_SPEC = _spec(
    "exec",
    f"Run a batch of ops in order; each op is one probe or one guess. {_RESULT_LEGEND} {_GUESS_LEGEND} "
    "Only send a guess op once every digit is confirmed.",
    {"ops": {"type": "array", "items": {"type": "object", "properties": {"probe": _PACKED, "guess": _CODE_FINAL}, "minProperties": 1, "maxProperties": 1}}},
    required=["ops"],
)
_CHECK_SPEC = _spec(
    "check",
    f"Compare a full {ROOM_DIGITS}-digit trial code against the door code, one comparison per position. {_RESULT_LEGEND} A check is not a guess and never opens the door.",
    {"code": {"type": "string", "description": f"trial code, exactly {ROOM_DIGITS} digits"}},
    required=["code"],
)

def _pack(position, digit):
    """Encode one probe as the two-digit integer PD used by the exec variants."""
    return position * 10 + digit

class _RoomGame:
    """Per-run game state shared by tool handlers and text parsers: probes,
    guesses, live counters, and the terminal guard. Once `over`, every
    further action reports "skipped: game ended" instead of mutating state."""

    def __init__(self, secret, stats, counters):
        self._secret = secret
        self._stats = stats
        self._counters = counters

    @property
    def over(self):
        """True once the room is escaped or both guesses are spent."""
        return self._counters["won"] or self._counters["guesses"] >= 2

    def bump(self, field):
        self._counters[field] += 1
        self._stats[field] += 1

    def try_probe(self, position, digit):
        """Count one comparison and light up the position once it lands exactly."""
        if self.over:
            return "skipped: game ended"
        self.bump("probes")
        if not (1 <= position <= ROOM_DIGITS and 0 <= digit <= 9):
            return "invalid"
        result = compare_door_digit(int(self._secret[position - 1]), digit)
        if result == 0:
            self._stats["solved"].add(position)
        return result

    def submit_guess(self, value):
        """Consume one of two attempts for an exact-length digit-string code.

        Malformed guesses (wrong type, wrong length, non-digits) are rejected
        without consuming an attempt, so a schema-violating placeholder such
        as `guess: 0` cannot burn the game."""
        if self.over:
            return "skipped: game ended"
        code = value.strip() if isinstance(value, str) else None
        if code is None or len(code) != ROOM_DIGITS or not code.isdigit():
            return f"invalid code: pass the full {ROOM_DIGITS}-digit code as a string"
        self.bump("guesses")
        if code == self._secret:
            self._counters["won"] = True
            return "escaped room"
        if self._counters["guesses"] >= 2:
            return "failed to escape: two wrong guesses"
        return "wrong: one guess remaining"

# Tool handlers take (game, arguments) with already-parsed JSON arguments; the
# dispatch loop reports (TypeError, ValueError) escapes as "invalid arguments".
def _handle_probe(game, arguments):
    position = int(arguments.get("position", 0))
    digit = int(arguments.get("digit", -1))
    return f"position={position}, digit={digit}, result={game.try_probe(position, digit)}"

def _handle_probe_many(game, arguments):
    probes = arguments.get("probes")
    if not isinstance(probes, list):
        return 'invalid probes: expected an array of {"position": int, "digit": int} objects'
    results = []
    for item in probes:
        try:
            position = int(item.get("position", 0))
            digit = int(item.get("digit", 0))
        except (AttributeError, TypeError, ValueError):
            game.bump("probes")
            results.append({"probe": item, "result": "invalid"})
            continue
        results.append({"position": position, "digit": digit, "result": game.try_probe(position, digit)})
    return json.dumps(results)

def _handle_guess(game, arguments):
    return game.submit_guess(arguments.get("code"))

def _handle_check(game, arguments):
    code = str(arguments.get("code", "")).strip()
    if not (len(code) == ROOM_DIGITS and code.isdigit()):
        return f"invalid code: pass exactly {ROOM_DIGITS} digits"
    return json.dumps([game.try_probe(position, int(digit)) for position, digit in enumerate(code, 1)])

def _packed_probe(game, value):
    """Run one packed-PD probe; returns -1/0/+1, 'invalid', or a skip notice."""
    if game.over:
        return "skipped: game ended"
    try:
        packed = int(value)
    except (TypeError, ValueError):
        game.bump("probes")
        return "invalid"
    return game.try_probe(packed // 10, packed % 10)

def _handle_exec(game, arguments):
    results = {}
    probes = arguments.get("probes")
    if probes is not None:
        if not isinstance(probes, list):
            return "invalid probes: expected an array of packed PD integers"
        results["probes"] = [{"probe": value, "result": _packed_probe(game, value)} for value in probes]
    if arguments.get("guess") is not None:
        results["guess"] = game.submit_guess(arguments["guess"])
    if not results:
        return "no-op: pass probes and/or guess"
    return json.dumps(results)

def _handle_exec_ops(game, arguments):
    ops = arguments.get("ops")
    if not isinstance(ops, list):
        return 'invalid ops: expected an array of {"probe": PD} or {"guess": code} objects'
    results = []
    for op in ops:
        if isinstance(op, dict) and set(op) == {"probe"}:
            results.append({"probe": op["probe"], "result": _packed_probe(game, op["probe"])})
        elif isinstance(op, dict) and set(op) == {"guess"}:
            results.append({"guess": op["guess"], "result": game.submit_guess(op["guess"])})
        else:
            results.append({"op": op, "result": 'invalid: use {"probe": PD} or {"guess": code}'})
    return json.dumps(results)

def _parse_unicode(game, text):
    """Text-protocol hook: pull 🔍/🔑 command lines out of the assistant reply."""
    observations = []
    for position_text, digit_text in re.findall(r"🔍\s*(\d+)\s*=\s*(\d+)", text):
        position, digit = int(position_text), int(digit_text)
        observations.append(f"🔍{position}={digit} -> {game.try_probe(position, digit)}")
    for code in re.findall(r"🔑\s*([0-9]+)", text):
        observations.append(f"🔑{code} -> {game.submit_guess(code)}")
    if not observations:
        observations.append("No valid commands found. Use lines like 🔍2=7 or 🔑13579.")
    return observations

def _variant(style, render_probes, render_guess, tools=(), parse=None):
    """One benchmark row: probe-style sentence (with a {probe} hole), example
    renderers shared by all prompts, tool (spec, handler) pairs, and an
    optional text-parsing hook for tool-less control channels."""
    return {
        "style": style,
        "render_probes": render_probes,
        "render_guess": render_guess,
        "tools": list(tools),
        "parse": parse,
    }

VARIANTS = {
    "naive": _variant(
        "Use the probe tool to probe one digit, for example {probe}. "
        "Batch: issue one probe tool call for every unknown position in the same reply — never probe only one position per turn.",
        lambda pairs: ", ".join(f"probe(position={p}, digit={d})" for p, d in pairs),
        lambda code: f'guess(code="{code}")',
        tools=[(_PROBE_SPEC, _handle_probe), (_GUESS_SPEC, _handle_guess)],
    ),
    "optimized": _variant(
        "Use the probe_many tool to batch probes, for example {probe}. "
        "Batch: cover every unknown position in one probe_many call per reply.",
        lambda pairs: "probe_many(probes=[" + ", ".join(f'{{"position": {p}, "digit": {d}}}' for p, d in pairs) + "])",
        lambda code: f'guess(code="{code}")',
        tools=[(_PROBE_MANY_SPEC, _handle_probe_many), (_GUESS_SPEC, _handle_guess)],
    ),
    "exec": _variant(
        "Use the exec tool to batch probes as packed PD integers, tens digit = position and units digit = digit guess; "
        "for example {probe} probes digit 7 at position 2. "
        "Batch: cover every unknown position in one exec call per reply. "
        "While probing, omit the guess field entirely — a placeholder guess such as 0 wastes one of your two attempts.",
        lambda pairs: "exec(probes=[" + ", ".join(str(_pack(p, d)) for p, d in pairs) + "])",
        lambda code: f'exec(guess="{code}")',
        tools=[(_EXEC_SPEC, _handle_exec)],
    ),
    "exec_ops": _variant(
        "Use the exec tool to batch ops; each probe op is a packed PD integer, tens digit = position and units digit = digit guess; "
        "for example {probe} probes digit 7 at position 2. "
        "Batch: cover every unknown position's probe op in one exec call per reply. "
        "Only send a guess op once every digit is confirmed.",
        lambda pairs: "exec(ops=[" + ", ".join(f'{{"probe": {_pack(p, d)}}}' for p, d in pairs) + "])",
        lambda code: f'exec(ops=[{{"guess": "{code}"}}])',
        tools=[(_EXEC_OPS_SPEC, _handle_exec_ops)],
    ),
    "vector": _variant(
        "Use the check tool to compare one full trial code against the door code, one -1/0/+1 result per position; "
        "for example {probe} probes digit 7 at position 2 and digit 4 everywhere else. "
        "Checks are probes, not guesses. Batch: a single check probes all positions at once.",
        lambda pairs: 'check(code="' + "".join(str(dict(pairs).get(p, 4)) for p in range(1, ROOM_DIGITS + 1)) + '")',
        lambda code: f'guess(code="{code}")',
        tools=[(_CHECK_SPEC, _handle_check), (_GUESS_SPEC, _handle_guess)],
    ),
    "unicode": _variant(
        "Use 🔍<position>=<digit> to probe one digit, for example {probe}. "
        "Batch: include one 🔍 line for every unknown position in each reply.",
        lambda pairs: ", ".join(f"🔍{p}={d}" for p, d in pairs),
        lambda code: f"🔑{code}",
        parse=_parse_unicode,
    ),
}

def _instructions(variant):
    """Shared system prompt; only the control-channel sentences differ per variant."""
    tool_mode = bool(variant["tools"])
    return (
        f"You are trapped in an escape room and must escape this room by unlocking the {ROOM_DIGITS}-digit door keypad. "
        "Each digit is 0 through 9. "
        + ("You have tools. " if tool_mode else "You have no tools. Instead, write plain-text command lines only. ")
        + variant["style"].format(probe=variant["render_probes"]([(2, 7)])) + " "
        "In every reply, probe all unknown positions in parallel, binary-searching each digit on the -1/+1 feedback. "
        "A result of -1 means probe a higher digit there next; +1 means probe a lower digit next; 0 pins that position. "
        "Track each position independently. Never guess until every position has returned an exact 0 result. "
        "Each round, move every unresolved position to its next binary-search midpoint; "
        "keep confirmed digits fixed and never repeat a comparison you have already made. "
        f"When every digit is known exactly, submit the full code, for example {variant['render_guess']('13579')}. "
        "You may only submit two guesses; a second wrong guess fails to escape. "
        "Do not explain your reasoning. "
        + ("Output only tool calls. Do not stop until your guess returns win or lose." if tool_mode else "Output only command lines.")
    )

def _kickoff(variant):
    """Shared opening user message with a variant-rendered all-positions example."""
    example = variant["render_probes"]([(position, 4) for position in range(1, ROOM_DIGITS + 1)])
    return (
        f"You must escape this room. Unlock the hidden {ROOM_DIGITS}-digit door keypad code. "
        f"Start by probing all positions in one reply, for example: {example}. "
        "Do not guess until observations show an exact digit for every position."
    )

async def run_room(name, secret, stats):
    """One escape attempt. Mutates `stats` as it goes so the dashboard can watch it live."""
    variant = VARIANTS[name]
    counters = {"tool_calls": 0, "probes": 0, "guesses": 0, "won": False}
    usage_before = (stats["model_calls"], stats["input_tokens"], stats["output_tokens"])

    game = _RoomGame(secret, stats, counters)
    tools = [spec for spec, _ in variant["tools"]]
    dispatch = {spec["function"]["name"]: handler for spec, handler in variant["tools"]}
    messages = [
        {"role": "system", "content": _instructions(variant)},
        {"role": "user", "content": _kickoff(variant)},
    ]
    final_output = ""
    lost = False
    trace = []

    def report(error=""):
        return {
            "won": counters["won"],
            "error": error,
            "model_calls": stats["model_calls"] - usage_before[0],
            "tool_calls": counters["tool_calls"],
            "probes": counters["probes"],
            "guesses": counters["guesses"],
            "input_tokens": stats["input_tokens"] - usage_before[1],
            "output_tokens": stats["output_tokens"] - usage_before[2],
            "final_output": final_output,
            "trace": trace,
            "messages": messages,
            "secret": secret,
        }

    for turn_index in range(1, MAX_TURNS + 1):
        try:
            message = await chat(messages, stats, tools=tools or None)
        except httpx.HTTPError as api_error:
            # A transport failure or error status ends this room; it does not end
            # the benchmark, so record it and let the other variants keep running.
            trace.append({
                "turn": turn_index,
                "assistant": "",
                "thinking": "",
                "observations": [f"api error: {api_error}"],
                "won": False,
                "lost": True,
            })
            return report(f"{type(api_error).__name__}: {api_error}")
        tool_calls = message.get("tool_calls") or []
        final_output = message.get("content") or ""
        thinking = message.get("reasoning_content") or message.get("reasoning") or ""
        assistant_message = {"role": "assistant", "content": final_output}
        if tool_calls:
            assistant_message["tool_calls"] = [
                {
                    "id": tool_call.get("id"),
                    "type": "function",
                    "function": {
                        "name": tool_call["function"]["name"],
                        "arguments": tool_call["function"].get("arguments") or "{}",
                    },
                }
                for tool_call in tool_calls
            ]
        messages.append(assistant_message)

        observations = []
        for tool_call in tool_calls:
            game.bump("tool_calls")
            function = tool_call["function"]
            handler = dispatch.get(function["name"])
            try:
                arguments = json.loads(function.get("arguments") or "{}")
            except json.JSONDecodeError:
                arguments = None
            if game.over:
                result = "skipped: game ended"
            elif handler is None:
                result = f"unknown tool: {function['name']}"
            elif not isinstance(arguments, dict):
                result = "invalid arguments: expected a JSON object"
            else:
                try:
                    result = handler(game, arguments)
                except (TypeError, ValueError):
                    result = "invalid arguments"
            observations.append(f"{function['name']}({function.get('arguments') or ''}) -> {result}")
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.get("id"),
                "content": result,
            })
        if variant["parse"] is not None:
            observations.extend(variant["parse"](game, final_output))

        lost = not counters["won"] and counters["guesses"] >= 2
        trace.append({
            "turn": turn_index,
            "assistant": final_output,
            "thinking": thinking,
            "observations": observations,
            "won": counters["won"],
            "lost": lost,
        })
        if counters["won"] or lost:
            break
        if not tool_calls:
            if variant["parse"] is None:
                # The model stopped issuing tool calls, so the run is over either way.
                break
            messages.append({"role": "user", "content": "Observations:\n" + "\n".join(observations)})

    return report()

_STATE_COLORS = {"waiting": "#94a3b8", "running": "#6366f1", "escaped": "#10b981", "trapped": "#f59e0b"}

def new_room_stats(name):
    """Live counters for one variant, mutated in place by run_room while it works."""
    return {
        "variant": name,
        "state": "waiting",
        "trial": 0,
        "trials_done": 0,
        "escaped": 0,
        "solved": set(),
        "model_calls": 0,
        "tool_calls": 0,
        "probes": 0,
        "guesses": 0,
        "input_tokens": 0,
        "output_tokens": 0,
        "started": time.perf_counter(),
        "finished": None,
    }

def _compact(count):
    return f"{count / 1000:.1f}k" if count >= 1000 else str(count)

def _room_card(stats, digits, trials):
    # Progress counts digits pinned down, not just trials finished, so the bar keeps moving mid-room.
    solved_now = len(stats["solved"])
    percent = 100.0 * (stats["trials_done"] * digits + solved_now) / (digits * trials)
    color = _STATE_COLORS[stats["state"]]
    keypad = "".join("▮" if position in stats["solved"] else "▯" for position in range(1, digits + 1))
    end = time.perf_counter() if stats["finished"] is None else stats["finished"]
    line = " · ".join([
        f"room {max(stats['trial'], 1)}/{trials} {keypad}",
        f"escaped {stats['escaped']}",
        f"{stats['model_calls']} model",
        f"{stats['tool_calls']} tool",
        f"{stats['probes']} probes",
        f"{stats['guesses']} guesses",
        f"{_compact(stats['input_tokens'])}/{_compact(stats['output_tokens'])} tok",
        f"{end - stats['started']:.1f}s",
    ])
    return (
        '<div style="margin-bottom:.8rem;font-family:ui-monospace,SFMono-Regular,Menlo,monospace">'
        '<div style="display:flex;justify-content:space-between;font-size:.8rem;font-weight:600">'
        f'<span>{stats["variant"]}</span><span style="color:{color}">{stats["state"]}</span></div>'
        '<div style="height:.5rem;margin:.3rem 0;border-radius:999px;background:rgba(127,127,127,.18);overflow:hidden">'
        f'<div style="height:100%;width:{percent:.1f}%;background:{color};transition:width .2s linear"></div></div>'
        f'<div style="font-size:.72rem;opacity:.7">{line}</div></div>'
    )

def render_room_dashboard(live, digits, trials):
    """Stacked progress bars, one per variant, with live escape-room counters."""
    return mo.Html("".join(_room_card(stats, digits, trials) for stats in live.values()))