import asyncio
import concurrent.futures
import contextlib
import datetime
import json
import threading

import pytest

from tui_gateway import server
from tui_gateway import ws as ws_mod




def _run_disconnect(monkeypatch, seed):
    """Drive handle_ws to its disconnect `finally`, seeding sessions against the
    live WSTransport the moment it exists. Returns nothing; inspect _sessions."""
    # Disable the grace-reap Timer: detached sessions normally schedule a
    # threading.Timer via _schedule_ws_orphan_reap, which would outlive the test
    # and fire _reap during interpreter teardown — touching _sessions/DB and
    # producing spurious post-run errors under the per-file CI runner. Grace=0
    # short-circuits the Timer (see _schedule_ws_orphan_reap) so the test leaves
    # no lingering thread.
    monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0)

    # Mirror the real _finalize_session chokepoint: it is the single place that
    # closes the slash-worker (#38095). Stub it but keep that behavior so the
    # disconnect-reap path still exercises worker teardown.
    def _fake_finalize(s, end_reason="tui_close"):
        w = s.get("slash_worker")
        if w:
            w.close()

    monkeypatch.setattr(server, "_finalize_session", _fake_finalize)

    created = []
    real_transport = ws_mod.WSTransport
    monkeypatch.setattr(
        ws_mod, "WSTransport",
        lambda ws, loop, **kw: created.append(real_transport(ws, loop, **kw)) or created[-1],
    )

    class FakeWS:
        async def accept(self):
            pass

        async def send_text(self, line):
            pass

        async def receive_text(self):
            seed(created[0])  # transport now exists; attach it to sessions
            raise ws_mod._WebSocketDisconnect()

        async def close(self):
            pass

    asyncio.run(ws_mod.handle_ws(FakeWS()))


def test_ws_disconnect_reaps_flagged_session_and_closes_worker(monkeypatch):
    closed = []

    class FakeWorker:
        def close(self):
            closed.append(True)

    server._sessions.clear()
    try:
        _run_disconnect(
            monkeypatch,
            lambda t: server._sessions.update(
                flagged={
                    "transport": t,
                    "close_on_disconnect": True,
                    "slash_worker": FakeWorker(),
                    "session_key": "k",
                }
            ),
        )
        assert "flagged" not in server._sessions
        assert closed == [True]
    finally:
        server._sessions.clear()




def test_ws_connection_registers_then_disconnect_unregisters_live_transport(monkeypatch):
    """A connected client must be tracked in the live-transport registry so a
    session-less global broadcast (skin.changed from the background watcher)
    reaches it, and dropped on disconnect so no stale write targets a dead peer.
    This is the WS half of the cross-surface live-theme fix."""
    server._sessions.clear()
    server._live_transports.clear()
    seen = {}
    try:
        _run_disconnect(
            monkeypatch,
            lambda t: seen.__setitem__("registered", t in server._live_transports),
        )
        # Seeded at receive_text time — i.e. after gateway.ready registered it.
        assert seen["registered"] is True
        # handle_ws's finally must have unregistered it.
        assert not server._live_transports
    finally:
        server._sessions.clear()
        server._live_transports.clear()


def test_ws_disconnect_releases_wake_word_owner(monkeypatch):
    released = []
    created = []
    monkeypatch.setattr(
        server,
        "_release_wake_for_transport",
        lambda transport: released.append(transport) or True,
    )

    _run_disconnect(monkeypatch, lambda transport: created.append(transport))

    assert released == created






def test_ws_ready_advertises_heartbeat_and_ping_is_inline(monkeypatch):
    sent = []
    inbound = iter(
        [
            json.dumps(
                {
                    "jsonrpc": "2.0",
                    "id": "heartbeat-1",
                    "method": "gateway.ping",
                    "params": {},
                }
            )
        ]
    )
    monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0)

    class FakeWS:
        async def accept(self):
            pass

        async def send_text(self, line):
            sent.append(json.loads(line))

        async def receive_text(self):
            try:
                return next(inbound)
            except StopIteration:
                raise ws_mod._WebSocketDisconnect()

        async def close(self):
            pass

    asyncio.run(ws_mod.handle_ws(FakeWS()))

    ready = sent[0]["params"]
    assert ready["type"] == "gateway.ready"
    assert ready["payload"]["heartbeat"] is True
    assert sent[1] == {
        "jsonrpc": "2.0",
        "result": {"ok": True},
        "id": "heartbeat-1",
    }


