"""Tests for ``relay_bus.RelayBus``.

Strategy: spin up a local ``websockets.serve`` mock relay on 127.0.0.1, run
the bus against it. The mock implements just enough of NIP-01 to drive the
bus through its happy and edge paths (REQ → EVENT → EOSE → CLOSE; EVENT →
OK; CLOSED; abrupt disconnect for reconnect testing).

Each test spins up its own mock relay on a free port so they can run in
parallel without colliding.
"""

from __future__ import annotations

import asyncio
import base64
import hashlib
import json
import random
import socket
import threading
import time
from typing import Any

import pytest

import relay_bus  # noqa: E402  (sys.path patched in conftest)
from mock_relay import MockRelay


# ---- fixtures --------------------------------------------------------------


@pytest.fixture
def bus():
    b = relay_bus.RelayBus()
    b.start()
    yield b
    b.stop(timeout=3.0)


# ---- tests -----------------------------------------------------------------


def test_start_stop_clean(bus):
    # smoke: the bus came up in the fixture without raising.
    assert bus._loop is not None


def test_req_happy_path(bus):
    """Mock relay returns 3 EVENTs and an EOSE for any REQ. The bus should
    collect all three events and send CLOSE on the way out."""
    events_to_send = [
        {"id": "e1", "kind": 1, "content": "first", "created_at": 1, "pubkey": "p", "sig": "s", "tags": []},
        {"id": "e2", "kind": 1, "content": "second", "created_at": 2, "pubkey": "p", "sig": "s", "tags": []},
        {"id": "e3", "kind": 1, "content": "third", "created_at": 3, "pubkey": "p", "sig": "s", "tags": []},
    ]

    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            sub_id = msg[1]
            for e in events_to_send:
                await ws.send(json.dumps(["EVENT", sub_id, e]))
            await ws.send(json.dumps(["EOSE", sub_id]))

    relay = MockRelay(handler)
    relay.start()
    try:
        got = bus.req_sync(relay.url, {"kinds": [1], "limit": 10}, timeout=5.0)
    finally:
        relay.stop()

    assert [e["id"] for e in got] == ["e1", "e2", "e3"]
    # Last frame the bus sent must be a CLOSE referencing the same sub_id.
    sub_ids = [f[1] for f in relay.received if f[0] == "REQ"]
    closes = [f[1] for f in relay.received if f[0] == "CLOSE"]
    assert sub_ids and closes == sub_ids


def test_req_timeout_returns_partial(bus):
    """Relay sends one event and then hangs (no EOSE). The req should time
    out, send CLOSE, and return the single event it managed to collect."""
    sent_event = {"id": "only", "kind": 1, "content": "x", "created_at": 1,
                  "pubkey": "p", "sig": "s", "tags": []}

    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EVENT", msg[1], sent_event]))
            # deliberately no EOSE

    relay = MockRelay(handler)
    relay.start()
    try:
        t0 = time.monotonic()
        got = bus.req_sync(relay.url, {"kinds": [1]}, timeout=0.5)
        elapsed = time.monotonic() - t0
    finally:
        relay.stop()

    assert [e["id"] for e in got] == ["only"]
    # We don't want to wait dramatically longer than the timeout itself.
    assert elapsed < 2.0
    assert any(f[0] == "CLOSE" for f in relay.received)


def test_req_multiplex_on_one_ws(bus):
    """Two sequential req_sync calls to the same URL should reuse one WS and
    use distinct sub_ids."""
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EOSE", msg[1]]))

    relay = MockRelay(handler)
    relay.start()
    try:
        bus.req_sync(relay.url, {"a": 1}, timeout=2.0)
        bus.req_sync(relay.url, {"a": 2}, timeout=2.0)
    finally:
        relay.stop()

    reqs = [f for f in relay.received if f[0] == "REQ"]
    assert len(reqs) == 2
    assert reqs[0][1] != reqs[1][1]  # different sub_ids
    # Only one WS upgrade — `_connections` is the list of every server-side
    # connection ever accepted.
    assert len(relay._connections) == 1


def test_publish_ok(bus):
    async def handler(relay, ws, msg):
        if msg[0] == "EVENT":
            eid = msg[1].get("id")
            await ws.send(json.dumps(["OK", eid, True, ""]))

    relay = MockRelay(handler)
    relay.start()
    try:
        bus.publish_sync(
            relay.url,
            {"id": "abc", "kind": 1, "content": "hi", "created_at": 1,
             "pubkey": "p", "sig": "s", "tags": []},
            timeout=2.0,
        )
    finally:
        relay.stop()


def test_publish_rejected_raises(bus):
    async def handler(relay, ws, msg):
        if msg[0] == "EVENT":
            await ws.send(json.dumps(["OK", msg[1]["id"], False, "rate-limited"]))

    relay = MockRelay(handler)
    relay.start()
    try:
        with pytest.raises(relay_bus.RelayPublishRejected, match="rate-limited"):
            bus.publish_sync(
                relay.url,
                {"id": "abc", "kind": 1, "content": "no", "created_at": 1,
                 "pubkey": "p", "sig": "s", "tags": []},
                timeout=2.0,
            )
    finally:
        relay.stop()


