"""Discovery + activation: pick ≤3 maximally-complementary relays for a pubkey,
persist the result, then leave a long-lived live subscription open on each.

Two-phase algorithm:

  PHASE 1 — probe relays in parallel (`PROBE_WORKERS` workers). Relay set
  depends on role:
    - primary:   ALL active relays. Two REQs per relay share a fresh
        per-relay id set:
          `{authors:[hex], limit}`  — events authored by the pubkey
          `{"#p":[hex], limit}`     — events tagging the pubkey
    - secondary: only the author's NIP-65 outbox (write relays). One REQ
        per relay, `{authors:[hex], limit}`. If outbox is unknown, a
        one-shot bootstrap REQ for kind 10002 across all relays runs
        first; only when bootstrap also fails do we fall back to
        full-broadcast.
  Why outbox-first for secondary: with hundreds of secondaries queued by
  p-tag harvesting, fanning out 120 REQs × ~10 s timeout per pubkey
  starves the auto-tracker for hours. Outbox cuts that to 2-4 REQs.
  Why authors-only for secondary: secondaries are 1-hop frontier
  candidates; pulling every mention/reply/reaction worldwide that p-tags
  them dwarfs author traffic for popular pubkeys (~10× volume) for
  marginal value.
  `_backfill_one` ingests every event into the DB along the way (so even
  relays we end up rejecting still contribute their unique events to
  history, and primary discovery seeds secondaries from p-tags as usual).
  Per-relay output: full set of event ids that relay knows about for
  this pubkey.

  PHASE 2 — greedy weighted max-coverage. Iteratively pick the relay that
  adds the most NEW ids relative to already-picked relays, until either
  `PRODUCTIVE_RELAYS_TARGET` are picked or no remaining relay contributes
  ≥`MIN_MARGINAL_GAIN`. This deliberately prefers a relay returning 5
  events that nobody else has over a relay returning 500 events that are
  already fully covered — the live sub keeps that relay's unique data
  flowing forward, which the discarded one wouldn't.

Tie-breaks (equal marginal gain): larger total set first, then
lexicographic url — deterministic across runs.

Storage in `tracked_pubkeys` is the source of truth across restarts. On boot
the daemon replays each row by reopening live subs on its recorded
`productive_relays`. Re-running this tool for the same pubkey:
  - clears the row's existing live subs (rescan from scratch),
  - persists the fresh `productive_relays`,
  - bumps role to 'primary' (promotion if it was 'secondary'),
  - retroactively registers every `p`-tag pubkey from the now-primary's
    historical events as a secondary so the auto-tracker picks them up.
"""

from __future__ import annotations

import logging
from concurrent.futures import ThreadPoolExecutor
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from ingest_sql import IngestDaemonSQL

log = logging.getLogger("track_pubkey")

PRODUCTIVE_RELAYS_TARGET = 3
# A relay must add at least this many event ids vs. already-picked relays
# to be picked. 1 = include even a single unique event (matches the user
# requirement: "fewer results but uncovered → still pick it"). Bumping this
# would intentionally drop low-yield relays at the cost of coverage.
MIN_MARGINAL_GAIN = 1

# Phase 1 fans out one REQ per relay; each REQ has its own
# `BACKFILL_TOTAL_TIMEOUT_S` (10 s in `ingest.py`) so a serial loop over ~80
# active relays approached half an hour per discovery and starved the
# auto-tracker. _backfill_one only touches `conn.ws` of its own relay and
# `_ingest_event` already serializes through the storage lock, so the
# probes are safe to run concurrently. 8 keeps total wall-time bounded by
# the slowest few relays without flooding the storage lock with writers.
PROBE_WORKERS = 8

# Replaceable event kinds carrying profile-shaped state — one row per
# (pubkey, kind) by NIP-16. These are *critical* to ingest for tracked
# pubkeys (UI shows display name + picture from kind 0; relay routing
# follows kind 10002; graph features key off kind 3), but for an active
# user they can easily fall outside the per-relay `limit` window the
# main REQ uses. So we issue a separate, unbounded REQ for these kinds
# in Phase 1 — typically 3 events back, idempotent across relays, and
# applies regardless of role (primary AND secondary need their profile
# snapshot).
PROFILE_KINDS = [0, 3, 10002]


