import unittest
from typing import Any, Dict, List, Optional

from abstra_internals.cloud.browser_service.transport import (
    BrowserServer,
    SessionOpen,
    handle_command,
)


class FakePool:
    def __init__(self, result: str = "ok", raises: Optional[Exception] = None) -> None:
        self.result = result
        self.raises = raises
        self.handled: List[tuple] = []
        self.closed: List[str] = []

    def handle(
        self, session_id, project_id, editor_token, tool_name, args, epoch
    ) -> str:
        self.handled.append(
            (session_id, project_id, editor_token, tool_name, args, epoch)
        )
        if self.raises is not None:
            raise self.raises
        return self.result

    def close(self, session_id: str, project_id: str, epoch: int) -> None:
        self.closed.append(session_id)


SESSION = SessionOpen(session_id="s1", project_id="p1")


class HandleCommandTest(unittest.TestCase):
    def test_success_returns_result_and_routes_session_context(self) -> None:
        pool = FakePool(result="the-serialized-string")
        reply = handle_command(
            pool,
            SESSION,
            {"corr_id": "c1", "tool_name": "browser_navigate", "args": {"url": "u"}},
        )
        self.assertEqual(
            reply, {"corr_id": "c1", "ok": True, "result": "the-serialized-string"}
        )
        # editor_token is not carried over the wire (dropped from the protocol);
        # the pool receives None until Phase 4 adds it back. The session's epoch
        # (default 1) is threaded through so close is epoch-guarded.
        self.assertEqual(
            pool.handled[-1], ("s1", "p1", None, "browser_navigate", {"url": "u"}, 1)
        )

    def test_tool_error_is_returned_as_tool_output_text(self) -> None:
        pool = FakePool(raises=ValueError("Page 't' does not exist."))
        reply = handle_command(
            pool,
            SESSION,
            {"corr_id": "c2", "tool_name": "browser_get_text", "args": {}},
        )
        self.assertFalse(reply["ok"])
        self.assertEqual(
            reply["result"], "Error executing tool: Page 't' does not exist."
        )

    def test_close_command_short_circuits(self) -> None:
        pool = FakePool()
        reply = handle_command(
            pool, SESSION, {"corr_id": "c3", "tool_name": "__close__"}
        )
        self.assertTrue(reply["ok"])
        self.assertTrue(reply["closed"])
        self.assertEqual(pool.handled, [])  # not dispatched


class FakeConn:
    """Stands in for a RabbitMQConnection session channel."""

    def __init__(self, incoming: List[Any]) -> None:
        self._incoming = list(incoming)
        self.sent: List[Dict[str, Any]] = []
        self.closed = False

    def recv(self) -> Any:
        return self._incoming.pop(0) if self._incoming else None

    def send(self, obj: Dict[str, Any]) -> None:
        self.sent.append(obj)

    def close(self) -> None:
        self.closed = True


class ServeSessionTest(unittest.TestCase):
    def test_serves_commands_then_closes_on_close_command(self) -> None:
        pool = FakePool(result="R")
        server = BrowserServer("amqp://unused", pool)
        conn = FakeConn(
            [
                {
                    "corr_id": "c1",
                    "tool_name": "browser_navigate",
                    "args": {"url": "u"},
                },
                {"corr_id": "c2", "tool_name": "__close__"},
                {
                    "corr_id": "c3",
                    "tool_name": "browser_navigate",
                },  # after close: ignored
            ]
        )
        server._serve_session(conn, SESSION)

        # Two replies (command + close), loop stopped at close.
        self.assertEqual([r["corr_id"] for r in conn.sent], ["c1", "c2"])
        self.assertEqual(conn.sent[0]["result"], "R")
        self.assertTrue(conn.sent[1]["closed"])
        # Session cleaned up + channel closed.
        self.assertEqual(pool.closed, ["s1"])
        self.assertTrue(conn.closed)

    def test_none_ends_the_loop_and_cleans_up(self) -> None:
        pool = FakePool()
        server = BrowserServer("amqp://unused", pool)
        conn = FakeConn([])  # recv() -> None immediately
        server._serve_session(conn, SESSION)
        self.assertEqual(conn.sent, [])
        self.assertEqual(pool.closed, ["s1"])
        self.assertTrue(conn.closed)

    def test_json_string_commands_are_parsed(self) -> None:
        pool = FakePool(result="R")
        server = BrowserServer("amqp://unused", pool)
        conn = FakeConn(
            ['{"corr_id": "c1", "tool_name": "browser_list_tabs", "args": {}}', None]
        )
        server._serve_session(conn, SESSION)
        self.assertEqual(conn.sent[0], {"corr_id": "c1", "ok": True, "result": "R"})