def test_publish_timeout_raises(bus):
    async def handler(relay, ws, msg):
        if msg[0] == "EVENT":
            # ignore — no OK reply
            return

    relay = MockRelay(handler)
    relay.start()
    try:
        with pytest.raises(relay_bus.RelayTimeout):
            bus.publish_sync(
                relay.url,
                {"id": "abc", "kind": 1, "content": "x", "created_at": 1,
                 "pubkey": "p", "sig": "s", "tags": []},
                timeout=0.5,
            )
    finally:
        relay.stop()


def test_subscribe_yields_events_and_can_cancel(bus):
    """Live subscription receives a couple of pushed events, then the consumer
    cancels and the bus sends CLOSE."""
    delivered: list[dict] = []
    delivered_evt = threading.Event()

    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            sub_id = msg[1]
            await ws.send(json.dumps(["EVENT", sub_id, {
                "id": "live1", "kind": 1, "content": "a", "created_at": 1,
                "pubkey": "p", "sig": "s", "tags": [],
            }]))
            await ws.send(json.dumps(["EOSE", sub_id]))
            await asyncio.sleep(0.1)
            await ws.send(json.dumps(["EVENT", sub_id, {
                "id": "live2", "kind": 1, "content": "b", "created_at": 2,
                "pubkey": "p", "sig": "s", "tags": [],
            }]))

    def on_event(evt: dict) -> None:
        delivered.append(evt)
        if len(delivered) >= 2:
            delivered_evt.set()

    relay = MockRelay(handler)
    relay.start()
    try:
        handle = bus.subscribe_sync(relay.url, {"kinds": [1]}, on_event)
        assert delivered_evt.wait(timeout=3.0), "did not receive two live events"
        handle.cancel(timeout=2.0)
    finally:
        relay.stop()

    ids = [e["id"] for e in delivered]
    assert ids == ["live1", "live2"]
    assert any(f[0] == "CLOSE" for f in relay.received)


def test_concurrent_first_open_is_serialized(bus):
    """Two ``req_sync`` calls firing in parallel on the SAME URL must NOT
    each start their own recv task. Regression test for the ConcurrencyError
    'cannot call recv while another coroutine is already running recv' that
    surfaced under load when ``ensure_open`` had no internal lock.
    """
    handler_called = threading.Event()

    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            handler_called.set()
            await ws.send(json.dumps(["EOSE", msg[1]]))

    relay = MockRelay(handler)
    relay.start()
    try:
        errors: list[Exception] = []

        def _go():
            try:
                bus.req_sync(relay.url, {"kinds": [1]}, timeout=3.0)
            except Exception as e:
                errors.append(e)

        threads = [threading.Thread(target=_go) for _ in range(5)]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=5.0)
        assert handler_called.is_set()
        assert errors == [], f"unexpected errors: {errors}"
        # Exactly one WS handshake despite five concurrent callers — the
        # open_lock in ``_RelayConnection`` collapses them onto one connection.
        assert len(relay._connections) == 1
    finally:
        relay.stop()


def test_subscribe_reconnects_and_reissues_req(bus):
    """Live sub: after the relay drops the WS, the bus reconnects and
    re-issues the REQ with an updated since cursor."""
    delivered_evt = threading.Event()
    delivered: list[dict] = []
    drop_done = threading.Event()

    async def handler(relay, ws, msg):
        if msg[0] != "REQ":
            return
        sub_id = msg[1]
        flt = msg[2] if len(msg) >= 3 else {}
        # First REQ (no since): deliver a single event then drop the WS.
        # Second REQ (after reconnect): must carry since > 0; deliver another.
        if "since" not in flt:
            await ws.send(json.dumps(["EVENT", sub_id, {
                "id": "before", "kind": 1, "content": "x", "created_at": 100,
                "pubkey": "p", "sig": "s", "tags": [],
            }]))
            await ws.send(json.dumps(["EOSE", sub_id]))
            await asyncio.sleep(0.1)
            drop_done.set()
            await ws.close(code=1011, reason="test-drop")
        else:
            assert flt["since"] >= 101, f"expected since >= 101, got {flt['since']!r}"
            await ws.send(json.dumps(["EVENT", sub_id, {
                "id": "after", "kind": 1, "content": "y", "created_at": 200,
                "pubkey": "p", "sig": "s", "tags": [],
            }]))
            await ws.send(json.dumps(["EOSE", sub_id]))

    def on_event(evt: dict) -> None:
        delivered.append(evt)
        if any(e["id"] == "after" for e in delivered):
            delivered_evt.set()

    relay = MockRelay(handler)
    relay.start()
    try:
        handle = bus.subscribe_sync(relay.url, {"kinds": [1]}, on_event)
        assert delivered_evt.wait(timeout=8.0), (
            f"reconnect+resub never delivered post-drop event; got: "
            f"{[e['id'] for e in delivered]}"
        )
        handle.cancel(timeout=2.0)
    finally:
        relay.stop()

    ids = [e["id"] for e in delivered]
    assert "before" in ids and "after" in ids


# ---- leak / hardening fixtures + helpers -----------------------------------


@pytest.fixture
def fast_grace(monkeypatch):
    """Shorten the idle-close grace so leak assertions don't have to wait
    the full 5 s every time. Also bound the WS close-handshake timeout so
    tests that exercise the transport-abort fallback run in <1 s."""
    monkeypatch.setattr(relay_bus, "IDLE_CLOSE_GRACE_S", 0.15)
    monkeypatch.setattr(relay_bus, "WS_CLOSE_TIMEOUT_S", 0.3)
    yield


