"""Single async client for all NIP-01 traffic with Nostr relays.

The old code held one ``websockets.sync`` WS per use case (ingest, live
subscription, ad-hoc query, repair) and ran 1–2 OS threads per connection.
Under load that produced thousands of threads racing on the same SSL object,
which is undefined behaviour in OpenSSL and triggered native ``abort()``
faults at runtime.

``RelayBus`` replaces all four clients with a single asyncio component:

* one event-loop thread for the whole process,
* one WS per relay URL (``_RelayConnection``),
* arbitrary number of overlapping NIP-01 REQ subscriptions per WS, each
  identified by its own ``sub_id``,
* one in-process queue per subscription, fed by a single recv coroutine,
* automatic reconnect with exponential backoff for long-lived subs.

Public sync facade (``req_sync`` / ``subscribe_sync`` / ``publish_sync``)
uses ``asyncio.run_coroutine_threadsafe`` to bridge into the loop. Callers
already running inside an asyncio loop can use the ``await``-friendly
methods directly.
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import threading
import uuid
from collections.abc import AsyncIterator
from concurrent.futures import Future
from dataclasses import dataclass, field
from typing import Any, Callable

import websockets
from websockets.asyncio.client import connect as ws_connect
from websockets.exceptions import ConnectionClosed, WebSocketException

log = logging.getLogger("relay_bus")

CONNECT_TIMEOUT_S = 10.0
PING_INTERVAL_S = 30.0
PING_TIMEOUT_S = 10.0
REQ_TIMEOUT_DEFAULT_S = 10.0
PUBLISH_TIMEOUT_DEFAULT_S = 5.0
RECONNECT_BACKOFF_INITIAL_S = 1.0
RECONNECT_BACKOFF_MAX_S = 60.0
# How long to keep a relay connection open after its last subscription closes.
# A short grace window avoids tearing down and reopening the same WS when an
# ad-hoc REQ is immediately followed by another (e.g. probe loops).
IDLE_CLOSE_GRACE_S = 5.0
# Upper bound on the close-handshake wait. A peer that goes silent can hang
# `ClientConnection.close()` indefinitely past the library's own internal
# timer, leaving the keepalive task + socket fd alive. After this we abort
# the underlying transport so those resources actually get freed.
WS_CLOSE_TIMEOUT_S = 3.0


# ---- public exceptions -----------------------------------------------------


class RelayBusError(Exception):
    """Base for all RelayBus-originated failures."""


class RelayConnectError(RelayBusError):
    """Failed to open a websocket to a relay."""


class RelayTimeout(RelayBusError):
    """A REQ or publish did not complete within the deadline."""


class RelayPublishRejected(RelayBusError):
    """Relay returned ``OK <id> false <message>`` for a published event."""


# ---- internal: per-subscription state --------------------------------------


@dataclass
class _Subscription:
    sub_id: str
    filter: dict
    mode: str  # "req" (close on EOSE) | "subscribe" (live)
    queue: asyncio.Queue
    # Set when EOSE arrives for a "req" sub, or when the connection drops
    # mid-fetch — readers await this to know they should stop pulling from the
    # queue. For "subscribe" subs, EOSE is forwarded as a sentinel but the sub
    # stays open.
    eose: asyncio.Event = field(default_factory=asyncio.Event)
    closed: asyncio.Event = field(default_factory=asyncio.Event)
    # Number of events forwarded — used by subscribe-mode to update the
    # ``since`` cursor on reconnect so we don't replay history.
    last_event_at: int | None = None


# Sentinel placed on a subscription's queue to signal "the relay said EOSE".
# Distinct object identity, not a dict, so subscribers can pattern-match it
# without colliding with a real event.
EOSE_SENTINEL: dict = {"__eose__": True}


# ---- internal: per-relay connection ----------------------------------------


class _RelayConnection:
    """One WS to one relay URL. Multiplexes any number of NIP-01 REQs.

    All async methods on this object MUST be called from the bus event loop
    thread; cross-thread callers go through ``RelayBus`` sync facade.
    """

    def __init__(
        self,
        url: str,
        loop: asyncio.AbstractEventLoop,
        on_done: Callable[[str], Any] | None = None,
    ) -> None:
        self.url = url
        self._loop = loop
        self._ws: Any | None = None  # websockets.asyncio.client.ClientConnection
        self._subs: dict[str, _Subscription] = {}
        # ``OK`` responses for outgoing EVENTs — keyed by event id, value is
        # the future the publish coroutine is waiting on.
        self._pending_ok: dict[str, asyncio.Future] = {}
        self._send_lock = asyncio.Lock()
        # Serializes connect / reconnect with concurrent ``ensure_open``
        # callers. Without it, two callers can both observe ``_opened`` as
        # clear, race past the guard, each open its own WS and start its own
        # recv task — two coroutines then call ``recv()`` on the same
        # connection and the websockets library raises ConcurrencyError.
        self._open_lock = asyncio.Lock()
        self._recv_task: asyncio.Task | None = None
        # Connection is considered "open" once the recv task is running. The
        # ``opened`` event lets request coroutines wait for the WS to become
        # available without racing on ``_ws is None``.
        self._opened = asyncio.Event()
        self._close_requested = False
        self._idle_close_task: asyncio.Task | None = None
        # Async callback the bus passes in so a self-terminating recv loop
        # can evict its own registry entry. Without this the connection
        # lingers in `bus._connections` as a zombie (no recv task, no WS)
        # and accumulates indefinitely once relays drop their subs.
        self._on_done = on_done

    # ---- lifecycle ----------------------------------------------------------

    async def ensure_open(self) -> None:
        """Open the WS if it isn't already. Idempotent and concurrency-safe.

        Raises ``RelayConnectError`` on first-time connect failure. Long-lived
        reconnects after an initial successful connect are handled inside
        ``_run`` and do NOT raise to the caller — they just flip ``_opened``
        again when they succeed.
        """
        async with self._open_lock:
            if self._close_requested:
                # Conn has been torn down (either explicit close() or recv
                # loop self-terminated). Refuse to revive — bus._get_or_open
                # will replace us with a fresh entry on the caller's retry.
                raise RelayConnectError(f"{self.url}: conn torn down")
            self._cancel_idle_close()
            if self._opened.is_set():
                return
            # A reconnect cycle inside the existing recv task may already be
            # in flight. Wait briefly for it to flip ``_opened``; only do our
            # own connect if it doesn't recover in time.
            if self._recv_task is not None and not self._recv_task.done():
                try:
                    await asyncio.wait_for(
                        self._opened.wait(), timeout=CONNECT_TIMEOUT_S,
                    )
                    return
                except TimeoutError:
                    pass
            try:
                self._ws = await asyncio.wait_for(
                    ws_connect(
                        self.url,
                        open_timeout=CONNECT_TIMEOUT_S,
                        ping_interval=PING_INTERVAL_S,
                        ping_timeout=PING_TIMEOUT_S,
                        max_size=2**20,  # 1 MiB; raise if real events overflow
                    ),
                    timeout=CONNECT_TIMEOUT_S,
                )
            except (TimeoutError, OSError, WebSocketException) as e:
                raise RelayConnectError(f"connect {self.url}: {e}") from e
            self._opened.set()
            if self._recv_task is None or self._recv_task.done():
                self._recv_task = asyncio.create_task(
                    self._run(), name=f"recv:{self.url}",
                )
            log.info("connected %s", self.url)

    async def close(self) -> None:
        """Tear down the WS and cancel the recv task. Idempotent."""
        if self._close_requested:
            return
        self._close_requested = True
        self._cancel_idle_close()
        # Wake any subscribers that are still waiting on this conn.
        for sub in self._subs.values():
            sub.closed.set()
            with contextlib.suppress(Exception):
                sub.queue.put_nowait(EOSE_SENTINEL)
        self._subs.clear()
        for fut in self._pending_ok.values():
            if not fut.done():
                fut.set_exception(RelayConnectError(f"{self.url}: closed"))
        self._pending_ok.clear()
        await self._close_ws_safely()
        self._opened.clear()
        if self._recv_task is not None:
            self._recv_task.cancel()
            with contextlib.suppress(asyncio.CancelledError, Exception):
                await self._recv_task
            self._recv_task = None

    async def _close_ws_safely(self) -> None:
        """Close ``_ws`` with a hard timeout and a transport-abort fallback.

        The `websockets` library's ``close()`` waits for the close handshake;
        a peer that goes silent can stall it past its own internal timer,
        and the keepalive task it spawned keeps the socket fd alive even
        after we drop our reference to the connection. Bounding the wait
        and forcing ``transport.abort()`` on hang lets the keepalive task
        observe the dead socket and exit, so the fd is actually released.
        """
        if self._ws is None:
            return
        ws = self._ws
        self._ws = None
        try:
            await asyncio.wait_for(ws.close(), timeout=WS_CLOSE_TIMEOUT_S)
        except asyncio.CancelledError:
            self._abort_transport(ws)
            raise
        except Exception as e:
            log.info("ws.close %s did not finish (%s); aborting transport",
                     self.url, e)
            self._abort_transport(ws)

    @staticmethod
    def _abort_transport(ws: Any) -> None:
        with contextlib.suppress(Exception):
            t = getattr(ws, "transport", None)
            if t is not None:
                t.abort()

    def _cancel_idle_close(self) -> None:
        if self._idle_close_task is not None and not self._idle_close_task.done():
            self._idle_close_task.cancel()
        self._idle_close_task = None

    def schedule_idle_close_if_empty(
        self, on_close: Callable[[str], Any],
    ) -> None:
        """If no subscriptions remain, schedule a delayed close.

        ``on_close`` is an async function (``async def``) that removes this
        connection from the bus registry.
        """
        if self._subs or self._pending_ok:
            return
        self._cancel_idle_close()

        async def _delayed():
            try:
                await asyncio.sleep(IDLE_CLOSE_GRACE_S)
            except asyncio.CancelledError:
                return
            if self._subs or self._pending_ok or self._close_requested:
                return
            await self.close()
            await on_close(self.url)

        self._idle_close_task = asyncio.create_task(
            _delayed(), name=f"idle-close:{self.url}",
        )

    # ---- send helpers -------------------------------------------------------

    async def _send(self, frame: list) -> None:
        # The asyncio websockets API serializes writes per connection, but
        # we still want explicit mutual exclusion so that a half-written
        # frame can't be interleaved with a CLOSE during teardown.
        if self._ws is None:
            raise RelayConnectError(f"{self.url}: not connected")
        payload = json.dumps(frame)
        async with self._send_lock:
            try:
                await self._ws.send(payload)
            except (ConnectionClosed, OSError) as e:
                raise RelayConnectError(f"{self.url}: send failed: {e}") from e

    # ---- REQ / subscribe ---------------------------------------------------

    async def open_sub(self, flt: dict, *, mode: str) -> _Subscription:
        await self.ensure_open()
        sub_id = "bus_" + uuid.uuid4().hex[:12]
        sub = _Subscription(
            sub_id=sub_id, filter=dict(flt), mode=mode,
            queue=asyncio.Queue(maxsize=0),
        )
        self._subs[sub_id] = sub
        try:
            await self._send(["REQ", sub_id, flt])
        except Exception:
            self._subs.pop(sub_id, None)
            raise
        return sub

    async def close_sub(self, sub: _Subscription) -> None:
        self._subs.pop(sub.sub_id, None)
        sub.closed.set()
        if self._ws is not None:
            with contextlib.suppress(Exception):
                await self._send(["CLOSE", sub.sub_id])

    # ---- publish -----------------------------------------------------------

    async def publish(self, event: dict, *, timeout: float) -> None:
        await self.ensure_open()
        eid = event.get("id")
        if not isinstance(eid, str) or not eid:
            raise RelayBusError("publish: event missing 'id'")
        fut: asyncio.Future = self._loop.create_future()
        self._pending_ok[eid] = fut
        try:
            await self._send(["EVENT", event])
            try:
                ok, msg = await asyncio.wait_for(fut, timeout=timeout)
            except TimeoutError as e:
                raise RelayTimeout(f"{self.url}: no OK for event {eid[:8]}") from e
            if not ok:
                raise RelayPublishRejected(f"{self.url}: {msg or '(no message)'}")
        finally:
            self._pending_ok.pop(eid, None)

    # ---- recv loop ---------------------------------------------------------

    async def _run(self) -> None:
        """Main recv loop. On WS drop, walks through reconnect/backoff and
        re-issues every still-open REQ so subscribe-mode callers see a
        seamless stream (modulo a brief gap)."""
        backoff = RECONNECT_BACKOFF_INITIAL_S
        first_pass = True
        try:
            while not self._close_requested:
                if not first_pass:
                    try:
                        await self._reconnect()
                    except RelayConnectError as e:
                        log.warning("reconnect %s failed: %s; backoff %.1fs",
                                    self.url, e, backoff)
                        try:
                            await asyncio.sleep(backoff)
                        except asyncio.CancelledError:
                            return
                        backoff = min(backoff * 2, RECONNECT_BACKOFF_MAX_S)
                        continue
                    backoff = RECONNECT_BACKOFF_INITIAL_S
                first_pass = False

                try:
                    async for raw in self._ws:  # type: ignore[union-attr]
                        self._handle_message(raw)
                except ConnectionClosed as e:
                    log.info("ws %s closed: %s", self.url, e)
                except Exception:
                    log.exception("recv loop error %s", self.url)
                # Tell subscribers their stream paused; subscribe-mode will
                # resume after reconnect, req-mode treats this as failure.
                for sub in list(self._subs.values()):
                    if sub.mode == "req":
                        sub.closed.set()
                        with contextlib.suppress(Exception):
                            sub.queue.put_nowait(EOSE_SENTINEL)
                # Fail any in-flight publishes; the relay can't ack them now.
                for fut in self._pending_ok.values():
                    if not fut.done():
                        fut.set_exception(RelayConnectError(f"{self.url}: dropped"))
                self._pending_ok.clear()
                # Drop any "req" subs that asked us to close; "subscribe" subs
                # remain so we can re-issue REQ after reconnect.
                self._subs = {sid: s for sid, s in self._subs.items() if s.mode != "req"}
                self._opened.clear()
                await self._close_ws_safely()
                if not self._subs:
                    # Nothing to keep us around for. Flag the conn as torn
                    # down and tell the bus to drop us from its registry so
                    # a future req/subscribe gets a fresh _RelayConnection
                    # instead of inheriting this zombie (no recv task, no
                    # WS). Without this the registry grows unboundedly:
                    # every reconnect cycle that ends with subs == 0 leaks
                    # one entry + the keepalive socket the closed WS was
                    # still pinning.
                    self._close_requested = True
                    if self._on_done is not None:
                        try:
                            await self._on_done(self.url)
                        except Exception:
                            log.exception(
                                "recv-loop on_done failed for %s", self.url,
                            )
                    return
        finally:
            self._opened.clear()

    async def _reconnect(self) -> None:
        try:
            self._ws = await asyncio.wait_for(
                ws_connect(
                    self.url,
                    open_timeout=CONNECT_TIMEOUT_S,
                    ping_interval=PING_INTERVAL_S,
                    ping_timeout=PING_TIMEOUT_S,
                    max_size=2**20,
                ),
                timeout=CONNECT_TIMEOUT_S,
            )
        except (TimeoutError, OSError, WebSocketException) as e:
            raise RelayConnectError(f"reconnect {self.url}: {e}") from e
        self._opened.set()
        log.info("reconnected %s", self.url)
        # Replay every still-live subscription. For subscribe-mode subs we
        # advance ``since`` past the most recent event we've already seen so
        # the relay doesn't redeliver the entire history.
        for sub in list(self._subs.values()):
            flt = dict(sub.filter)
            if sub.last_event_at is not None:
                flt["since"] = sub.last_event_at + 1
            try:
                await self._send(["REQ", sub.sub_id, flt])
            except RelayConnectError:
                # The send raised — let the outer loop detect the drop on the
                # next iteration and retry.
                return

    def _handle_message(self, raw: str | bytes) -> None:
        try:
            msg = json.loads(raw)
        except (TypeError, ValueError):
            return
        if not isinstance(msg, list) or not msg:
            return
        kind = msg[0]
        if kind == "EVENT" and len(msg) >= 3:
            sub_id, evt = msg[1], msg[2]
            sub = self._subs.get(sub_id)
            if sub is None or not isinstance(evt, dict):
                return
            ts = evt.get("created_at")
            if isinstance(ts, int):
                if sub.last_event_at is None or ts > sub.last_event_at:
                    sub.last_event_at = ts
            try:
                sub.queue.put_nowait(evt)
            except asyncio.QueueFull:
                # We don't bound queues today, but defensively log if we ever do.
                log.warning("sub %s queue full; dropped event", sub_id)
        elif kind == "EOSE" and len(msg) >= 2:
            sub = self._subs.get(msg[1])
            if sub is None:
                return
            sub.eose.set()
            if sub.mode == "req":
                # req-mode: signal completion via sentinel and let the reader
                # close the sub and remove it.
                with contextlib.suppress(Exception):
                    sub.queue.put_nowait(EOSE_SENTINEL)
        elif kind == "CLOSED" and len(msg) >= 2:
            sub = self._subs.pop(msg[1], None)
            if sub is None:
                return
            sub.closed.set()
            with contextlib.suppress(Exception):
                sub.queue.put_nowait(EOSE_SENTINEL)
        elif kind == "NOTICE" and len(msg) >= 2:
            log.info("%s NOTICE: %s", self.url, msg[1])
        elif kind == "OK" and len(msg) >= 3:
            eid = msg[1]
            ok = bool(msg[2])
            err = msg[3] if len(msg) >= 4 and isinstance(msg[3], str) else ""
            fut = self._pending_ok.pop(eid, None)
            if fut is not None and not fut.done():
                fut.set_result((ok, err))


# ---- public bus -----------------------------------------------------------


class RelayBus:
    """Process-wide client for all NIP-01 relay traffic.

    Start once at process boot, stop on shutdown. Methods are safe to call
    from any thread; the async ones must be awaited from within the bus loop
    (FastMCP tool handlers can ``await`` them directly when the bus loop is
    the same one FastMCP runs in).
    """

    def __init__(self) -> None:
        self._loop: asyncio.AbstractEventLoop | None = None
        self._thread: threading.Thread | None = None
        self._ready = threading.Event()
        self._connections: dict[str, _RelayConnection] = {}
        # Guards ``_connections`` access from the loop thread itself; calls
        # from outside threads always go through ``run_coroutine_threadsafe``
        # so they're naturally serialized on the loop.
        self._reg_lock = asyncio.Lock  # placeholder; created in loop

    # ---- lifecycle ----------------------------------------------------------

    def start(self) -> None:
        if self._thread is not None:
            return
        self._thread = threading.Thread(
            target=self._run_loop, name="relay-bus-loop", daemon=True,
        )
        self._thread.start()
        if not self._ready.wait(timeout=5.0):
            raise RelayBusError("relay bus loop failed to start")

    def stop(self, timeout: float = 5.0) -> None:
        if self._loop is None or self._thread is None:
            return
        fut = asyncio.run_coroutine_threadsafe(self._shutdown(), self._loop)
        try:
            fut.result(timeout=timeout)
        except Exception:
            log.exception("relay bus shutdown error")
        self._loop.call_soon_threadsafe(self._loop.stop)
        self._thread.join(timeout=timeout)
        self._thread = None
        self._loop = None

    def _run_loop(self) -> None:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        self._loop = loop
        self._reg_lock = asyncio.Lock()
        self._ready.set()
        try:
            loop.run_forever()
        finally:
            loop.close()

    async def _shutdown(self) -> None:
        urls = list(self._connections.keys())
        for url in urls:
            conn = self._connections.pop(url, None)
            if conn is not None:
                with contextlib.suppress(Exception):
                    await conn.close()

    # ---- registry ----------------------------------------------------------

    async def _get_or_open(self, url: str) -> _RelayConnection:
        async with self._reg_lock:  # type: ignore[union-attr]
            conn = self._connections.get(url)
            if conn is not None and conn._close_requested:
                # Recv loop self-terminated between `_drop` and our lookup,
                # or we caught it mid-teardown. Replace pre-emptively so
                # the caller doesn't get a conn whose `ensure_open` will
                # just raise.
                self._connections.pop(url, None)
                conn = None
            we_created = conn is None
            if we_created:
                conn = _RelayConnection(
                    url, self._loop,  # type: ignore[arg-type]
                    on_done=self._drop,
                )
                self._connections[url] = conn
        try:
            await conn.ensure_open()
        except (Exception, asyncio.CancelledError):
            # A failed connect on a freshly-created entry must NOT leave
            # the half-built conn lingering in the registry — otherwise
            # every transient outage at startup (refused, handshake
            # timeout, dns) turns into a permanent zombie that the next
            # caller will re-use and re-fail forever. Only evict the entry
            # if it's still ours (concurrent callers may have replaced it).
            if we_created:
                try:
                    async with self._reg_lock:  # type: ignore[union-attr]
                        if self._connections.get(url) is conn:
                            self._connections.pop(url, None)
                except Exception:
                    pass
            raise
        return conn

    async def _drop(self, url: str) -> None:
        async with self._reg_lock:  # type: ignore[union-attr]
            self._connections.pop(url, None)

    # ---- async API ---------------------------------------------------------

    async def req(
        self, url: str, filter: dict, *, timeout: float = REQ_TIMEOUT_DEFAULT_S,
    ) -> list[dict]:
        """Send a NIP-01 REQ, return all EVENTs up to EOSE or ``timeout``.

        After EOSE the bus sends ``CLOSE`` automatically. On timeout it sends
        ``CLOSE`` too and returns whatever events arrived so far.
        """
        conn = await self._get_or_open(url)
        sub = await conn.open_sub(filter, mode="req")
        events: list[dict] = []
        deadline = self._loop.time() + timeout  # type: ignore[union-attr]
        try:
            while True:
                remaining = deadline - self._loop.time()  # type: ignore[union-attr]
                if remaining <= 0:
                    break
                try:
                    item = await asyncio.wait_for(sub.queue.get(), timeout=remaining)
                except TimeoutError:
                    break
                if item is EOSE_SENTINEL:
                    break
                events.append(item)
        finally:
            await conn.close_sub(sub)
            conn.schedule_idle_close_if_empty(self._drop)
        return events

    async def subscribe(
        self, url: str, filter: dict,
    ) -> AsyncIterator[dict]:
        """Open a long-lived subscription. Yields events as they arrive. The
        subscription is automatically re-issued (with an updated ``since``)
        after reconnects. ``CLOSE`` is sent when the consumer stops iterating
        (the generator is closed)."""
        conn = await self._get_or_open(url)
        sub = await conn.open_sub(filter, mode="subscribe")
        try:
            while not sub.closed.is_set():
                item = await sub.queue.get()
                # Live mode: EOSE means "you're caught up to now". Don't end the
                # stream — pass the sentinel through so the consumer can stamp a
                # coverage watermark, then keep waiting for further pushes.
                yield item
        finally:
            await conn.close_sub(sub)
            conn.schedule_idle_close_if_empty(self._drop)

    async def publish(
        self, url: str, event: dict, *, timeout: float = PUBLISH_TIMEOUT_DEFAULT_S,
    ) -> None:
        """Send an ``EVENT`` frame and wait for the relay's ``OK`` ack.

        Raises ``RelayTimeout`` if no OK arrives within ``timeout`` and
        ``RelayPublishRejected`` if the relay returns ``OK <id> false``.
        """
        conn = await self._get_or_open(url)
        try:
            await conn.publish(event, timeout=timeout)
        finally:
            conn.schedule_idle_close_if_empty(self._drop)

    # ---- sync facade -------------------------------------------------------

    def _submit(self, coro) -> Future:
        if self._loop is None:
            raise RelayBusError("bus not started")
        return asyncio.run_coroutine_threadsafe(coro, self._loop)

    def req_sync(
        self, url: str, filter: dict, *, timeout: float = REQ_TIMEOUT_DEFAULT_S,
    ) -> list[dict]:
        fut = self._submit(self.req(url, filter, timeout=timeout))
        return fut.result(timeout=timeout + 5.0)

    def publish_sync(
        self, url: str, event: dict, *, timeout: float = PUBLISH_TIMEOUT_DEFAULT_S,
    ) -> None:
        fut = self._submit(self.publish(url, event, timeout=timeout))
        fut.result(timeout=timeout + 5.0)

    def stats(self) -> dict:
        """Cheap in-memory snapshot for observability. Thread-safe — reads
        only the registry dict and per-connection counters; Python dict
        reads are atomic and len() tolerates concurrent mutation by
        sampling slightly stale numbers.

        Shape:
          {"connections": N,             # open + reconnecting WSes
           "subscriptions": N,           # live + in-flight req subs total
           "pending_publishes": N,       # EVENTs awaiting an OK ack
           "asyncio_tasks": N,           # all_tasks on the bus loop
           "tasks_by_name": {prefix: N}, # histogram of task name prefixes
           "connections_by_state": {     # zombie-detection breakdown
               "open":           N,      # _opened set, recv task alive
               "reconnecting":   N,      # recv task alive but _opened cleared
               "idle_close_armed": N,    # idle-close timer running
               "stale":          N,      # no recv task and no ws (should be 0)
           },
           "subs_by_mode": {"req": N, "subscribe": N}}

        The histogram is the key leak-attribution lever: `recv:` should
        track `connections`, `idle-close:` should be bounded by the same,
        and an unbounded `drive:` count means `subscribe_sync` callers
        aren't cancelling their handles. A growing `Task-` bucket points
        at unnamed tasks the rest of this module forgot to label.
        """
        conns = list(self._connections.values())
        subs_total = sum(len(c._subs) for c in conns)
        pending_total = sum(len(c._pending_ok) for c in conns)

        state = {"open": 0, "reconnecting": 0, "idle_close_armed": 0, "stale": 0}
        modes: dict[str, int] = {"req": 0, "subscribe": 0}
        for c in conns:
            recv_alive = c._recv_task is not None and not c._recv_task.done()
            if c._opened.is_set() and recv_alive:
                state["open"] += 1
            elif recv_alive:
                state["reconnecting"] += 1
            elif c._ws is None and c._recv_task is None:
                state["stale"] += 1
            else:
                state["stale"] += 1
            if c._idle_close_task is not None and not c._idle_close_task.done():
                state["idle_close_armed"] += 1
            for sub in c._subs.values():
                modes[sub.mode] = modes.get(sub.mode, 0) + 1

        task_count = 0
        tasks_by_name: dict[str, int] = {}
        if self._loop is not None:
            try:
                tasks = list(asyncio.all_tasks(self._loop))
            except RuntimeError:
                tasks = []
            task_count = len(tasks)
            for t in tasks:
                raw = t.get_name() or "?"
                # Bucket by stable prefix: "recv:wss://relay/" → "recv",
                # "Task-1234" → "Task". Keeps the histogram compact.
                if ":" in raw:
                    bucket = raw.split(":", 1)[0]
                elif raw.startswith("Task-"):
                    bucket = "Task"
                else:
                    bucket = raw
                tasks_by_name[bucket] = tasks_by_name.get(bucket, 0) + 1

        return {
            "connections": len(conns),
            "subscriptions": subs_total,
            "pending_publishes": pending_total,
            "asyncio_tasks": task_count,
            "tasks_by_name": tasks_by_name,
            "connections_by_state": state,
            "subs_by_mode": modes,
        }

    def subscribe_sync(
        self, url: str, filter: dict,
        on_event: Callable[[dict], Any],
        on_eose: Callable[[], Any] | None = None,
    ) -> "_SubscriptionHandle":
        """Start a live subscription whose events are dispatched to
        ``on_event`` (called on the bus loop thread). Returns a handle whose
        ``cancel()`` tears down the underlying iterator.

        ``on_eose`` (optional) fires every time the relay signals end-of-stored-
        events for this sub — i.e. "you're caught up to now". Re-fired after
        each reconnect's re-issued REQ. Used to advance the coverage watermark.
        """
        if self._loop is None:
            raise RelayBusError("bus not started")
        handle = _SubscriptionHandle(self)

        async def _drive():
            asyncio.current_task().set_name(f"drive:{url}")
            # Bind the generator to a local so we can `aclose()` it
            # explicitly — without that, exiting via `break` or
            # CancelledError orphans the generator and its `finally`
            # (which calls `close_sub` + `schedule_idle_close_if_empty`,
            # i.e. the only thing that ever evicts a one-shot subscribe
            # from the bus registry) never runs. The bus loop runs via
            # `run_forever()` so the default asyncio async-gen finalizer
            # isn't installed; relying on GC here was the real driver
            # of the `connections.stale` leak we saw at runtime.
            gen = self.subscribe(url, filter)
            try:
                async for evt in gen:
                    if handle._cancelled.is_set():
                        break
                    if evt is EOSE_SENTINEL:
                        if on_eose is not None:
                            try:
                                on_eose()
                            except Exception:
                                log.exception("subscribe_sync on_eose raised")
                        continue
                    try:
                        on_event(evt)
                    except Exception:
                        log.exception("subscribe_sync callback raised")
            except asyncio.CancelledError:
                pass
            finally:
                with contextlib.suppress(Exception):
                    await gen.aclose()

        task = asyncio.run_coroutine_threadsafe(_drive(), self._loop)
        handle._task = task
        return handle


class _SubscriptionHandle:
    def __init__(self, bus: RelayBus) -> None:
        self._bus = bus
        self._task: Future | None = None
        self._cancelled = threading.Event()

    def cancel(self, timeout: float = 5.0) -> None:
        self._cancelled.set()
        if self._task is not None:
            self._task.cancel()
            with contextlib.suppress(Exception):
                self._task.result(timeout=timeout)


# ---- module-level singleton (optional convenience) -------------------------

_DEFAULT_BUS: RelayBus | None = None
_DEFAULT_BUS_LOCK = threading.Lock()


def get_default_bus() -> RelayBus:
    """Return a process-wide ``RelayBus`` instance, starting it on first use."""
    global _DEFAULT_BUS
    with _DEFAULT_BUS_LOCK:
        if _DEFAULT_BUS is None:
            _DEFAULT_BUS = RelayBus()
            _DEFAULT_BUS.start()
    return _DEFAULT_BUS


__all__ = [
    "RelayBus",
    "RelayBusError",
    "RelayConnectError",
    "RelayTimeout",
    "RelayPublishRejected",
    "get_default_bus",
    "EOSE_SENTINEL",
]