def _slow_dispatch_harness(monkeypatch):
    """handle_ws over a scripted FakeWS whose ``slow`` RPC blocks inside dispatch() until released.

    Returns (inbound queue, sent frames, event log, release Event). Push ``None`` to disconnect."""
    monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0)
    sent, log, release = [], [], threading.Event()

    def fake_dispatch(req, transport):
        method = req.get("method")
        log.append(f"dispatch:{method}")
        if method in {"slow", "session.context_breakdown"}:
            release.wait(5)
        log.append(f"done:{method}")
        return {"jsonrpc": "2.0", "id": req.get("id"), "result": {"method": method}}

    def fake_close_sessions(transport, end_reason):
        log.append("teardown")
        return 0, 0

    monkeypatch.setattr(server, "dispatch", fake_dispatch)
    monkeypatch.setattr(server, "_close_sessions_for_transport", fake_close_sessions)
    inbound: asyncio.Queue = asyncio.Queue()

    class FakeWS:
        async def accept(self):
            pass

        async def send_text(self, line):
            sent.append(json.loads(line))

        async def receive_text(self):
            frame = await inbound.get()
            if frame is None:
                raise ws_mod._WebSocketDisconnect()
            return json.dumps(frame)

        async def close(self):
            pass

    return FakeWS(), inbound, sent, log, release


async def _wait_for(predicate, timeout=2.0):
    deadline = asyncio.get_running_loop().time() + timeout
    while not predicate():
        if asyncio.get_running_loop().time() > deadline:
            return False
        await asyncio.sleep(0.01)
    return True


def _rpc(req_id, method):
    return {"jsonrpc": "2.0", "id": req_id, "method": method, "params": {}}


def test_ws_ping_is_answered_while_an_earlier_rpc_blocks_dispatch(monkeypatch):
    """#108325: a handler blocked for minutes (lock wait behind a long compaction, GIL-heavy turn) must not
    starve gateway.ping — the client's 45s heartbeat deadline would otherwise tear down a busy but healthy
    backend. Non-ping RPCs keep their serial arrival order."""
    ws, inbound, sent, log, release = _slow_dispatch_harness(monkeypatch)
    ids = lambda: [f.get("id") for f in sent if "id" in f]

    async def scenario():
        task = asyncio.create_task(ws_mod.handle_ws(ws))
        try:
            for frame in (_rpc("r1", "slow"), _rpc("r2", "fast"), _rpc("heartbeat-1", "gateway.ping")):
                inbound.put_nowait(frame)
            ping_answered = await _wait_for(lambda: "heartbeat-1" in ids(), timeout=1.0)
            assert ping_answered, f"gateway.ping went unanswered while dispatch was busy; sent ids={ids()}"
            assert "done:slow" not in log and "dispatch:fast" not in log
        finally:
            release.set()
        assert await _wait_for(lambda: {"r1", "r2"} <= set(ids()))
        inbound.put_nowait(None)
        await asyncio.wait_for(task, 5)

    asyncio.run(scenario())
    assert log.index("done:slow") < log.index("dispatch:fast")
    assert ids().index("r1") < ids().index("r2")


def test_ws_context_breakdown_does_not_block_a_later_control_read(monkeypatch):
    """Regression for #140667: the expensive read-only context meter may finish after controls."""
    ws, inbound, sent, log, release = _slow_dispatch_harness(monkeypatch)
    ids = lambda: [frame.get("id") for frame in sent if "id" in frame]

    async def scenario():
        task = asyncio.create_task(ws_mod.handle_ws(ws))
        try:
            inbound.put_nowait(_rpc("context", "session.context_breakdown"))
            inbound.put_nowait(_rpc("control", "session.control.read"))
            assert await _wait_for(lambda: "control" in ids())
            assert "context" not in ids()
        finally:
            release.set()
        assert await _wait_for(lambda: "context" in ids())
        inbound.put_nowait(None)
        await asyncio.wait_for(task, 5)

    asyncio.run(scenario())
    assert log.index("done:session.context_breakdown") > log.index("done:session.control.read")