def _wait_for(predicate: Callable[[], bool], *, timeout: float = 3.0,
              interval: float = 0.02) -> bool:
    """Poll ``predicate`` until it's truthy or ``timeout`` elapses."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(interval)
    return False


def _registry_size(bus: relay_bus.RelayBus) -> int:
    return len(bus._connections)


# ---- hanging-close server (raw TCP) ----------------------------------------


class HangingRelay:
    """Raw TCP server that fakes a WebSocket relay in pathological ways.

    Two modes (set via ``send_handshake``):
      * ``True`` (default): accept the upgrade, then silently swallow every
        frame including the client's CLOSE. Exercises the
        ``_close_ws_safely`` / ``transport.abort`` path.
      * ``False``: accept the TCP connect, read the HTTP upgrade request,
        but never send the 101 response. Exercises the connect-side timeout
        (``CONNECT_TIMEOUT_S``).
    """

    _WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

    def __init__(self, send_handshake: bool = True) -> None:
        self.send_handshake = send_handshake
        self.url: str | None = None
        self._loop: asyncio.AbstractEventLoop | None = None
        self._thread: threading.Thread | None = None
        self._ready = threading.Event()
        self._stopped = threading.Event()
        self._server: asyncio.base_events.Server | None = None

    @staticmethod
    def _free_port() -> int:
        with socket.socket() as s:
            s.bind(("127.0.0.1", 0))
            return s.getsockname()[1]

    def start(self) -> None:
        port = self._free_port()
        self.url = f"ws://127.0.0.1:{port}"

        def _run():
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            self._loop = loop

            async def _main():
                self._server = await asyncio.start_server(
                    self._on_client, "127.0.0.1", port,
                )
                self._ready.set()
                while not self._stopped.is_set():
                    await asyncio.sleep(0.05)
                self._server.close()
                await self._server.wait_closed()

            try:
                loop.run_until_complete(_main())
            finally:
                loop.close()

        self._thread = threading.Thread(
            target=_run, name="hanging-relay", daemon=True,
        )
        self._thread.start()
        if not self._ready.wait(timeout=3.0):
            raise RuntimeError("hanging relay didn't start")

    def stop(self) -> None:
        self._stopped.set()
        if self._thread is not None:
            self._thread.join(timeout=3.0)

    @classmethod
    def _accept_value(cls, key: str) -> str:
        return base64.b64encode(
            hashlib.sha1((key + cls._WS_GUID).encode()).digest(),
        ).decode()

    async def _on_client(self, reader: asyncio.StreamReader,
                         writer: asyncio.StreamWriter) -> None:
        try:
            key = None
            while True:
                line = await reader.readline()
                if not line or line == b"\r\n":
                    break
                if line.lower().startswith(b"sec-websocket-key:"):
                    key = line.split(b":", 1)[1].strip().decode()
            if key is None:
                writer.close()
                return
            if not self.send_handshake:
                # Hold the socket open without replying. The bus's connect
                # path should hit `CONNECT_TIMEOUT_S` and bail.
                while not self._stopped.is_set():
                    await asyncio.sleep(0.05)
                return
            accept = self._accept_value(key)
            writer.write(
                b"HTTP/1.1 101 Switching Protocols\r\n"
                b"Upgrade: websocket\r\n"
                b"Connection: Upgrade\r\n"
                b"Sec-WebSocket-Accept: " + accept.encode() + b"\r\n\r\n"
            )
            await writer.drain()
            # Read everything and drop it. Critically: do NOT send a Close
            # frame back, even when the client sends one. The client's
            # `ws.close()` should hang until our hard timeout fires.
            while not self._stopped.is_set():
                try:
                    data = await asyncio.wait_for(reader.read(4096), timeout=0.5)
                except asyncio.TimeoutError:
                    continue
                if not data:
                    break
        except Exception:
            pass
        finally:
            with contextlib_suppress():
                writer.close()


class contextlib_suppress:
    """Tiny no-deps suppressor — keeps the test file from importing
    contextlib just for one finally clause."""
    def __enter__(self): return self
    def __exit__(self, *_): return True


# ---- leak / hardening tests ------------------------------------------------


def test_req_idle_close_drops_from_registry(bus, fast_grace):
    """Baseline: a one-shot req leaves nothing in `bus._connections` once
    the idle-close grace has elapsed. Anchor for the rest of the leak
    suite — if this fails, the cleanup contract is broken at the surface."""
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EOSE", msg[1]]))

    relay = MockRelay(handler)
    relay.start()
    try:
        bus.req_sync(relay.url, {"kinds": [1]}, timeout=2.0)
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"req conn lingered: {bus.stats()}"
        )
    finally:
        relay.stop()


def test_subscribe_cancel_drops_from_registry(bus, fast_grace):
    """Regression for the orphaned-async-generator leak: cancelling a
    `subscribe_sync` handle MUST end with the conn evicted from the bus
    registry. Pre-fix `_drive` exited without `aclose()`-ing the inner
    `subscribe()` generator, so its finally (close_sub +
    schedule_idle_close_if_empty) never ran and the conn lingered forever.
    """
    delivered: list[dict] = []
    delivered_first = threading.Event()

    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            sub_id = msg[1]
            await ws.send(json.dumps(["EVENT", sub_id, {
                "id": "e1", "kind": 1, "content": "x", "created_at": 1,
                "pubkey": "p", "sig": "s", "tags": [],
            }]))
            await ws.send(json.dumps(["EOSE", sub_id]))
            # Hold the connection open so the only way out is our cancel.
            await asyncio.sleep(30)

    def on_event(evt: dict) -> None:
        delivered.append(evt)
        delivered_first.set()

    relay = MockRelay(handler)
    relay.start()
    try:
        handle = bus.subscribe_sync(relay.url, {"kinds": [1]}, on_event)
        assert delivered_first.wait(timeout=3.0), "live sub never produced"
        handle.cancel(timeout=2.0)
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"subscribe conn lingered: {bus.stats()}"
        )
        # The bus-loop task histogram must also drain (drive + recv + idle).
        s = bus.stats()
        assert s["subs_by_mode"] == {"req": 0, "subscribe": 0}, s
        assert s["tasks_by_name"].get("drive", 0) == 0, s
    finally:
        relay.stop()


def test_server_drops_during_req_drops_from_registry(bus, fast_grace):
    """Server force-closes the WS mid-req. The recv loop exits via
    `if not self._subs: return` after the req-mode sub is filtered out.
    Pre-fix this self-termination path left the entry in the registry;
    post-fix it calls `on_done` so the registry stays clean.
    """
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EVENT", msg[1], {
                "id": "e", "kind": 1, "content": "x", "created_at": 1,
                "pubkey": "p", "sig": "s", "tags": [],
            }]))
            # No EOSE; drop instead.
            await asyncio.sleep(0.05)
            await ws.close(code=1011, reason="test-drop")

    relay = MockRelay(handler)
    relay.start()
    try:
        # The req may finish (with the one event) or raise; either is fine,
        # what matters is the post-state of the registry.
        try:
            bus.req_sync(relay.url, {"kinds": [1]}, timeout=1.0)
        except Exception:
            pass
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=3.0), (
            f"stale conn after server drop: {bus.stats()}"
        )
    finally:
        relay.stop()


def test_concurrent_subscribe_cancel_no_leak(bus, fast_grace):
    """N parallel subscribers on the same relay, all cancelled within a
    tight window. Stresses the cancel race that pre-fix accumulated zombie
    `drive:` tasks and stale conn-registry entries (the `Task` + socket
    histogram blow-up we observed at runtime)."""
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EOSE", msg[1]]))
            await asyncio.sleep(30)

    relay = MockRelay(handler)
    relay.start()
    try:
        handles = [
            bus.subscribe_sync(relay.url, {"kinds": [i]}, lambda e: None)
            for i in range(15)
        ]
        # Let each subscription get past REQ → EOSE so the bus actually
        # has 15 live subs before we tear them all down.
        time.sleep(0.3)
        assert bus.stats()["subscriptions"] == 15

        for h in handles:
            h.cancel(timeout=2.0)

        # After cancel everything should drain.
        assert _wait_for(lambda: bus.stats()["subscriptions"] == 0, timeout=2.0), (
            f"sub registry leaked: {bus.stats()}"
        )
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"conn registry leaked: {bus.stats()}"
        )
        # `drive:` tasks must all be gone.
        assert bus.stats()["tasks_by_name"].get("drive", 0) == 0, bus.stats()
    finally:
        relay.stop()


def test_hanging_close_handshake_uses_transport_abort(bus, fast_grace):
    """Server accepts the upgrade but never replies to anything — including
    the client's close frame. `_close_ws_safely` must bound `ws.close()`
    via `WS_CLOSE_TIMEOUT_S` and force `transport.abort()` so we tear down
    in <1 s instead of waiting for the library's internal 10-s close
    timeout (the real-world signature of the keepalive-task / socket-fd
    leak we observed)."""
    server = HangingRelay()
    server.start()
    try:
        handle = bus.subscribe_sync(server.url, {"kinds": [1]}, lambda e: None)
        # Let the bus actually open the WS (the upgrade is synchronous on
        # the server side, but the bus loop's connect coroutine needs a
        # tick to run).
        assert _wait_for(
            lambda: bus.stats()["connections_by_state"]["open"] >= 1,
            timeout=2.0,
        ), f"WS never opened: {bus.stats()}"

        t0 = time.monotonic()
        handle.cancel(timeout=2.0)
        ok = _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0)
        elapsed = time.monotonic() - t0

        assert ok, f"conn lingered after hanging close: {bus.stats()}"
        # WS_CLOSE_TIMEOUT_S is 0.3 in fast_grace; full teardown should be
        # well under a second. Without transport.abort the websockets
        # library's own close-handshake wait stretches this to ~10 s.
        assert elapsed < 1.5, f"hanging close not bounded: {elapsed:.2f}s"
    finally:
        server.stop()


def test_malformed_frames_do_not_corrupt_bus(bus, fast_grace):
    """Server sprays junk (non-JSON, wrong-shape arrays, unknown verbs) at
    the bus. The req should still time out cleanly, the registry should
    drain, and no exception should escape the bus loop."""
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            sub_id = msg[1]
            await ws.send("this is not json")
            await ws.send(json.dumps({"not": "a list"}))
            await ws.send(json.dumps([]))  # empty list
            await ws.send(json.dumps(["GARBAGE", "extra", 42]))
            await ws.send(json.dumps(["EVENT", sub_id, "not-a-dict"]))
            await ws.send(json.dumps(["EVENT", "unknown-sub", {
                "id": "x", "kind": 1, "content": "x", "created_at": 1,
                "pubkey": "p", "sig": "s", "tags": [],
            }]))
            # Finally a valid EOSE.
            await ws.send(json.dumps(["EOSE", sub_id]))

    relay = MockRelay(handler)
    relay.start()
    try:
        got = bus.req_sync(relay.url, {"kinds": [1]}, timeout=2.0)
        assert got == []  # no valid events delivered for our sub_id
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"registry didn't drain after garbage: {bus.stats()}"
        )
    finally:
        relay.stop()


def test_ensure_open_raises_on_torn_down_conn(bus):
    """White-box invariant: after a conn has self-terminated, calling
    `ensure_open` on the same object must raise rather than try to revive
    it. `_get_or_open` relies on this to detect the race where a caller
    grabbed the conn from the registry just before `_drop` ran."""
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EOSE", msg[1]]))

    relay = MockRelay(handler)
    relay.start()
    try:
        bus.req_sync(relay.url, {"kinds": [1]}, timeout=2.0)
        # Grab the conn (still in registry; idle-close not fired yet here
        # because we use the default 5 s grace) and force-close it.
        conn = bus._connections.get(relay.url)
        assert conn is not None
        fut = asyncio.run_coroutine_threadsafe(conn.close(), bus._loop)
        fut.result(timeout=2.0)
        assert conn._close_requested

        async def _call():
            await conn.ensure_open()

        with pytest.raises(relay_bus.RelayConnectError, match="torn down"):
            asyncio.run_coroutine_threadsafe(_call(), bus._loop).result(timeout=2.0)
    finally:
        relay.stop()


def test_stale_conn_in_registry_is_replaced_on_next_get(bus):
    """If a conn self-terminates between recv-loop exit and the registry
    pop firing, the next caller asking the bus for the same URL must get
    a *fresh* `_RelayConnection`, not the torn-down one. Without the
    `_close_requested` check in `_get_or_open` the caller would receive
    a zombie and `ensure_open` would raise on every retry."""
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EOSE", msg[1]]))

    relay = MockRelay(handler)
    relay.start()
    try:
        bus.req_sync(relay.url, {"kinds": [1]}, timeout=2.0)
        # Manually mark the registry conn as torn down without going through
        # close() — simulates the brief window where _on_done has not yet
        # popped the entry.
        conn = bus._connections.get(relay.url)
        assert conn is not None
        conn._close_requested = True

        # Next req must succeed by replacing the stale entry under reg_lock.
        bus.req_sync(relay.url, {"kinds": [1]}, timeout=2.0)
        new_conn = bus._connections.get(relay.url)
        assert new_conn is not None
        assert new_conn is not conn, "stale conn was reused"
    finally:
        relay.stop()


def test_burst_req_then_drain(bus, fast_grace):
    """Stress: many parallel req_syncs against several relays. After the
    burst settles, the registry must drain to zero and `asyncio_tasks`
    must return to the bus-loop baseline (the bus loop itself runs a
    couple of internal tasks; anything beyond that is a leak)."""
    relays: list[MockRelay] = []

    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EOSE", msg[1]]))

    try:
        for _ in range(4):
            r = MockRelay(handler)
            r.start()
            relays.append(r)

        rng = random.Random(0xC0FFEE)
        errors: list[Exception] = []

        def _hit():
            r = rng.choice(relays)
            try:
                bus.req_sync(r.url, {"kinds": [1]}, timeout=2.0)
            except Exception as e:
                errors.append(e)

        threads = [threading.Thread(target=_hit) for _ in range(40)]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=8.0)

        assert errors == [], f"req errors: {errors}"
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"registry leaked after burst: {bus.stats()}"
        )
        s = bus.stats()
        # `recv:` and `drive:` buckets must be empty; `Task` bucket holds
        # only transient `wait_for`-wrappers, which should all be settled.
        assert s["tasks_by_name"].get("recv", 0) == 0, s
        assert s["tasks_by_name"].get("drive", 0) == 0, s
        assert s["asyncio_tasks"] < 5, s
    finally:
        for r in relays:
            r.stop()


# ---- aggressive bug-hunting tests ------------------------------------------
# Below tests are written to FAIL on plausible-but-undetected leaks. Some
# target paths in the base code that pre-date the recent fixes, so a failure
# here is informative — it means a real bug surfaced, not a regression in
# the diagnostics layer. Each test names what specifically it's checking
# and what direction a failure points to.


def _unreachable_url() -> str:
    """A `ws://127.0.0.1:<port>` URL where nothing is listening. Picking a
    port by binding+closing leaves a race window but for these tests the
    bus's connect timeout dominates anyway."""
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        port = s.getsockname()[1]
    return f"ws://127.0.0.1:{port}"


