server-26/drb-c2-core/app/internal/discord_streams.py

"""
Discord radio channels as desired state (server-26#203 MVP).

One ``discord_streams/{guild_id}:{channel_id}`` doc per Discord voice channel
that should carry radio. Fields:

  org_id, guild_id, channel_id, channel_name, guild_name,
  system_id, system_name,
  mode: "persistent" | "adhoc"
      persistent — stays connected and reconnects by itself (node reboot,
                   Discord voice drop, C2 restart) until someone stops it.
      adhoc      — the Discord bot sets desired=off after ``idle_minutes``
                   with no humans in the voice channel.
  desired: "on" | "off"
  node_id, token_id, lease_id — placement (written with the lease, atomically)
  actual: {state, error, since, at} — reported by the node, never guessed here
      state: pending | connecting | connected | error | stopped | node_offline
  tg_filter: None — phase-2 seam (per-channel talkgroup gate); always None now
  now_playing_message_id, created_by_discord_id, created_at, updated_at

Flow: ``start_stream`` picks a node + leases a token -> ``publish_desired``
retains ``nodes/{node}/discord/desired`` (full list for that node, with a
monotonically increasing ``seq``) -> the edge reconciles and publishes
``nodes/{node}/discord/actual`` -> ``handle_actual`` updates ``actual`` and
settles token releases. ``reconcile_loop`` re-places persistent streams whose
node went away.

The desired payload carries token **ids**, never token strings: the edge
fetches the string over HTTPS (GET /discord/node/tokens/{id}) with its own
node API key, and C2 hands it out only to the node holding the lease. A
retained MQTT message therefore never stores a bot credential on the broker.
"""
from __future__ import annotations

import asyncio
from datetime import datetime, timezone, timedelta
from typing import Optional

from app.config import settings
from app.internal import firestore as fstore
from app.internal import discord_leases as leases
from app.internal.logger import logger

STREAMS = leases.STREAMS
NODE_STATE = "discord_nodes"   # discord_nodes/{node_id}: {seq, applied_seq, actual_at}
ONLINE = ("online", "recording")


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _effective_system(node: dict) -> Optional[str]:
    if node.get("is_overridden") and node.get("override_system_id"):
        return node["override_system_id"]
    return node.get("assigned_system_id")


async def _node_stream_counts(org_id: Optional[str]) -> dict[str, int]:
    streams = await fstore.collection_list(STREAMS, desired="on")
    counts: dict[str, int] = {}
    for s in streams:
        if s.get("node_id"):
            counts[s["node_id"]] = counts.get(s["node_id"], 0) + 1
    return counts


async def candidate_nodes(org_id: str, system_id: str, exclude: Optional[str] = None) -> list[dict]:
    """Online, approved nodes of this org currently listening to ``system_id``,
    best first: nodes already carrying streams (one gateway session can serve
    several guilds) but under the per-node cap, then the rest."""
    nodes = await fstore.collection_list("nodes", org_id=org_id)
    counts = await _node_stream_counts(org_id)
    cap = settings.discord_node_max_streams
    out = []
    for n in nodes:
        if n.get("node_id") == exclude:
            continue
        if n.get("status") not in ONLINE or n.get("approval_status") not in (None, "approved"):
            continue
        if _effective_system(n) != system_id:
            continue
        if counts.get(n["node_id"], 0) >= cap:
            continue
        out.append(n)
    out.sort(key=lambda n: (-counts.get(n["node_id"], 0), n.get("name") or n["node_id"]))
    return out


# ---------------------------------------------------------------------------
# Desired state -> node
# ---------------------------------------------------------------------------

async def _bump_seq(node_id: str) -> int:
    db = fstore.db

    def _txn(txn):
        ref = db.collection(NODE_STATE).document(node_id)
        snap = ref.get(transaction=txn)
        seq = int((snap.to_dict() or {}).get("seq") or 0) + 1 if snap.exists else 1
        txn.set(ref, {"node_id": node_id, "seq": seq}, merge=True)
        return seq

    return await fstore.run_transaction(_txn)


async def desired_for_node(node_id: str) -> list[dict]:
    streams = await fstore.collection_list(STREAMS, node_id=node_id)
    out = []
    for s in streams:
        if s.get("desired") != "on" or not s.get("token_id"):
            continue
        out.append({
            "stream_id": s["stream_id"],
            "guild_id": s["guild_id"],
            "channel_id": s["channel_id"],
            "token_id": s["token_id"],
            "lease_id": s.get("lease_id"),
            "system_name": s.get("system_name"),
            "mode": s.get("mode", "adhoc"),
            # Phase-2 seam: the edge ignores this today (no TG gate yet).
            "tg_filter": s.get("tg_filter"),
        })
    out.sort(key=lambda d: d["stream_id"])
    return out


async def publish_desired(node_id: str) -> int:
    """Retain the node's full stream list; stamp releasing tokens with the seq."""
    from app.internal.mqtt_handler import mqtt_handler
    seq = await _bump_seq(node_id)
    payload = {"seq": seq, "streams": await desired_for_node(node_id), "issued_at": _now_iso()}
    mqtt_handler.publish_discord_desired(node_id, payload)
    await leases.stamp_release_seq(node_id, seq)
    return seq


