"""Long-lived NIP-01 REQ subscriptions, layered on top of ``relay_bus``.

The bus owns one WS per relay URL and multiplexes any number of sub_ids over
it. Live subs are just registered subscriptions whose events feed the
daemon's ``_ingest_event`` hook. Reconnects with backoff + re-issue of REQ
(with an advanced ``since``) are handled inside the bus; this module only
holds the registry of (pubkey, relay, handle) tuples and the bookkeeping
for filter shapes and ``since`` cursors.

Filter sets per role:
  * ``primary``   — TWO subscriptions per (pubkey, relay): ``{authors:[hex]}``
                    (direction 'out') AND ``{"#p":[hex]}`` (direction 'in').
                    The bus uses one sub_id per filter, both sit on the same WS.
  * ``secondary`` — ONE subscription: ``{authors:[hex]}`` ('out'). 1-hop frontier.

``since`` per (pubkey, relay, direction) comes from the daemon's coverage
watermark (``_coverage_since``) — how far we KNOW we're caught up, minus a
safety margin — not from the last event seen. The relay's end-of-stored-events
signal feeds that watermark via ``on_eose`` → ``_mark_caught_up``.
"""

from __future__ import annotations

import logging
import threading
from dataclasses import dataclass
from typing import TYPE_CHECKING

from relay_bus import RelayBusError

if TYPE_CHECKING:
    from ingest_sql import IngestDaemonSQL
    from relay_bus import _SubscriptionHandle

log = logging.getLogger("live_subscription")


# Coverage directions. 'out' = the pubkey's own events (`authors` filter),
# 'in' = mentions of it (`#p` filter). A secondary only tracks 'out'.
DIR_OUT = "out"
DIR_IN = "in"


@dataclass
class _RelayEntry:
    """Per-(pubkey, relay) live-sub state: the bus handles plus a per-direction
    'have we received the end-of-stored-events signal yet' flag. The flag gates
    the coverage flush — only caught-up directions get their watermark stamped,
    so a sub still pulling history is never marked as covered."""
    handles: list["_SubscriptionHandle"]
    caught_up: dict[str, bool]


def _directions_for(role: str) -> tuple[str, ...]:
    return (DIR_OUT,) if role != "primary" else (DIR_OUT, DIR_IN)


def _filter_for(pubkey_hex: str, direction: str, limit: int,
                since: int | None) -> dict:
    """Build the NIP-01 filter for one (pubkey, direction)."""
    key = "authors" if direction == DIR_OUT else "#p"
    flt: dict = {key: [pubkey_hex], "limit": limit}
    if since is not None:
        flt["since"] = since
    return flt