def test_ws_disconnect_teardown_waits_for_in_flight_dispatch(monkeypatch):
    """A client that drops while a handler runs must not have its sessions torn down under that handler, and
    frames it sent before dropping are still dispatched (as the serial read loop did)."""
    ws, inbound, _sent, log, release = _slow_dispatch_harness(monkeypatch)

    async def scenario():
        task = asyncio.create_task(ws_mod.handle_ws(ws))
        inbound.put_nowait(_rpc("r1", "slow"))
        assert await _wait_for(lambda: "dispatch:slow" in log)
        inbound.put_nowait(_rpc("r2", "queued"))
        inbound.put_nowait(None)
        await asyncio.sleep(0.1)
        assert "teardown" not in log
        release.set()
        await asyncio.wait_for(task, 5)

    asyncio.run(scenario())
    assert log == ["dispatch:slow", "done:slow", "dispatch:queued", "done:queued", "teardown"]


def test_ws_failed_reply_from_dispatcher_ends_the_connection(monkeypatch):
    """A response the dispatcher cannot send ends the connection even while the reader waits on the socket."""
    ws, inbound, _sent, log, _release = _slow_dispatch_harness(monkeypatch)
    real_send = type(ws).send_text

    async def send_text(self, line):
        if json.loads(line).get("id") == "r1":
            raise RuntimeError("peer gone")
        await real_send(self, line)

    monkeypatch.setattr(type(ws), "send_text", send_text)

    async def scenario():
        inbound.put_nowait(_rpc("r1", "fast"))
        await asyncio.wait_for(ws_mod.handle_ws(ws), 5)

    asyncio.run(scenario())
    assert log == ["dispatch:fast", "done:fast", "teardown"]


def test_ws_transport_serializes_concurrent_sends():
    active_sends = 0
    max_active_sends = 0
    sent = []

    class FakeWS:
        async def send_text(self, line):
            nonlocal active_sends, max_active_sends
            active_sends += 1
            max_active_sends = max(max_active_sends, active_sends)
            try:
                await asyncio.sleep(0.05)
                sent.append(line)
            finally:
                active_sends -= 1

    loop = asyncio.new_event_loop()
    thread = threading.Thread(target=loop.run_forever, daemon=True)
    thread.start()
    try:
        transport = ws_mod.WSTransport(FakeWS(), loop, peer="serialize-test")
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
            futures = [
                pool.submit(transport.write, {"idx": 1}),
                pool.submit(transport.write, {"idx": 2}),
            ]
            assert [f.result(timeout=2) for f in futures] == [True, True]

        assert len(sent) == 2
        assert max_active_sends == 1
        assert transport._closed is False
    finally:
        loop.call_soon_threadsafe(loop.stop)
        thread.join(timeout=2)
        loop.close()


def test_ws_transport_replies_with_error_for_unserializable_response(caplog):
    """#92506: an unserializable payload (datetime from profile.yaml ui_meta) must surface as a
    JSON-RPC error frame with the original id plus a log line — on both the worker-thread
    ``write`` and the loop-side ``write_async`` twin — and leave the transport open, instead of
    killing the pool worker silently so the client waits forever."""
    sent = []

    class FakeWS:
        async def send_text(self, line):
            sent.append(json.loads(line))

    bad = {"jsonrpc": "2.0", "id": "profiles", "result": {"created": datetime.datetime(2026, 8, 22)}}

    loop = asyncio.new_event_loop()
    thread = threading.Thread(target=loop.run_forever, daemon=True)
    thread.start()
    try:
        transport = ws_mod.WSTransport(FakeWS(), loop, peer="serialization-test")
        assert transport.write(bad) is True
        assert transport.write({"jsonrpc": "2.0", "id": "next", "result": {}}) is True
        assert asyncio.run_coroutine_threadsafe(transport.write_async(bad), loop).result(5) is True
        assert transport.closed is False
    finally:
        loop.call_soon_threadsafe(loop.stop)
        thread.join(timeout=2)
        loop.close()

    assert [m.get("id") for m in sent] == ["profiles", "next", "profiles"]
    for frame in (sent[0], sent[2]):
        assert frame["error"]["code"] == -32603
        assert frame["error"]["message"].startswith("response serialization error")
        assert "datetime" in frame["error"]["message"]
    assert sent[1] == {"jsonrpc": "2.0", "id": "next", "result": {}}
    assert caplog.text.count("frame serialization failed") == 2