def test_unreachable_url_does_not_leak_registry_entry(bus, monkeypatch):
    """A failed connect must NOT leave a stale entry in `bus._connections`.

    Direction of failure: `_get_or_open` inserts the conn into the registry
    BEFORE calling `ensure_open`. If `ensure_open` raises, the half-built
    entry never gets cleaned up — each fresh attempt re-uses the same dead
    conn and keeps failing forever, while the registry slowly grows by one
    every time the caller chooses a new URL.
    """
    monkeypatch.setattr(relay_bus, "CONNECT_TIMEOUT_S", 0.3)

    url = _unreachable_url()
    with pytest.raises(relay_bus.RelayConnectError):
        bus.req_sync(url, {"kinds": [1]}, timeout=2.0)

    # We do NOT want this URL to live in the registry — it's not a valid
    # conn, there's no recv task, and re-using it would just retry the
    # broken connect on every call.
    assert url not in bus._connections, (
        f"failed-connect left registry entry: {list(bus._connections)}"
    )


def test_repeated_failed_connects_do_not_accumulate(bus, monkeypatch):
    """Caller hammers many different unreachable URLs. The registry must
    not grow unboundedly. A failure here means every failed connect is
    leaving a zombie; over time a real deployment with flaky relays would
    accumulate them indefinitely."""
    monkeypatch.setattr(relay_bus, "CONNECT_TIMEOUT_S", 0.2)

    urls = [_unreachable_url() for _ in range(8)]
    for url in urls:
        with pytest.raises(relay_bus.RelayConnectError):
            bus.req_sync(url, {"kinds": [1]}, timeout=1.0)

    leftover = [u for u in urls if u in bus._connections]
    assert leftover == [], f"registry accumulated dead entries: {leftover}"