class LiveSubscriptionPool:
    """Registry of live subscriptions, one entry per (pubkey, relay) pair.

    Each entry can hold multiple bus subscription handles (two for primary,
    one for secondary). ``remove_pubkey`` cancels every handle and lets the
    bus drop the underlying WS if nothing else is keeping it open.
    """

    def __init__(self, daemon: "IngestDaemonSQL") -> None:
        self.daemon = daemon
        # pubkey -> relay_url -> list of bus handles
        self._subs: dict[str, dict[str, list["_SubscriptionHandle"]]] = {}
        self._lock = threading.Lock()

    # ---- registration -------------------------------------------------------

    def add(
        self, pubkey_hex: str, relay_url: str, limit: int,
        role: str = "primary",
    ) -> bool:
        if role not in ("primary", "secondary"):
            raise ValueError(
                f"LiveSub role must be 'primary' or 'secondary', got {role!r}",
            )
        with self._lock:
            if pubkey_hex in self._subs and relay_url in self._subs[pubkey_hex]:
                raise RuntimeError(
                    f"live sub already active for {pubkey_hex[:6]}@{relay_url}",
                )

        # Callback runs on the bus event loop thread. `_ingest_event` is
        # serialized by the daemon's own storage lock so multiple bus subs
        # can ingest concurrently without corrupting DuckDB state.
        def _on_event(evt: dict) -> None:
            try:
                self.daemon._ingest_event(evt, relay_url)
            except Exception:
                log.exception(
                    "live sub %s@%s: ingest failed for event %s",
                    pubkey_hex[:6], relay_url, evt.get("id"),
                )

        directions = _directions_for(role)
        handles: list["_SubscriptionHandle"] = []
        caught_up: dict[str, bool] = {}
        bus = self.daemon.bus
        for direction in directions:
            since = self.daemon._coverage_since(
                pubkey_hex, relay_url, direction, role,
            )
            flt = _filter_for(pubkey_hex, direction, limit, since)
            caught_up[direction] = False

            def _on_eose(_d: str = direction) -> None:
                self._mark_caught_up(pubkey_hex, relay_url, _d)

            try:
                handles.append(
                    bus.subscribe_sync(relay_url, flt, _on_event, _on_eose)
                )
            except RelayBusError as e:
                # Bus already torn down (shutdown race: auto-tracker has an
                # in-flight discovery whose worker thread reaches `pool.add`
                # after `daemon.stop()` started). Drop any partially-attached
                # handles and bail — there is nothing to subscribe to.
                log.warning(
                    "live sub %s@%s: bus unavailable, skipping subscribe (%s)",
                    pubkey_hex[:6], relay_url, e,
                )
                for h in handles:
                    h.unsubscribe()
                return False

        with self._lock:
            self._subs.setdefault(pubkey_hex, {})[relay_url] = _RelayEntry(
                handles=handles, caught_up=caught_up,
            )

        log.debug(
            "live sub attached: %s@%s role=%s (limit=%d, filters=%d)",
            pubkey_hex[:6], relay_url, role, limit, len(handles),
        )
        return True

    def _mark_caught_up(
        self, pubkey_hex: str, relay_url: str, direction: str,
    ) -> None:
        """Flag a (pubkey, relay, direction) as caught up — the relay sent its
        end-of-stored-events signal. Fires on the bus loop thread; cheap (RAM
        only). The periodic flush turns set flags into persisted watermarks."""
        with self._lock:
            entry = self._subs.get(pubkey_hex, {}).get(relay_url)
            if entry is not None and direction in entry.caught_up:
                entry.caught_up[direction] = True

    def caught_up_triples(self) -> list[tuple[str, str, str]]:
        """Every (pubkey, relay, direction) currently live AND past its
        end-of-stored-events signal — i.e. safe to stamp as covered-through-now.
        """
        with self._lock:
            return [
                (pk, url, direction)
                for pk, per_relay in self._subs.items()
                for url, entry in per_relay.items()
                for direction, ok in entry.caught_up.items()
                if ok
            ]

    def remove_pubkey(self, pubkey_hex: str) -> int:
        """Cancel every handle for a pubkey across all relays. Returns the
        number of handles that were cancelled."""
        with self._lock:
            per_relay = self._subs.pop(pubkey_hex, {})
        total = 0
        for relay_url, entry in per_relay.items():
            for h in entry.handles:
                try:
                    h.cancel(timeout=2.0)
                except Exception:
                    log.exception(
                        "cancel failed for %s@%s", pubkey_hex[:6], relay_url,
                    )
                total += 1
        if total:
            log.info("stopped %d live sub(s) for %s", total, pubkey_hex[:6])
        return total

    # ---- inspection ---------------------------------------------------------

    def has(self, pubkey_hex: str, relay_url: str) -> bool:
        with self._lock:
            return relay_url in self._subs.get(pubkey_hex, {})

    def count_for(self, pubkey_hex: str) -> int:
        with self._lock:
            return len(self._subs.get(pubkey_hex, {}))

    def count_pubkeys(self) -> int:
        with self._lock:
            return sum(1 for r in self._subs.values() if r)

    def list_subs(self) -> list[dict]:
        with self._lock:
            pairs = [
                (pk, url)
                for pk, per_relay in self._subs.items()
                for url in per_relay
            ]
        pairs.sort()
        return [{"pubkey": pk, "relay": url} for pk, url in pairs]

    # ---- teardown -----------------------------------------------------------

    def stop_all(self) -> None:
        with self._lock:
            snapshot = self._subs
            self._subs = {}
        total = 0
        for pk, per_relay in snapshot.items():
            for entry in per_relay.values():
                for h in entry.handles:
                    try:
                        h.cancel(timeout=2.0)
                    except Exception:
                        log.exception("cancel failed for %s", pk[:6])
                    total += 1
        log.info("stopped %d live sub handle(s)", total)