# ---------------------------------------------------------------------------
# Start / stop
# ---------------------------------------------------------------------------

class StreamError(Exception):
    def __init__(self, code: str, message: str, invite: Optional[list] = None, status: int = 409):
        super().__init__(message)
        self.code = code
        self.message = message
        self.invite = invite or []
        self.status = status


ERROR_TEXT = {
    "no_node": "No online node in your organization is listening to that system right now.",
    "pool_empty": "No radio bots have been added to DRB yet. A platform admin adds them under Radio bots.",
    "no_bot_in_guild": "No free radio bot is a member of this server. Invite one with the link below, then try again.",
    "pool_exhausted": "Every radio bot that can join this server is already busy. Stop another radio channel or add a bot.",
    "limit": "This server already has its maximum number of radio channels. Stop one with /radio stop first.",
}


async def start_stream(
    *,
    org_id: str,
    guild_id: str,
    channel_id: str,
    system_id: str,
    mode: str = "adhoc",
    channel_name: Optional[str] = None,
    guild_name: Optional[str] = None,
    created_by: Optional[str] = None,
    node_id: Optional[str] = None,
    preferred_token_id: Optional[str] = None,
    max_streams: Optional[int] = None,
) -> dict:
    """Make ``guild_id:channel_id`` carry ``system_id``. Idempotent: re-running
    it for a channel that is already on reuses the node and the lease."""
    if mode not in ("persistent", "adhoc"):
        raise StreamError("bad_mode", "mode must be persistent or adhoc", status=400)
    stream_id = leases.stream_id_for(guild_id, channel_id)
    existing = await fstore.doc_get(STREAMS, stream_id) or {}

    if max_streams is not None and existing.get("desired") != "on":
        live = [s for s in await fstore.collection_list(STREAMS, guild_id=guild_id) if s.get("desired") == "on"]
        if len(live) >= max_streams:
            raise StreamError("limit", ERROR_TEXT["limit"])

    system = await fstore.doc_get_cached("systems", system_id) or {}

    if node_id:
        candidates = [await fstore.doc_get("nodes", node_id) or {"node_id": node_id}]
    else:
        candidates = await candidate_nodes(org_id, system_id)
        # Keep the current node when it is still a valid choice (reuse the lease).
        cur = existing.get("node_id")
        if cur and existing.get("system_id") == system_id:
            candidates.sort(key=lambda n: 0 if n["node_id"] == cur else 1)
    if not candidates:
        raise StreamError("no_node", ERROR_TEXT["no_node"])

    now = _now_iso()
    fields = {
        "stream_id": stream_id, "org_id": org_id, "guild_id": guild_id, "channel_id": channel_id,
        "system_id": system_id, "system_name": system.get("name") or system_id,
        "mode": mode, "desired": "on", "tg_filter": None,
        "actual": {"state": "pending", "error": None, "since": now, "at": now},
        "updated_at": now,
    }
    if channel_name:
        fields["channel_name"] = channel_name
    if guild_name:
        fields["guild_name"] = guild_name
    if not existing:
        fields["created_at"] = now
        fields["created_by_discord_id"] = created_by
        fields["now_playing_message_id"] = None

    last: Optional[leases.LeaseResult] = None
    prev_node = existing.get("node_id")
    for node in candidates:
        res = await leases.lease_for_stream(stream_id, node_id=node["node_id"], org_id=org_id,
                                            preferred_token_id=preferred_token_id, stream_fields=fields)
        if res.ok:
            await publish_desired(node["node_id"])
            if prev_node and prev_node != node["node_id"]:
                await publish_desired(prev_node)
            logger.info(f"Discord stream {stream_id} -> node {node['node_id']} token {res.token_id}"
                        f" ({'reused' if res.reused else 'new'} lease)")
            return {**existing, **fields, "node_id": node["node_id"], "node_name": node.get("name"),
                    "token_id": res.token_id, "lease_id": res.lease_id, "reused": res.reused}
        last = res
    code = last.error if last else "no_node"
    raise StreamError(code, ERROR_TEXT.get(code, code), invite=last.invite if last else [])


async def stop_stream(stream_id: str, reason: str = "stopped") -> bool:
    existing = await fstore.doc_get(STREAMS, stream_id)
    if not existing:
        return False
    now = _now_iso()
    node_id = await leases.detach_stream(stream_id, {
        "desired": "off", "updated_at": now, "stop_reason": reason,
        "actual": {**(existing.get("actual") or {}), "state": "stopping", "at": now},
    })
    if node_id:
        await publish_desired(node_id)
    logger.info(f"Discord stream {stream_id} stopped ({reason})")
    return True


async def stop_guild(guild_id: str, reason: str) -> int:
    n = 0
    for s in await fstore.collection_list(STREAMS, guild_id=guild_id):
        if s.get("desired") == "on":
            await stop_stream(s["stream_id"], reason)
            n += 1
    return n