def test_publish_timeout_drains_registry(bus, fast_grace):
    """Server reads EVENT but never sends OK. The publish must time out
    cleanly and the conn must drain like any other transient. Failure
    means the publish path's finally is wired wrong (no
    `schedule_idle_close_if_empty`) and we have a publish-specific leak."""
    async def handler(relay, ws, msg):
        if msg[0] == "EVENT":
            # Acknowledge nothing.
            return

    relay = MockRelay(handler)
    relay.start()
    try:
        with pytest.raises(relay_bus.RelayTimeout):
            bus.publish_sync(
                relay.url,
                {"id": "x", "kind": 1, "content": "n", "created_at": 1,
                 "pubkey": "p", "sig": "s", "tags": []},
                timeout=0.4,
            )
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"publish-timeout conn lingered: {bus.stats()}"
        )
    finally:
        relay.stop()


def test_publish_rejected_drains_registry(bus, fast_grace):
    """Server returns ``OK <id> false``. `RelayPublishRejected` is raised
    and the conn drains."""
    async def handler(relay, ws, msg):
        if msg[0] == "EVENT":
            await ws.send(json.dumps(["OK", msg[1]["id"], False, "nope"]))

    relay = MockRelay(handler)
    relay.start()
    try:
        with pytest.raises(relay_bus.RelayPublishRejected):
            bus.publish_sync(
                relay.url,
                {"id": "y", "kind": 1, "content": "n", "created_at": 1,
                 "pubkey": "p", "sig": "s", "tags": []},
                timeout=2.0,
            )
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"publish-rejected conn lingered: {bus.stats()}"
        )
    finally:
        relay.stop()