class EpochFencingTest(unittest.TestCase):
    def test_subject_names_are_scoped_by_epoch(self) -> None:
        from abstra_internals.cloud.browser_service.transport import (
            cmd_subject,
            reply_subject,
        )

        self.assertEqual(cmd_subject("s1", 1), "browser-service.cmd.s1.1")
        self.assertEqual(reply_subject("s1", 1), "browser-service.reply.s1.1")
        # A re-open (epoch+1) yields DISTINCT subjects, so no two holders ever
        # share one session's cmd subject — the fencing guarantee.
        self.assertNotEqual(cmd_subject("s1", 1), cmd_subject("s1", 2))
        self.assertNotEqual(reply_subject("s1", 1), reply_subject("s1", 2))

    def test_subject_token_sanitizes_non_token_chars(self) -> None:
        from abstra_internals.cloud.browser_service.transport import cmd_subject

        # Dots/wildcards in a session id must not break subject token structure.
        self.assertEqual(cmd_subject("a.b*c>d", 1), "browser-service.cmd.a_b_c_d.1")

    def test_session_open_epoch_defaults_to_1_and_parses(self) -> None:
        base = {"session_id": "s", "project_id": "p"}
        self.assertEqual(SessionOpen.from_dict(base).epoch, 1)  # back-compat default
        self.assertEqual(SessionOpen.from_dict({**base, "epoch": 5}).epoch, 5)


class BlockingConn(FakeConn):
    """A session channel whose recv() blocks until release() — keeps the serving
    thread (and its capacity slot) alive for the duration of a test."""

    def __init__(self) -> None:
        import queue

        super().__init__([])
        self._q: "queue.Queue" = queue.Queue()

    def recv(self) -> Any:
        return self._q.get()

    def release(self) -> None:
        self._q.put(None)  # recv -> None ends the serve loop


def _wait_until(predicate, timeout: float = 2.0) -> bool:
    import time

    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(0.01)
    return False


class AdmissionTest(unittest.TestCase):
    def test_full_replica_refuses_the_open_and_frees_on_session_end(self) -> None:
        pool = FakePool()
        server = BrowserServer("amqp://unused", pool, max_sessions=1)
        conn = BlockingConn()
        server.open_session_channel = lambda session: conn  # type: ignore[method-assign]

        self.assertTrue(server.try_take_session(SessionOpen("s1", "p1")))
        self.assertEqual(server.active_sessions, 1)
        # Handshake: the holder announces readiness BEFORE any command can come.
        self.assertEqual(
            conn.sent[0], {"corr_id": "__ready__", "ok": True, "result": ""}
        )
        # At capacity: the second open is refused (caller requeues it).
        self.assertFalse(server.try_take_session(SessionOpen("s2", "p1")))

        conn.release()  # session ends -> slot frees
        self.assertTrue(_wait_until(lambda: server.active_sessions == 0))
        self.assertTrue(server.try_take_session(SessionOpen("s2", "p1")))

    def test_channel_open_failure_refuses_without_leaking_the_slot(self) -> None:
        server = BrowserServer("amqp://unused", FakePool(), max_sessions=1)

        def boom(session: SessionOpen) -> Any:
            raise RuntimeError("broker hiccup")

        server.open_session_channel = boom  # type: ignore[method-assign]
        self.assertFalse(server.try_take_session(SessionOpen("s1", "p1")))
        self.assertEqual(server.active_sessions, 0)  # slot not leaked

    def test_uncapped_server_always_admits(self) -> None:
        server = BrowserServer("amqp://unused", FakePool(), max_sessions=None)
        conn = BlockingConn()
        server.open_session_channel = lambda session: conn  # type: ignore[method-assign]
        for i in range(5):
            self.assertTrue(server.try_take_session(SessionOpen(f"s{i}", "p1")))
        self.assertEqual(server.active_sessions, 5)


class NatsSessionChannelTest(unittest.TestCase):
    """recv() maps idle-timeout/closed to None — the serve loop's end signal
    (what queue expiry used to provide on the RabbitMQ session channel)."""

    def _channel(self, inner: Any):
        from abstra_internals.cloud.browser_service.transport import (
            _NatsSessionChannel,
        )

        return _NatsSessionChannel(inner, idle_seconds=0.01)

    def test_recv_passes_through_messages_with_the_idle_timeout(self) -> None:
        class Inner:
            def recv(self, timeout=None):
                assert timeout == 0.01
                return {"corr_id": "c1"}

        self.assertEqual(self._channel(Inner()).recv(), {"corr_id": "c1"})

    def test_recv_returns_none_on_idle_timeout_and_on_closed(self) -> None:
        class Idle:
            def recv(self, timeout=None):
                raise TimeoutError()

        class Closed:
            def recv(self, timeout=None):
                raise EOFError()

        self.assertIsNone(self._channel(Idle()).recv())
        self.assertIsNone(self._channel(Closed()).recv())


if __name__ == "__main__":
    unittest.main()