def test_ws_transport_preserves_cross_batch_order():
    async def scenario():
        entered = []
        first_entered = asyncio.Event()
        release_first = asyncio.Event()
        second_started = asyncio.Event()

        class FakeWS:
            async def send_text(self, line):
                entered.append(line)
                if line == "A1":
                    first_entered.set()
                    await release_first.wait()

        transport = ws_mod.WSTransport(
            FakeWS(), asyncio.get_running_loop(), peer="batch-order-test"
        )
        first = asyncio.create_task(transport._safe_send_many(["A1", "A2"]))
        await first_entered.wait()

        async def send_second():
            second_started.set()
            await transport._safe_send_many(["B1", "B2"])

        second = asyncio.create_task(send_second())
        await second_started.wait()

        # The second task has reached the transport. Without whole-batch
        # serialization it runs B1/B2 before this task can resume.
        assert entered == ["A1"]

        release_first.set()
        await asyncio.gather(first, second)
        assert entered == ["A1", "A2", "B1", "B2"]

    asyncio.run(scenario())


_INTERACTIVE = frozenset({"session:create", "session:read", "session:submit", "session:control",
                          "session:approve", "session:respond"})


def _drive_authority_fallback(monkeypatch, *, capabilities, profile_id, method):
    """One RPC through handle_ws on an authority connection whose dispatch answers -32601, so the
    request takes the legacy-fallback branch (R2-M3). Returns the reply frame for id 1."""
    from gateway import session_controls

    class FakeConnection:
        def __init__(self, authority, transport, identity, operator=False):
            self.actor = type("Actor", (), {"capabilities": frozenset(capabilities), "profile_id": profile_id})()

        async def dispatch(self, req):
            return {"jsonrpc": "2.0", "id": req["id"], "error": {"code": -32601, "message": "unknown method"}}

        async def close(self):
            pass

    monkeypatch.setattr(session_controls, "AuthorityConnection", FakeConnection)
    monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0)
    sent, inbound = [], [json.dumps({"jsonrpc": "2.0", "id": 1, "method": method, "params": {}})]

    class FakeWS:
        scope = {"hermes.session_authority": object()}

        async def accept(self, **_kw):
            pass

        async def send_text(self, line):
            sent.extend(json.loads(part) for part in line.splitlines() if part.strip())

        async def receive_text(self):
            if inbound:
                return inbound.pop()
            raise ws_mod._WebSocketDisconnect()

        async def close(self, **_kw):
            pass

    asyncio.run(ws_mod.handle_ws(FakeWS(), auth_identity={"user_id": "u"}))
    return next(frame for frame in sent if frame.get("id") == 1)


def test_ready_frame_advertises_a_session_authority_only_when_one_is_bound(monkeypatch):
    """An explicit TUI attach (HERMES_TUI_GATEWAY_URL) negotiates the canonical wire from this flag."""
    from gateway import session_controls

    class FakeConnection:
        def __init__(self, authority, transport, identity, operator=False):
            self.actor = type("Actor", (), {"capabilities": frozenset(), "profile_id": "p"})()

        async def close(self):
            pass

    monkeypatch.setattr(session_controls, "AuthorityConnection", FakeConnection)
    monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0)

    def ready(scope):
        sent = []

        class FakeWS:
            async def accept(self, **_kw):
                pass

            async def send_text(self, line):
                sent.extend(json.loads(part) for part in line.splitlines() if part.strip())

            async def receive_text(self):
                raise ws_mod._WebSocketDisconnect()

            async def close(self, **_kw):
                pass

        FakeWS.scope = scope
        asyncio.run(ws_mod.handle_ws(FakeWS(), auth_identity={"user_id": "u"}))
        return sent[0]["params"]["payload"]

    assert ready({"hermes.session_authority": object()})["session_authority"] is True
    assert "session_authority" not in ready({})