def test_cancel_during_reconnect_backoff_drops_conn(bus, fast_grace, monkeypatch):
    """Subscribe → relay gone → bus enters reconnect-backoff sleep → caller
    cancels the handle while the bus is waiting. The conn must still
    drain: the cancel + aclose path schedules idle-close, which cancels
    the backoff-sleeping recv task and pops the conn from the registry.

    A failure here means the cancel race with reconnect-backoff doesn't
    properly tear down — the most common shape of a stuck recv-task.
    """
    monkeypatch.setattr(relay_bus, "RECONNECT_BACKOFF_INITIAL_S", 0.3)

    delivered_first = threading.Event()

    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EVENT", msg[1], {
                "id": "e", "kind": 1, "content": "x", "created_at": 1,
                "pubkey": "p", "sig": "s", "tags": [],
            }]))
            await ws.send(json.dumps(["EOSE", msg[1]]))

    relay = MockRelay(handler)
    relay.start()

    def on_event(_evt):
        delivered_first.set()

    handle = bus.subscribe_sync(relay.url, {"kinds": [1]}, on_event)
    assert delivered_first.wait(timeout=2.0)

    # Kill the relay so the bus enters reconnect-backoff. Give it a brief
    # moment to actually start the backoff sleep.
    relay.stop()
    time.sleep(0.4)

    handle.cancel(timeout=2.0)
    assert _wait_for(lambda: _registry_size(bus) == 0, timeout=3.0), (
        f"cancel-during-backoff didn't drain: {bus.stats()}"
    )
    s = bus.stats()
    assert s["tasks_by_name"].get("recv", 0) == 0, s
    assert s["tasks_by_name"].get("drive", 0) == 0, s


