"""Bridge mode — real chat-plan, then a real ``/ai <agent> !task`` build.

Where ``loop.py`` (direct mode) shortcuts the build through ``completion`` +
``vm.run``, this drives the *actual* product path end-to-end on a real encrypted
room, the way ``bench-sandbox.py`` does:

  1. boot the relay, spawn a real ``AgentBridge`` agent (the builder), connect an
     owner/referee client and a planner teammate client;
  2. BRIEF + DELIBERATE happen as real chat frames — the planner proposes and the
     real agent is asked (via the real ``/ai <agent>`` *chat* path) to confirm the
     approach;
  3. IMPLEMENT grants drive and issues ``/ai <agent> !<task>`` — the real agent
     turns it into shell, clears the destructive guard, and injects ``_sbx:input``
     keystrokes, which we capture off the wire;
  4. GRADE materializes the captured commands, reads the produced source, and runs
     it against the held-out hidden tests in the VM.

No ``cmd_chat`` code is modified — this only *uses* the tool. The reusable wire
primitives (Owner/grant/task/collect, execute, spawn_agent) are imported from
``bench-sandbox.py`` so the bridge path is byte-identical to the sandbox bench.
"""

from __future__ import annotations

import asyncio
import importlib
import json
import re
import shutil
import sys
import time
from pathlib import Path

import websockets

from . import infer, scoring, transcript as T
from .. import completion
from .challenge import Challenge
from .scoring import EventOutcome
from .team import Team
from .vm import VM

_SCRIPTS = Path(__file__).resolve().parents[2]
if str(_SCRIPTS) not in sys.path:
    sys.path.insert(0, str(_SCRIPTS))
_sb = importlib.import_module("bench-sandbox")   # reuse proven wire primitives

_FN_RE = re.compile(r"assert\s+([A-Za-z_]\w*)\s*\(")


def _fn_name(public_assert: str) -> str:
    m = _FN_RE.search(public_assert or "")
    return m.group(1) if m else "solve"


async def _collect_chat(owner, ws, agent: str, deadline: float,
                        quiet: float = 3.0) -> str:
    """Read the agent's plain chat reply (not a sandbox outcome) until a quiet
    gap. Skips control/`_sbx` frames and the audit line."""
    parts: list[str] = []
    while time.time() < deadline:
        try:
            raw = await asyncio.wait_for(ws.recv(),
                                         timeout=min(quiet, deadline - time.time()))
        except asyncio.TimeoutError:
            if parts:
                break
            continue
        except websockets.ConnectionClosed:
            break
        try:
            data = json.loads(raw)
        except json.JSONDecodeError:
            continue
        if data.get("type") != "message":
            continue
        dec = owner.decrypt_message(data.get("data", {}))
        if dec.get("username") != agent:
            continue
        text = _sb._strip_think(dec.get("text", ""))
        if text.startswith('{"_'):     # control / sbx frame
            continue
        parts.append(text)
    return "\n".join(parts).strip()


async def _agent_alive(owner, ws, agent: str, window: float = 1.5) -> bool:
    """Fast-fail liveness probe. The relay broadcasts a `user_left` + fresh
    `roster` whenever a client drops (e.g. the agent killed by a 1011 keepalive
    timeout). Drain whatever is already queued on the owner ws for a short
    window; if a roster snapshot arrives without the agent, it's gone — so the
    caller can abort instead of waiting out the full step deadline."""
    deadline = time.time() + window
    alive = True
    while time.time() < deadline:
        try:
            raw = await asyncio.wait_for(ws.recv(),
                                         timeout=max(0.05, deadline - time.time()))
        except (asyncio.TimeoutError, websockets.ConnectionClosed):
            break
        try:
            data = json.loads(raw)
        except json.JSONDecodeError:
            continue
        if data.get("type") == "roster":
            names = {u.get("username") for u in data.get("users", [])}
            alive = agent in names
    return alive


def _find_source(cwd: str | None, fn: str) -> str:
    """Read the python source the agent built. Prefer solution.py, else any .py
    that defines the target function, else the largest .py."""
    if not cwd:
        return ""
    pys = list(Path(cwd).rglob("*.py"))
    if not pys:
        return ""
    named = [p for p in pys if p.name == "solution.py"]
    if named:
        return named[0].read_text()
    for p in pys:
        try:
            if f"def {fn}" in p.read_text():
                return p.read_text()
        except Exception:  # noqa: BLE001
            continue
    return max(pys, key=lambda p: p.stat().st_size).read_text()