async def stop_node(node_id: str, reason: str) -> int:
    n = 0
    for s in await fstore.collection_list(STREAMS, node_id=node_id):
        if s.get("desired") == "on":
            await stop_stream(s["stream_id"], reason)
            n += 1
    return n


# ---------------------------------------------------------------------------
# Actual state <- node
# ---------------------------------------------------------------------------

_VALID_STATES = {"connecting", "connected", "error", "stopped", "paused_locally"}


async def handle_actual(node_id: str, payload: dict) -> None:
    """``nodes/{node_id}/discord/actual``. A node may only report on streams
    currently placed on it — anything else is ignored, so a compromised node
    cannot rewrite another node's (or org's) stream state."""
    reported = payload.get("streams") or {}
    now = _now_iso()
    if isinstance(reported, dict):
        for stream_id, st in list(reported.items())[:50]:
            if not isinstance(st, dict):
                continue
            doc = await fstore.doc_get(STREAMS, str(stream_id))
            if not doc or doc.get("node_id") != node_id or doc.get("desired") != "on":
                continue
            state = st.get("state") if st.get("state") in _VALID_STATES else "error"
            error = st.get("error")
            error = str(error)[:300] if error else None
            prev = doc.get("actual") or {}
            since = prev.get("since") if prev.get("state") == state else now
            await fstore.doc_update(STREAMS, stream_id, {
                "actual": {"state": state, "error": error, "since": since, "at": now},
            })

    if "sessions" in payload:
        sessions = {str(t) for t in (payload.get("sessions") or [])}
        applied = payload.get("applied_seq")
        applied = int(applied) if isinstance(applied, (int, float)) else None
        freed = await leases.settle_releases(node_id, sessions, applied)
        for tok in freed:
            logger.info(f"Radio bot token {tok} released — node {node_id} confirmed it logged out")
        await fstore.doc_set(NODE_STATE, node_id, {"applied_seq": applied, "actual_at": now})


# ---------------------------------------------------------------------------
# Self-healing: re-place persistent streams whose node is gone
# ---------------------------------------------------------------------------

async def release_offline_node(node_id: str) -> None:
    """Node offline past the lease grace: free its tokens (the edge already
    logged them out — see release_node_leases) and re-place its streams."""
    orphaned = await leases.release_node_leases(node_id)
    if not orphaned:
        return
    now = _now_iso()
    for sid in orphaned:
        await fstore.doc_set(STREAMS, sid, {
            "token_id": None, "lease_id": None,
            "actual": {"state": "node_offline", "error": "The node carrying this channel went offline.",
                       "since": now, "at": now},
        })
    await publish_desired(node_id)  # retained: the node sees an accurate (empty) list when it returns
    logger.warning(f"Node {node_id} offline past lease grace — released {len(orphaned)} radio stream(s)")


async def reconcile_once() -> None:
    nodes = {n["node_id"]: n for n in await fstore.collection_list("nodes") if n.get("node_id")}
    grace = timedelta(seconds=settings.discord_lease_offline_release_seconds)
    now = datetime.now(timezone.utc)
    for s in await fstore.collection_list(STREAMS, desired="on"):
        node = nodes.get(s.get("node_id") or "")
        online = bool(node) and node.get("status") in ONLINE and _effective_system(node) == s.get("system_id")
        if online and s.get("token_id"):
            continue
        if node and s.get("token_id"):
            # Node offline (or retuned): hold the lease until the grace passes.
            last_seen = node.get("last_seen")
            try:
                ls = datetime.fromisoformat(last_seen) if isinstance(last_seen, str) else last_seen
                if ls and ls.tzinfo is None:
                    ls = ls.replace(tzinfo=timezone.utc)
            except Exception:
                ls = None
            if node.get("status") in ONLINE or (ls and now - ls < grace):
                if node.get("status") not in ONLINE:
                    await _mark(s, "node_offline", "The node carrying this channel is offline; waiting for it.")
                continue
            await release_offline_node(node["node_id"])
        if s.get("mode") != "persistent":
            await stop_stream(s["stream_id"], "node_lost")
            continue
        try:
            await start_stream(org_id=s["org_id"], guild_id=s["guild_id"], channel_id=s["channel_id"],
                               system_id=s["system_id"], mode="persistent")
        except StreamError as e:
            await _mark(s, "error", e.message)


async def _mark(s: dict, state: str, error: str) -> None:
    prev = s.get("actual") or {}
    if prev.get("state") == state and prev.get("error") == error:
        return
    now = _now_iso()
    await fstore.doc_set(STREAMS, s["stream_id"], {"actual": {"state": state, "error": error,
                                                              "since": now, "at": now}})


async def reconcile_loop() -> None:
    logger.info("Discord stream reconciler started.")
    while True:
        await asyncio.sleep(settings.discord_reconcile_interval_seconds)
        try:
            await reconcile_once()
        except Exception as e:
            logger.error(f"Discord reconcile error: {e}")