def test_legacy_fallback_requires_the_interactive_grant(monkeypatch):
    """A worker-adoption ticket (or any connection without the authority's interactive grant) must
    not reach general legacy dispatch after an authority -32601; an interactive one still does."""
    ran = []
    monkeypatch.setitem(server._methods, "probe.legacy",
                        lambda rid, params: ran.append(rid) or {"jsonrpc": "2.0", "id": rid, "result": {}})
    launch = str(server._launch_home())
    for restricted in ({"worker:adopt"}, set()):
        reply = _drive_authority_fallback(monkeypatch, capabilities=restricted, profile_id=launch, method="probe.legacy")
        assert reply["error"]["code"] == -32601 and ran == []
    reply = _drive_authority_fallback(monkeypatch, capabilities=_INTERACTIVE, profile_id=launch, method="probe.legacy")
    assert reply["result"] == {} and ran == [1]


def test_legacy_fallback_keeps_the_ticket_profile_for_sessionless_scoped_handlers(monkeypatch, tmp_path):
    """A secondary-profile ticket's sessionless ``@_profile_scoped`` legacy call runs in THAT profile's
    home, not the launch profile's."""
    from hermes_constants import get_hermes_home

    secondary = tmp_path / "profiles" / "l106742sec"
    secondary.mkdir(parents=True)
    monkeypatch.setitem(server._methods, "probe.scoped", server._profile_scoped(
        lambda rid, params: {"jsonrpc": "2.0", "id": rid, "result": {"home": str(get_hermes_home())}}))
    reply = _drive_authority_fallback(monkeypatch, capabilities=_INTERACTIVE, profile_id=str(secondary),
                                      method="probe.scoped")
    assert reply["result"]["home"] == str(secondary)


@pytest.mark.parametrize("method,params", [
    ("session.delete", {"session_id": "owned"}),
    ("session.archive", {"session_id": "owned", "archived": True}),
])
def test_legacy_fallback_never_writes_an_authority_owned_session(monkeypatch, tmp_path, method, params):
    """A session write the authority answers -32601 for must not reach the legacy writer: it would
    change the authority's state.db behind its receipts and revision fence (and a delete would leave
    the authority holding a live session whose row is gone)."""
    from types import SimpleNamespace
    from gateway.session_authority import LiveSession, SessionAuthority
    from hermes_state import SessionDB
    from hermes_state_runtime import begin_runtime_epoch

    monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0)
    with SessionDB(db_path=tmp_path / "state.db") as db:
        db.create_session("owned", source="tui")
        revision = db.get_session("owned")["runtime_revision"]
        authority = SessionAuthority(SimpleNamespace(_draining=False), profile_id=str(tmp_path),
                                     instance_id="current", db=db, epoch=begin_runtime_epoch(db, instance_id="current"))
        authority.sessions["owned"] = LiveSession(None, "route")
        monkeypatch.setattr(server, "_profile_db", lambda *_a, **_k: contextlib.nullcontext(db))
        sent, inbound = [], [json.dumps({"jsonrpc": "2.0", "id": 1, "method": method, "params": params})]

        class FakeWS:
            scope = {"hermes.session_authority": authority}

            async def accept(self, **_kw):
                pass

            async def send_text(self, line):
                sent.extend(json.loads(part) for part in line.splitlines() if part.strip())

            async def receive_text(self):
                if inbound:
                    return inbound.pop()
                raise ws_mod._WebSocketDisconnect()

            async def close(self, **_kw):
                pass

        asyncio.run(ws_mod.handle_ws(FakeWS(), auth_identity={"user_id": "owner", "instance_id": "current"}))
        reply = next(frame for frame in sent if frame.get("id") == 1)
        assert reply["error"]["code"] == -32601, reply
        row = db.get_session("owned")
        assert row is not None and not row.get("archived") and row["runtime_revision"] == revision
        assert "owned" in authority.sessions