def run(
    daemon: "IngestDaemonSQL",
    pubkey_hex: str,
    limit: int,
    role: str = "primary",
    outbox_bootstrap: bool = True,
) -> dict:
    """Run discovery for `pubkey_hex` and persist with the given role.

    `role` is what the row should carry AFTER this run completes:
      - "primary"   — used by `subscribe_to_pubkey` (explicit user intent)
                      and by promotion. Triggers retroactive p-tag expansion.
      - "secondary" — used by the auto-tracker when it picks up a row
                      registered as secondary by ingest. Must NOT promote
                      to primary, must NOT expand p-tags (1-hop frontier).

    Without the `role` param this used to default-promote every secondary
    the auto-tracker touched, which created an exponential frontier
    explosion: every secondary got promoted, started expanding its own
    p-tags, registered new secondaries, which got promoted, and so on.

    `outbox_bootstrap` controls whether this call may issue a kind-10002
    fan-out when storage has no outbox for the pubkey yet. The
    auto-tracker turns this off because it batches the kind-10002
    lookup across all pubkeys claimed in the tick (one REQ per relay
    instead of N) and calls `run` afterwards. Direct callers leave it
    True so a single-pubkey discovery still self-heals if the outbox
    is missing.
    """
    if role not in ("primary", "secondary"):
        raise ValueError(f"role must be 'primary' or 'secondary', got {role!r}")
    if limit <= 0:
        raise ValueError(f"limit must be > 0, got {limit}")

    pool = daemon.live_subs
    if pool is None:
        raise RuntimeError("daemon has no live_subs pool — wrong backend?")

    log.info("track_pubkey.run %s: role=%s limit=%d", pubkey_hex[:8], role, limit)

    # Tear down any stale live subs for this pubkey before rescan, otherwise
    # add() would collide on (pubkey, relay) duplicates.
    pool.remove_pubkey(pubkey_hex)

    # ---- choose which relays to probe in PHASE 1 ---------------------------
    # Primary discovery probes every active relay — that's how we find the
    # best 3 globally. Secondary discovery is the hot path for the
    # auto-tracker (hundreds of pubkeys queued from p-tag harvesting) and
    # cannot afford to fan out to all ~120 relays per pubkey, so we
    # restrict it to the author's NIP-65 outbox (write relays). If the
    # outbox isn't in storage yet, do a one-shot bootstrap REQ for kind
    # 10002 across all relays, then re-check; only fall back to a full
    # probe when even bootstrap finds nothing (orphan pubkey with no
    # advertised relay list).
    outbox_used = False
    outbox_bootstrapped = False
    if role == "secondary":
        write_urls = daemon._author_write_urls(pubkey_hex)
        if not write_urls and outbox_bootstrap:
            daemon.req_and_ingest(
                {"authors": [pubkey_hex], "kinds": [10002], "limit": 1},
            )
            write_urls = daemon._author_write_urls(pubkey_hex)
            outbox_bootstrapped = True
        active = daemon._read_active_relay_urls()
        if write_urls:
            urls = sorted(active & write_urls)
            outbox_used = True
            log.info(
                "discovery %s: outbox-restricted to %d relay(s)%s",
                pubkey_hex[:6], len(urls),
                " (bootstrapped)" if outbox_bootstrapped else "",
            )
        else:
            urls = sorted(active)
            log.info(
                "discovery %s: no outbox known after bootstrap, "
                "falling back to full probe over %d relay(s)",
                pubkey_hex[:6], len(urls),
            )
    else:
        urls = sorted(daemon._read_active_relay_urls())

    # ---- PHASE 1: probe each chosen relay (in parallel) -------------------
    # Per-relay event-id sets are kept independently so PHASE 2 can run a
    # weighted greedy max-cover over them. Events are ingested into the DB
    # by `_backfill_one` regardless of whether the relay survives the pick.
    def _probe(url: str) -> tuple[dict, set[str] | None]:
        ids: set[str] = set()
        c_out, err_out = daemon._backfill_one(
            url, {"authors": [pubkey_hex], "limit": limit}, ids,
        )
        # Secondary skips the incoming `#p` REQ entirely — see module docstring.
        # If the outgoing REQ failed, also don't bother with the incoming.
        if role == "primary" and err_out is None:
            c_in, err_in = daemon._backfill_one(
                url, {"#p": [pubkey_hex], "limit": limit}, ids,
            )
        else:
            c_in, err_in = 0, None
        # Profile-snapshot REQ: replaceable kinds 0/3/10002 may be older
        # than the per-relay `limit` cursor on `authors`, so a dedicated
        # REQ without `limit`/`since` guarantees we ingest them. Run
        # regardless of role; skip only if the conn already broke.
        prev_err = err_out or err_in
        if prev_err is None:
            c_profile, err_profile = daemon._backfill_one(
                url,
                {"authors": [pubkey_hex], "kinds": PROFILE_KINDS},
                ids,
            )
        else:
            c_profile, err_profile = 0, None
        err = err_out or err_in or err_profile
        row = {
            "url": url,
            "events_received": c_out + c_in + c_profile,
            "events_authored": c_out,
            "events_referencing": c_in,
            "events_profile": c_profile,
            "unique_event_ids": len(ids),
            "error": err,
            "skipped": False,
            "selected": False,
            "marginal_gain": None,
        }
        return row, (ids if err is None and ids else None)

    per_relay_ids: dict[str, set[str]] = {}
    per_relay: list[dict] = []

    if urls:
        workers = min(PROBE_WORKERS, len(urls))
        with ThreadPoolExecutor(
            max_workers=workers, thread_name_prefix="probe",
        ) as exec_pool:
            results = list(exec_pool.map(_probe, urls))
        for url, (row, ids) in zip(urls, results):
            per_relay.append(row)
            if ids is not None:
                per_relay_ids[url] = ids

    daemon._reconcile_after_backfill()

    # ---- PHASE 2: greedy weighted max-coverage -----------------------------
    # Pick relays one at a time; at each step pick the one whose ids set
    # has the largest intersection-complement with `covered`. Tie-breaks:
    # larger total |ids| first, then lex url. Stop on target or low gain.
    by_url = {row["url"]: row for row in per_relay}
    # Snapshot the candidate set BEFORE the greedy loop mutates per_relay_ids
    # (winners get popped). Candidates are the relays that completed PHASE 1
    # with a non-empty id-set — those that errored or were skipped don't
    # count as losers in `_record_discovery_outcome`.
    candidate_urls: set[str] = set(per_relay_ids.keys())
    productive_relays: list[str] = []
    covered: set[str] = set()
    while len(productive_relays) < PRODUCTIVE_RELAYS_TARGET and per_relay_ids:
        best_url = None
        best_gain = -1
        best_total = -1
        for url, ids in per_relay_ids.items():
            gain = len(ids - covered)
            total = len(ids)
            if (
                gain > best_gain
                or (gain == best_gain and total > best_total)
                or (gain == best_gain and total == best_total
                    and (best_url is None or url < best_url))
            ):
                best_gain = gain
                best_total = total
                best_url = url
        if best_url is None or best_gain < MIN_MARGINAL_GAIN:
            break
        productive_relays.append(best_url)
        covered |= per_relay_ids[best_url]
        row = by_url[best_url]
        row["selected"] = True
        row["marginal_gain"] = best_gain
        log.info(
            "discovery %s: pick #%d %s (+%d new, total=%d, covered=%d)",
            pubkey_hex[:6], len(productive_relays), best_url,
            best_gain, best_total, len(covered),
        )
        del per_relay_ids[best_url]

    daemon._persist_tracked_pubkey(
        pubkey_hex, productive_relays, limit, role=role,
    )

    winners_set = set(productive_relays)
    losers_set = candidate_urls - winners_set
    daemon._record_discovery_outcome(winners_set, losers_set)

    # Retroactive frontier expansion is a PRIMARY-only behavior. Secondaries
    # do not expand their own p-tags (the 1-hop frontier rule), so when the
    # auto-tracker runs discovery on a secondary we must skip this step —
    # otherwise every secondary's existing events would seed yet more
    # secondaries indirectly.
    #
    # `_p_tag_pubkeys_for_author` takes `_sql_lock` itself; calling it from
    # under the lock here would deadlock (the lock is a plain `threading.Lock`,
    # not reentrant). Read the targets first, then re-acquire the lock for the
    # batch of `_register_secondary_unlocked` writes.
    if role == "primary" and hasattr(daemon, "_p_tag_pubkeys_for_author"):
        targets = daemon._p_tag_pubkeys_for_author(pubkey_hex)
        if targets:
            with daemon._sql_lock:
                for target in targets:
                    daemon._register_secondary_unlocked(target)
            log.info(
                "discovery %s: %d secondary candidate(s) from p-tags",
                pubkey_hex[:6], len(targets),
            )

    for url in productive_relays:
        pool.add(pubkey_hex, url, limit, role=role)

    if len(productive_relays) >= PRODUCTIVE_RELAYS_TARGET:
        stop_reason = "target_reached"
    elif per_relay_ids:
        # Remaining candidates exist but none would add ≥MIN_MARGINAL_GAIN.
        stop_reason = "no_marginal_gain"
    else:
        stop_reason = "exhausted_relays"

    return {
        "pubkey": pubkey_hex,
        "limit_per_relay": limit,
        "productive_relays": productive_relays,
        "live_subs_active": len(productive_relays),
        "relays_scanned": len(per_relay),
        "relays_available": len(urls),
        "unique_events_covered": len(covered),
        "stop_reason": stop_reason,
        "outbox_used": outbox_used,
        "outbox_bootstrapped": outbox_bootstrapped,
        "per_relay": per_relay,
    }