def test_subscribe_through_many_server_drops_no_growth(bus, fast_grace):
    """Server drops the WS after every EOSE so the bus keeps reconnecting.
    Over N cycles, the per-conn invariants must hold:
      * exactly 1 entry in the bus registry (the same conn, reused),
      * exactly 1 `recv:` task,
      * exactly 1 `drive:` task,
      * exactly 1 (or 2 for primary-style filters) `subscribe` sub.
    Any growth means a reconnect leaks a coroutine or a sub registration.
    """
    cycles_seen = threading.Event()
    cycle_count = [0]

    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EOSE", msg[1]]))
            cycle_count[0] += 1
            if cycle_count[0] >= 5:
                cycles_seen.set()
            # Drop the WS so the bus reconnects.
            await asyncio.sleep(0.05)
            await ws.close(code=1011, reason="cycle-drop")

    relay = MockRelay(handler)
    relay.start()
    try:
        handle = bus.subscribe_sync(relay.url, {"kinds": [1]}, lambda e: None)
        # Wait until the bus has cycled through at least 5 drop/reconnect
        # rounds — give it more headroom than the cycle count so flakey
        # CI doesn't time-bomb us.
        assert cycles_seen.wait(timeout=10.0), (
            f"server-drop cycle never observed enough times: count={cycle_count[0]}"
        )
        s = bus.stats()
        assert _registry_size(bus) == 1, f"registry grew: {s}"
        assert s["tasks_by_name"].get("recv", 0) == 1, s
        assert s["tasks_by_name"].get("drive", 0) == 1, s
        assert s["subs_by_mode"]["subscribe"] == 1, s

        handle.cancel(timeout=2.0)
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"didn't drain after reconnect-storm cancel: {bus.stats()}"
        )
    finally:
        relay.stop()


def test_event_callback_exception_does_not_leak(bus, fast_grace):
    """The `on_event` callback raises on every call. The bus must keep
    delivering and the cancel path must still drain. Failure direction:
    `_drive` swallows the exception (we already do via try/except), but a
    queue-side or sub-side handler getting confused would manifest here.
    """
    fired = threading.Event()
    counter = [0]

    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            sub_id = msg[1]
            for i in range(5):
                await ws.send(json.dumps(["EVENT", sub_id, {
                    "id": f"e{i}", "kind": 1, "content": "x", "created_at": i,
                    "pubkey": "p", "sig": "s", "tags": [],
                }]))
            await ws.send(json.dumps(["EOSE", sub_id]))
            await asyncio.sleep(30)

    def on_event(_evt):
        counter[0] += 1
        fired.set()
        raise RuntimeError("synthetic callback failure")

    relay = MockRelay(handler)
    relay.start()
    try:
        handle = bus.subscribe_sync(relay.url, {"kinds": [1]}, on_event)
        assert fired.wait(timeout=2.0)
        # Give the bus a beat to push the rest of the events through the
        # raising callback — we want it to keep working, not bail out
        # after one exception.
        time.sleep(0.3)
        assert counter[0] >= 2, (
            f"bus stopped delivering after callback raised: count={counter[0]}"
        )
        handle.cancel(timeout=2.0)
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"raising-callback path leaked: {bus.stats()}"
        )
    finally:
        relay.stop()


def test_bus_stop_with_active_subscribers_completes_quickly(fast_grace):
    """A separate bus instance with live subs should `.stop()` quickly and
    leave its loop closed. Failure here means shutdown is waiting for some
    task that never settles."""
    b = relay_bus.RelayBus()
    b.start()
    try:
        async def handler(relay, ws, msg):
            if msg[0] == "REQ":
                await ws.send(json.dumps(["EOSE", msg[1]]))
                await asyncio.sleep(30)

        relay = MockRelay(handler)
        relay.start()
        try:
            handles = [
                b.subscribe_sync(relay.url, {"kinds": [i]}, lambda e: None)
                for i in range(5)
            ]
            time.sleep(0.2)
            assert b.stats()["subscriptions"] == 5
        finally:
            t0 = time.monotonic()
            b.stop(timeout=3.0)
            elapsed = time.monotonic() - t0
            relay.stop()
        # The bus should have torn down well before its own stop-timeout.
        assert elapsed < 2.0, f"bus.stop took too long: {elapsed:.2f}s"
        # Loop and thread must both be cleared.
        assert b._thread is None and b._loop is None
        del handles
    except Exception:
        try:
            b.stop(timeout=1.0)
        except Exception:
            pass
        raise


def test_immediate_cancel_after_subscribe_sync(bus, fast_grace):
    """Cancel the handle the instant `subscribe_sync` returns — before the
    `_drive` coroutine has a chance to register a sub. The bus must end
    up clean regardless of which side of the race wins."""
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EOSE", msg[1]]))
            await asyncio.sleep(30)

    relay = MockRelay(handler)
    relay.start()
    try:
        for _ in range(10):
            h = bus.subscribe_sync(relay.url, {"kinds": [1]}, lambda e: None)
            h.cancel(timeout=1.0)
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"rapid subscribe/cancel leaked: {bus.stats()}"
        )
    finally:
        relay.stop()


def test_many_subscribers_partial_cancel_keeps_conn(bus, fast_grace):
    """20 subs on the same URL share one WS. Cancelling half must keep the
    conn open (10 still alive); cancelling the rest must drain."""
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EOSE", msg[1]]))
            await asyncio.sleep(30)

    relay = MockRelay(handler)
    relay.start()
    try:
        handles = [
            bus.subscribe_sync(relay.url, {"kinds": [i]}, lambda e: None)
            for i in range(20)
        ]
        time.sleep(0.3)
        assert bus.stats()["subscriptions"] == 20
        assert _registry_size(bus) == 1
        # Exactly one server-side connection — all 20 subs multiplex.
        assert len(relay._connections) == 1

        for h in handles[:10]:
            h.cancel(timeout=1.0)
        # 10 still alive; conn must still be there.
        assert _wait_for(
            lambda: bus.stats()["subscriptions"] == 10, timeout=2.0,
        )
        assert _registry_size(bus) == 1, bus.stats()

        for h in handles[10:]:
            h.cancel(timeout=1.0)
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0), (
            f"didn't drain after all-cancel: {bus.stats()}"
        )
    finally:
        relay.stop()