def run_bridge_event(team: Team, challenge: Challenge, *,
                     host: str = "127.0.0.1", port: int = 4699,
                     password: str = "olympics-pass",
                     ollama: str = "http://127.0.0.1:11434",
                     code_model: str | None = None,
                     out_dir: str = "/tmp/hh-olympics/runs",
                     log_dir: str = "/tmp/hh-olympics",
                     step_timeout: float = 180.0, seed: int = 0,
                     profile: str = "balanced", agent_chat_confirm: bool = False,
                     progress=None) -> dict:
    """Run one real chat-plan -> !task-build event. Returns the score dict."""
    py = sys.executable
    driver = team.driver()
    planner = next((m for m in team.members if m is not driver), driver)
    build_model = code_model or driver.model
    agent_name = driver.model      # the agent joins under its model tag
    fn = _fn_name(challenge.public_tests[0] if challenge.public_tests else "")

    manifest = {"models": team.models, "seed": seed, "topology": team.topology,
                "framing": team.framing, "room_substrate": "real-bridge",
                "build_model": build_model, "agent": agent_name, "fn": fn,
                "profile": profile, "challenge_meta": challenge.meta,
                "created": time.strftime("%Y-%m-%dT%H:%M:%S")}
    tx = T.Transcript(team.id, challenge.id, manifest)
    vm = VM(challenge.lang)
    Path(log_dir).mkdir(parents=True, exist_ok=True)

    def emit(p):
        if progress:
            progress(p)

    # ── boot relay + spawn the real agent ────────────────────────────────
    srv_log = open(Path(log_dir) / f"bridge-server-{port}.log", "w")
    srv = _sb.subprocess.Popen(
        [py, "cmd_chat.py", "serve", host, str(port),
         "--password", password, "--no-tls"],
        cwd=str(_sb.REPO), stdout=srv_log, stderr=_sb.subprocess.STDOUT)
    agent = None
    alog = None
    tokens = {"in": 0, "out": 0}
    penalties: list[dict] = []
    t0 = time.time()
    try:
        if not _sb._wait_port(host, port, time.time() + 30):
            raise RuntimeError("relay never bound")
        emit("server up")
        alog = open(Path(log_dir) / f"bridge-agent-{port}.log", "w")
        agent = _sb.spawn_agent(py, host, port, password, agent_name,
                                build_model, alog)
        emit(f"agent '{agent_name}' spawned (code-model={build_model})")

        outcome = asyncio.run(_drive(
            tx, team, challenge, planner, driver, agent_name, fn,
            host, port, password, ollama, step_timeout, vm, tokens,
            penalties, agent_chat_confirm, build_model, emit))
    finally:
        if agent is not None:
            agent.terminate()
            try:
                agent.wait(timeout=10)
            except _sb.subprocess.TimeoutExpired:
                agent.kill()
        if alog:
            alog.close()
        srv.terminate()
        try:
            srv.wait(timeout=10)
        except _sb.subprocess.TimeoutExpired:
            srv.kill()
        srv_log.close()

    outcome.wall_clock_s = time.time() - t0
    run_dir = Path(out_dir) / f"{team.id}__{challenge.id}__bridge"
    tx_path = run_dir / "transcript.json"
    tx.save(tx_path)
    score = scoring.score_event(outcome, profile=profile, budget={"max_rounds": 1})
    score["transcript"] = str(tx_path)
    score["room_substrate"] = "real-bridge"
    (run_dir / "score.json").write_text(json.dumps(score, indent=2))
    return score