def replay_on_startup(daemon: "IngestDaemonSQL") -> dict:
    """Read tracked_pubkeys from storage and start a live sub on every
    (pubkey, productive_relay) pair. Called once after daemon boot.

    No discovery is re-run; productive relays are taken at face value from
    the persisted record. Re-discovery happens only when the user re-invokes
    `subscribe_to_pubkey` for that pubkey — or when the auto-tracker picks
    up a row with `last_scan_at IS NULL` (never-scanned secondary or freshly
    promoted primary).

    Both primary and secondary rows replay identically — the only role-
    dependent behavior is whether `_ingest_event` expands their `p`-tags.
    """
    rows = daemon._read_tracked_pubkeys()
    pool = daemon.live_subs
    started = 0
    pubkeys_with_relays = 0
    pubkeys_empty = 0
    pubkeys_ignored = 0

    for row in rows:
        # 'ignored' is a tombstone — no live-sub on startup, no discovery.
        # The row stays so frontier expansion's existence check keeps blocking
        # resurrection.
        if (row.get("role") or "primary") == "ignored":
            pubkeys_ignored += 1
            continue
        relays = row.get("productive_relays") or []
        if not relays:
            pubkeys_empty += 1
            continue
        pubkeys_with_relays += 1
        limit = int(row.get("limit_per_relay") or 100)
        role = row.get("role") or "primary"
        for url in relays:
            try:
                pool.add(row["hex"], url, limit, role=role)
                started += 1
            except RuntimeError as e:
                log.warning("replay skip %s@%s: %s", row["hex"][:6], url, e)

    log.info(
        "replay: %d tracked pubkeys (%d with relays, %d empty, %d ignored), "
        "started %d live sub(s)",
        len(rows), pubkeys_with_relays, pubkeys_empty, pubkeys_ignored, started,
    )
    return {
        "tracked_pubkeys_total": len(rows),
        "tracked_pubkeys_with_productive_relays": pubkeys_with_relays,
        "tracked_pubkeys_without_productive_relays": pubkeys_empty,
        "tracked_pubkeys_ignored": pubkeys_ignored,
        "live_subs_started": started,
    }