def test_silent_handshake_eventually_connect_timeouts(bus, monkeypatch):
    """Server accepts TCP but never sends the 101 response. The bus's
    `CONNECT_TIMEOUT_S` must fire and surface `RelayConnectError`."""
    monkeypatch.setattr(relay_bus, "CONNECT_TIMEOUT_S", 0.4)

    server = HangingRelay(send_handshake=False)
    server.start()
    try:
        t0 = time.monotonic()
        with pytest.raises(relay_bus.RelayConnectError):
            bus.req_sync(server.url, {"kinds": [1]}, timeout=2.0)
        elapsed = time.monotonic() - t0
        # Should fail close to CONNECT_TIMEOUT_S, definitely under 2 s.
        assert elapsed < 1.5, f"connect didn't time out promptly: {elapsed:.2f}s"
        assert server.url not in bus._connections, (
            f"silent-handshake URL stayed in registry: {list(bus._connections)}"
        )
    finally:
        server.stop()


def test_random_walk_workload_drains(bus, fast_grace):
    """Fuzz: many threads doing a random mix of req / subscribe+cancel /
    publish-ok / publish-reject against several relays. After the workload
    settles, the bus must drain to a clean state.

    This is the test most likely to catch a subtle leak the targeted tests
    miss — interleavings of cancel + drop + new-sub on the same URL trip
    paths the focused tests can't easily reach.
    """
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            sub_id = msg[1]
            await ws.send(json.dumps(["EVENT", sub_id, {
                "id": "x", "kind": 1, "content": "n", "created_at": 1,
                "pubkey": "p", "sig": "s", "tags": [],
            }]))
            await ws.send(json.dumps(["EOSE", sub_id]))
        elif msg[0] == "EVENT":
            # Echo OK true for any publish (some tests force-reject via
            # event id).
            eid = msg[1].get("id", "")
            ok = not eid.startswith("REJ")
            await ws.send(json.dumps(["OK", eid, ok, "" if ok else "fuzz-reject"]))

    relays: list[MockRelay] = []
    try:
        for _ in range(4):
            r = MockRelay(handler)
            r.start()
            relays.append(r)

        rng = random.Random(0xFEEDFACE)
        errors: list[tuple[str, str]] = []

        def _op():
            r = rng.choice(relays)
            choice = rng.choice(("req", "subscribe", "publish_ok", "publish_rej"))
            try:
                if choice == "req":
                    bus.req_sync(r.url, {"kinds": [1]}, timeout=1.5)
                elif choice == "subscribe":
                    h = bus.subscribe_sync(r.url, {"kinds": [1]}, lambda e: None)
                    time.sleep(rng.uniform(0.02, 0.15))
                    h.cancel(timeout=1.0)
                elif choice == "publish_ok":
                    bus.publish_sync(r.url, {
                        "id": f"OK{rng.randrange(1_000_000):06d}", "kind": 1,
                        "content": "f", "created_at": 1, "pubkey": "p",
                        "sig": "s", "tags": [],
                    }, timeout=1.5)
                else:
                    try:
                        bus.publish_sync(r.url, {
                            "id": f"REJ{rng.randrange(1_000_000):06d}",
                            "kind": 1, "content": "f", "created_at": 1,
                            "pubkey": "p", "sig": "s", "tags": [],
                        }, timeout=1.5)
                    except relay_bus.RelayPublishRejected:
                        pass
            except Exception as e:
                errors.append((choice, repr(e)))

        threads = [threading.Thread(target=_op) for _ in range(60)]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=10.0)

        assert errors == [], f"unexpected errors during fuzz: {errors[:5]}"
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=3.0), (
            f"fuzz workload left registry dirty: {bus.stats()}"
        )
        s = bus.stats()
        assert s["tasks_by_name"].get("recv", 0) == 0, s
        assert s["tasks_by_name"].get("drive", 0) == 0, s
        assert s["subs_by_mode"] == {"req": 0, "subscribe": 0}, s
    finally:
        for r in relays:
            r.stop()


def test_close_then_resubscribe_same_url_works(bus, fast_grace):
    """Cancel a sub, wait until the conn drops, then re-subscribe to the
    same URL. The bus must create a fresh `_RelayConnection`, not try to
    revive the torn-down one (which would raise from `ensure_open`)."""
    async def handler(relay, ws, msg):
        if msg[0] == "REQ":
            await ws.send(json.dumps(["EOSE", msg[1]]))
            await asyncio.sleep(30)

    relay = MockRelay(handler)
    relay.start()
    try:
        h1 = bus.subscribe_sync(relay.url, {"kinds": [1]}, lambda e: None)
        time.sleep(0.1)
        h1.cancel(timeout=1.0)
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0)

        # Now re-subscribe — must work.
        h2 = bus.subscribe_sync(relay.url, {"kinds": [2]}, lambda e: None)
        time.sleep(0.1)
        assert bus.stats()["subscriptions"] == 1
        h2.cancel(timeout=1.0)
        assert _wait_for(lambda: _registry_size(bus) == 0, timeout=2.0)
    finally:
        relay.stop()