async def _drive(tx, team, challenge, planner, driver, agent_name, fn,
                 host, port, password, ollama, step_to, vm, tokens,
                 penalties, agent_chat_confirm, build_model, emit) -> EventOutcome:
    owner = _sb.Owner(host, port, password)
    owner.srp_authenticate()
    # a planner teammate client (real chat voice alongside the real agent)
    from cmd_chat.client.client import Client  # noqa: E402
    mate = Client(host, port, username=planner.name, password=password, no_tls=True)
    mate.srp_authenticate()

    owner_url = f"{owner.ws_url}/ws/chat?user_id={owner.user_id}&ws_token={owner.ws_token}"
    mate_url = f"{mate.ws_url}/ws/chat?user_id={mate.user_id}&ws_token={mate.ws_token}"

    public_passed = False
    submitted = False
    correct = False
    completion_src = ""
    defect_class: str | None = None
    model_baseline: bool | None = None

    async with websockets.connect(owner_url) as ws, \
            websockets.connect(mate_url) as mws:
        if not await owner.wait_for_agent(ws, agent_name, time.time() + step_to):
            penalties.append({"kind": "agent_offline", "detail": agent_name})
            return EventOutcome(team.id, challenge.id, False, False, 0, 0.0,
                                tokens, False, None, penalties)

        async def mate_send(text):
            await mws.send(mate.room_fernet.encrypt(text.encode()).decode())

        # ── BRIEF ────────────────────────────────────────────────────────
        tx.phase_change("BRIEF")
        brief = challenge.brief()
        await owner._send(ws, brief)
        tx.add(T.KIND_MESSAGE, "BRIEF", "referee", {"text": brief}, role="referee")
        emit("brief posted")

        # ── DELIBERATE: planner proposes (real chat), agent confirms (real
        #    /ai chat path) ────────────────────────────────────────────────
        tx.phase_change("DELIBERATE")
        # blocking inference must run off the event loop, else the websocket
        # keepalive can't answer the server's ping and the room drops us (1011).
        plan_msg = await asyncio.to_thread(
            infer.chat,
            planner.model, planner.system_prompt(),
            [{"role": "user", "content":
              f"{brief}\n\nPropose, in 2-3 short lines, the approach and the "
              f"exact function signature (name it {fn}). End with PLAN-LOCKED."}],
            host=ollama)
        tokens["in"] += plan_msg.tokens.get("in", 0)
        tokens["out"] += plan_msg.tokens.get("out", 0)
        ptext = plan_msg.text if plan_msg.ok else f"[infer error: {plan_msg.error}]"
        await mate_send(ptext)
        tx.add(T.KIND_MESSAGE, "DELIBERATE", planner.name,
               {"text": ptext}, role=planner.role, model=planner.model)
        emit("planner proposed")

        # Optionally ask the REAL agent to weigh in via the real /ai *chat* path.
        # HARNESS FINDING (default off): the agent's streaming chat path
        # (_answer -> _stream_reply -> OllamaProvider.stream) starves its own
        # asyncio loop past the websockets 20s keepalive window under local CPU
        # inference, so the relay drops the agent with `1011 keepalive ping
        # timeout` and the build never happens. The non-streaming `!task` path
        # (_run_in_sandbox -> to_thread(complete)) does NOT have this defect, so
        # by default we keep planning to real planner room-frames and let the
        # agent's only inference be the robust `!task` build below.
        if agent_chat_confirm:
            await owner._send(
                ws, f"/ai {agent_name} Teammate proposed: {ptext[:300]} . In one "
                    f"line, confirm the function name {fn} and the core idea.")
            agent_reply = await _collect_chat(owner, ws, agent_name,
                                              time.time() + step_to)
            if not await _agent_alive(owner, ws, agent_name):
                penalties.append({"kind": "agent_dropped",
                                  "detail": "streaming /ai chat keepalive timeout"})
            tx.add(T.KIND_MESSAGE, "DELIBERATE", agent_name,
                   {"text": agent_reply or "(no chat reply)"}, role="builder",
                   model=agent_name)
            emit("agent chat reply captured")
        else:
            tx.add(T.KIND_MESSAGE, "DELIBERATE", "referee",
                   {"text": "agent /ai chat-confirm skipped (streaming chat path "
                            "drops the agent via 1011 keepalive timeout under CPU "
                            "inference); building via the robust !task path",
                    "note": "harness-workaround"}, role="referee")
            emit("agent chat-confirm skipped (streaming-path keepalive bug)")

        # ── IMPLEMENT: real /ai !task build ───────────────────────────────
        await owner.grant(ws, agent_name)
        await asyncio.sleep(0.6)
        tx.acl("IMPLEMENT", {"_perm": "acl", "drivers": [agent_name]})
        tx.phase_change("IMPLEMENT", note=f"driver={agent_name}")
        example = challenge.public_tests[0] if challenge.public_tests else ""
        task = (f"Create a file named solution.py that defines a python function "
                f"named {fn} solving this: {challenge.text.strip()} "
                f"It must satisfy: {example} . "
                f"Write only the function definition in solution.py — no tests, "
                f"no example calls, no printing.")
        await owner.task(ws, agent_name, task)
        res = await owner.collect(ws, agent_name, time.time() + step_to)
        tx.tool("IMPLEMENT", agent_name,
                {"label": "sbx-build", "outcome": res["outcome"],
                 "sbx_frames": res["sbx"], "commands": res["commands"]},
                role="builder")
        emit(f"build outcome={res['outcome']} sbx={res['sbx']} "
             f"cmds={len(res['commands'])}")

        if res["outcome"] == "destructive_gated":
            tx.guard("IMPLEMENT", agent_name, {"gated": res["message"][:80]})
            penalties.append({"kind": "destructive", "detail": "build gated"})

        # ── GRADE: materialize captured cmds, read source, run hidden ─────
        if res["commands"]:
            ex = await asyncio.to_thread(_sb.execute, res["commands"], 30.0)
            src = _find_source(ex.get("cwd"), fn)
            tx.tool("GRADE", "referee",
                    {"label": "materialize", "ran": ex["ran"],
                     "skipped": ex["skipped"], "found_source": bool(src),
                     "out": (ex["out"] or "")[-300:]})
            if src:
                completion_src = src
                # Phase 3a — does the as-injected source even compile? A syntax/
                # indentation failure here is the product keystroke path mangling
                # the code (e.g. _extract_commands stripping heredoc indentation),
                # not a model error. Attribute it so a matrix can separate them.
                compiles = True
                try:
                    compile(src, "<sbx>", "exec")
                except (SyntaxError, IndentationError) as e:
                    compiles = False
                    defect_class = f"product-path-defect:{type(e).__name__}"
                pub = await asyncio.to_thread(
                    vm.run, challenge.public_program(src), 30.0, label="public")
                public_passed = pub.ok
                tx.tool("TEST", "referee",
                        {"label": "public", "rc": pub.rc, "ok": pub.ok,
                         "out": pub.out[-300:]})
                hid = await asyncio.to_thread(
                    vm.run, challenge.hidden_program(src), 30.0,
                    label="hidden-grade")
                correct = hid.ok
                tx.tool("SUBMIT", "referee",
                        {"label": "hidden-grade", "rc": hid.rc, "ok": hid.ok,
                         "out": hid.out[-300:]})
                if not correct and not compiles:
                    penalties.append({"kind": "product_path_defect",
                                      "detail": defect_class})
            if ex.get("cwd"):
                shutil.rmtree(ex["cwd"], ignore_errors=True)
        submitted = bool(completion_src)

        # Phase 3b — control: the build model's *true* coding ability via the raw
        # completion path (direct mode, no keystroke stripping). The delta between
        # this and `correct` above quantifies exactly what the product path costs.
        ctrl = await asyncio.to_thread(
            completion.complete, build_model, challenge.gen_prompt(),
            challenge.stop_tokens(), host=ollama, timeout=180.0)
        if ctrl.ok and ctrl.text.strip():
            base = await asyncio.to_thread(
                vm.run, challenge.hidden_program(ctrl.text), 30.0,
                label="model-baseline")
            model_baseline = base.ok
            tx.tool("GRADE", "referee",
                    {"label": "model-baseline", "ok": base.ok, "rc": base.rc,
                     "model": build_model, "out": base.out[-300:]})
            emit(f"model-baseline (direct gen) correct={base.ok}")

        await owner.revoke(ws)
        tx.add(T.KIND_TOOL, "SUBMIT", "referee",
               {"label": "artifact", **vm.freeze()})

    rtg = 1 if public_passed else None
    return EventOutcome(team.id, challenge.id, submitted, correct,
                        rounds=1, wall_clock_s=0.0, tokens=tokens,
                        public_passed=public_passed, rounds_to_green=rtg,
                        penalties=penalties, defect_class=defect_class,
                        model_baseline_correct=model_baseline)