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

"""
Radio-bot token leases (server-26#204, node-26#25).

A pool "radio bot" is a Discord bot token that an edge node logs in with to
join voice channels. The rules this module enforces, all inside one Firestore
transaction per change so two concurrent requests can never both win:

* A token is leased to **at most one node** at a time (``lease.node_id``).
  Discord allows one gateway session per token to be useful; two nodes logged
  in on one token fight over the same voice connections.
* One leased token can serve several streams on its node, but **one per
  guild** — Discord allows one voice connection per bot user per guild.
* Re-requesting a stream that already holds a lease on the same node reuses
  it (``/radio play`` twice, a reconcile retry, a re-join) instead of taking a
  second token.
* Release is **ordered**: stopping a stream only marks the token
  ``releasing``. The token returns to the pool once the node reports (MQTT
  ``nodes/{id}/discord/actual``) that it applied a desired-state version at
  least as new as the one that dropped the token *and* has no gateway session
  on it any more (``settle_releases``) — or, as a last resort, once the node
  has been offline longer than the edge's own MQTT-loss fail-safe
  (``release_node_leases``), at which point the edge has already logged out.

Firestore layout:

  bot_tokens/{token_id}
    token, name, org_id (None = platform pool, else BYO for that org),
    guild_ids (list | None — None means membership never checked),
    application_id, bot_user_id, valid, validated_at, check_error,
    lease: None | {node_id, lease_id, since, streams: [stream_id],
                   releasing: bool, release_seq: int | None}
    in_use / assigned_node_id / assigned_at — legacy mirrors of ``lease``,
    still written so the old /tokens page and any reader of them stay right.

  discord_streams/{guild_id}:{channel_id} — see discord_streams.py.
"""
from __future__ import annotations

import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional

from app.internal import firestore as fstore

TOKENS = "bot_tokens"
STREAMS = "discord_streams"

# Connect (1<<20) + Speak (1<<21) + View Channel (1<<10). Nothing else — a
# radio bot never reads or sends messages, so it never asks to.
RADIO_BOT_PERMISSIONS = (1 << 10) | (1 << 20) | (1 << 21)


def invite_url(application_id: Optional[str], guild_id: Optional[str] = None) -> Optional[str]:
    if not application_id:
        return None
    url = (
        "https://discord.com/oauth2/authorize"
        f"?client_id={application_id}&scope=bot&permissions={RADIO_BOT_PERMISSIONS}"
    )
    if guild_id:
        url += f"&guild_id={guild_id}&disable_guild_select=true"
    return url


def stream_id_for(guild_id: str, channel_id: str) -> str:
    return f"{guild_id}:{channel_id}"


def guild_of(stream_id: str) -> str:
    return stream_id.split(":", 1)[0]


def _now() -> datetime:
    return datetime.now(timezone.utc)


@dataclass
class LeaseResult:
    ok: bool
    token_id: Optional[str] = None
    lease_id: Optional[str] = None
    reused: bool = False
    # error codes: pool_empty | no_bot_in_guild | pool_exhausted
    error: Optional[str] = None
    # bots that would work if invited to this guild: [{name, invite_url}]
    invite: list = field(default_factory=list)


def _lease_streams(tok: dict) -> list:
    return list(((tok.get("lease") or {}).get("streams")) or [])


def _legacy_mirror(lease: Optional[dict]) -> dict:
    if not lease:
        return {"lease": None, "in_use": False, "assigned_node_id": None, "assigned_at": None}
    return {
        "lease": lease,
        "in_use": True,
        "assigned_node_id": lease["node_id"],
        "assigned_at": lease.get("since"),
    }


def _candidate_order(tok_id: str, tok: dict, node_id: str, org_id: Optional[str],
                     preferred: Optional[str]) -> tuple:
    lease = tok.get("lease") or {}
    return (
        0 if tok_id == preferred else 1,
        0 if org_id and tok.get("org_id") == org_id else 1,     # BYO before pool
        0 if lease.get("node_id") == node_id else 1,            # reuse a session
        0 if tok.get("guild_ids") is not None else 1,           # known member first
        str(tok.get("name") or ""),
        tok_id,
    )


def pick_token(
    tokens: dict[str, dict],
    *,
    node_id: str,
    guild_id: str,
    org_id: Optional[str],
    preferred_token_id: Optional[str] = None,
    exclude_stream_id: Optional[str] = None,
) -> LeaseResult:
    """Pure choice over a snapshot of bot_tokens. No I/O — unit tested directly."""
    if not tokens:
        return LeaseResult(ok=False, error="pool_empty")
    usable: list[tuple] = []
    not_member: list[dict] = []
    for tok_id, tok in tokens.items():
        if tok.get("valid") is False:
            continue
        owner = tok.get("org_id")
        if owner is not None and owner != org_id:
            continue  # another org's branded bot — never leased across orgs
        lease = tok.get("lease") or {}
        if lease and lease.get("node_id") and lease.get("node_id") != node_id:
            continue  # live (or still releasing) on another node
        busy_guilds = {guild_of(s) for s in _lease_streams(tok) if s != exclude_stream_id}
        if guild_id in busy_guilds:
            continue  # this bot already has a voice connection in this guild
        members = tok.get("guild_ids")
        if members is not None and guild_id not in members:
            not_member.append({"name": tok.get("name") or tok_id,
                               "invite_url": invite_url(tok.get("application_id"), guild_id)})
            continue
        usable.append((_candidate_order(tok_id, tok, node_id, org_id, preferred_token_id), tok_id))
    if usable:
        usable.sort()
        return LeaseResult(ok=True, token_id=usable[0][1])
    if not_member:
        return LeaseResult(ok=False, error="no_bot_in_guild", invite=not_member)
    return LeaseResult(ok=False, error="pool_exhausted")


async def lease_for_stream(
    stream_id: str,
    *,
    node_id: str,
    org_id: Optional[str],
    preferred_token_id: Optional[str] = None,
    stream_fields: Optional[dict] = None,
) -> LeaseResult:
    """Bind ``stream_id`` to a token leased to ``node_id``, atomically.

    Also writes ``stream_fields`` (desired=on, node_id, token_id, lease_id...)
    onto the stream doc in the same transaction, so a stream doc can never name
    a token whose lease does not list it.
    """
    guild_id = guild_of(stream_id)
    db = fstore.db

    def _txn(txn):
        stream_ref = db.collection(STREAMS).document(stream_id)
        stream_snap = stream_ref.get(transaction=txn)
        stream = stream_snap.to_dict() if stream_snap.exists else {}
        tokens = {s.id: s.to_dict() for s in db.collection(TOKENS).stream(transaction=txn)}

        prev_tok_id = stream.get("token_id")
        prev_tok = tokens.get(prev_tok_id) if prev_tok_id else None
        if (
            prev_tok
            and stream.get("node_id") == node_id
            and (prev_tok.get("lease") or {}).get("node_id") == node_id
            and stream_id in _lease_streams(prev_tok)
            and prev_tok.get("valid") is not False
        ):
            lease = prev_tok["lease"]
            if lease.get("releasing"):
                lease = {**lease, "releasing": False, "release_seq": None}
                txn.update(db.collection(TOKENS).document(prev_tok_id), _legacy_mirror(lease))
            txn.set(stream_ref, {**(stream_fields or {}), "node_id": node_id,
                                 "token_id": prev_tok_id, "lease_id": lease["lease_id"]}, merge=True)
            return LeaseResult(ok=True, token_id=prev_tok_id, lease_id=lease["lease_id"], reused=True)

        result = pick_token(tokens, node_id=node_id, guild_id=guild_id, org_id=org_id,
                            preferred_token_id=preferred_token_id, exclude_stream_id=stream_id)
        if not result.ok:
            return result

        # Detach from the previous token (moved node, or token went invalid).
        if prev_tok and stream_id in _lease_streams(prev_tok) and prev_tok_id != result.token_id:
            txn.update(db.collection(TOKENS).document(prev_tok_id),
                       _legacy_mirror(_without_stream(prev_tok["lease"], stream_id)))

        tok = tokens[result.token_id]
        lease = tok.get("lease") or {}
        if lease.get("node_id") == node_id:
            streams = [s for s in _lease_streams(tok) if s != stream_id] + [stream_id]
            lease = {**lease, "streams": streams, "releasing": False, "release_seq": None}
        else:
            lease = {"node_id": node_id, "lease_id": uuid.uuid4().hex, "since": _now().isoformat(),
                     "streams": [stream_id], "releasing": False, "release_seq": None}
        txn.update(db.collection(TOKENS).document(result.token_id), _legacy_mirror(lease))
        txn.set(stream_ref, {**(stream_fields or {}), "node_id": node_id,
                             "token_id": result.token_id, "lease_id": lease["lease_id"]}, merge=True)
        result.lease_id = lease["lease_id"]
        return result

    return await fstore.run_transaction(_txn)


def _without_stream(lease: dict, stream_id: str) -> dict:
    streams = [s for s in (lease.get("streams") or []) if s != stream_id]
    # The token stays leased to the node — it may still be logged in there —
    # until settle_releases sees the node ack. ``releasing`` only once empty.
    return {**lease, "streams": streams, "releasing": not streams,
            "release_seq": lease.get("release_seq") if streams else None}


async def detach_stream(stream_id: str, stream_updates: Optional[dict] = None) -> Optional[str]:
    """Drop ``stream_id`` from its token's lease (token -> ``releasing`` if that
    was its last stream) and apply ``stream_updates`` to the stream doc, in one
    transaction. Returns the node_id the token is (still) leased to, if any."""
    db = fstore.db

    def _txn(txn):
        stream_ref = db.collection(STREAMS).document(stream_id)
        snap = stream_ref.get(transaction=txn)
        if not snap.exists:
            return None
        stream = snap.to_dict()
        tok_id = stream.get("token_id")
        tok_node = None
        if tok_id:
            tok_ref = db.collection(TOKENS).document(tok_id)
            tsnap = tok_ref.get(transaction=txn)
            tok = tsnap.to_dict() if tsnap.exists else None
            if tok and tok.get("lease") and stream_id in _lease_streams(tok):
                txn.update(tok_ref, _legacy_mirror(_without_stream(tok["lease"], stream_id)))
                tok_node = tok["lease"].get("node_id")
        txn.set(stream_ref, {**(stream_updates or {}), "token_id": None, "lease_id": None}, merge=True)
        return tok_node or stream.get("node_id")

    return await fstore.run_transaction(_txn)


async def delete_stopped_stream(stream_id: str) -> bool:
    """Delete a stream doc, but only while it is stopped and holds no token.

    True when the doc is gone (deleted now, or already absent). False when it
    is desired-on again or names a token — a concurrent start re-leased it,
    and deleting it would leave that token's lease listing a stream nobody can
    stop. ``lease_for_stream`` writes the stream doc in its own transaction,
    so reading it here is enough to serialise against it."""
    db = fstore.db

    def _txn(txn):
        ref = db.collection(STREAMS).document(stream_id)
        snap = ref.get(transaction=txn)
        if not snap.exists:
            return True
        stream = snap.to_dict() or {}
        if stream.get("desired") == "on" or stream.get("token_id"):
            return False
        txn.delete(ref)
        return True

    return await fstore.run_transaction(_txn)


async def stamp_release_seq(node_id: str, seq: int) -> None:
    """Record the desired-state version that dropped each releasing token, so a
    stale actual report (from before the node saw the drop) can never free it."""
    db = fstore.db

    def _txn(txn):
        tokens = [(s.id, s.to_dict()) for s in db.collection(TOKENS).stream(transaction=txn)]
        for tok_id, tok in tokens:
            lease = tok.get("lease") or {}
            if lease.get("node_id") == node_id and lease.get("releasing") and lease.get("release_seq") is None:
                txn.update(db.collection(TOKENS).document(tok_id),
                           _legacy_mirror({**lease, "release_seq": seq}))

    await fstore.run_transaction(_txn)


async def settle_releases(node_id: str, sessions: set, applied_seq: Optional[int]) -> list[str]:
    """Free tokens the node has confirmed it no longer runs. Returns freed ids."""
    db = fstore.db

    def _txn(txn):
        freed = []
        tokens = [(s.id, s.to_dict()) for s in db.collection(TOKENS).stream(transaction=txn)]
        for tok_id, tok in tokens:
            lease = tok.get("lease") or {}
            if lease.get("node_id") != node_id or not lease.get("releasing") or lease.get("streams"):
                continue
            rseq = lease.get("release_seq")
            if rseq is None or applied_seq is None or applied_seq < rseq:
                continue  # the node may not have seen the drop yet
            if tok_id in sessions:
                continue  # still logged in on the node
            txn.update(db.collection(TOKENS).document(tok_id), _legacy_mirror(None))
            freed.append(tok_id)
        return freed

    return await fstore.run_transaction(_txn)


async def release_node_leases(node_id: str) -> list[str]:
    """Last-resort release for a node that has been offline longer than the
    edge's MQTT-loss fail-safe (the edge logs every radio bot out by itself
    after DISCORD_MQTT_LOSS_STOP_SECONDS without the broker). Returns the
    stream ids that lost their token so the caller can re-place them."""
    db = fstore.db

    def _txn(txn):
        orphaned = []
        tokens = [(s.id, s.to_dict()) for s in db.collection(TOKENS).stream(transaction=txn)]
        for tok_id, tok in tokens:
            lease = tok.get("lease") or {}
            if lease.get("node_id") != node_id:
                continue
            orphaned.extend(lease.get("streams") or [])
            txn.update(db.collection(TOKENS).document(tok_id), _legacy_mirror(None))
        return orphaned

    return await fstore.run_transaction(_txn)


async def force_release_token(token_id: str) -> Optional[dict]:
    """Admin override (/tokens/flush, delete). Frees the token and detaches
    every stream that named it, in one transaction, so no stream doc is left
    pointing at a token whose lease no longer lists it. Returns the old lease
    (``None`` if the token was not leased or does not exist); the caller
    republishes desired state for ``lease["node_id"]`` so the edge logs out."""
    db = fstore.db

    def _txn(txn):
        ref = db.collection(TOKENS).document(token_id)
        snap = ref.get(transaction=txn)
        if not snap.exists:
            return None
        lease = snap.to_dict().get("lease")
        stream_refs = []
        for sid in (lease or {}).get("streams") or []:
            sref = db.collection(STREAMS).document(sid)
            ssnap = sref.get(transaction=txn)
            if ssnap.exists and (ssnap.to_dict() or {}).get("token_id") == token_id:
                stream_refs.append(sref)
        txn.update(ref, _legacy_mirror(None))
        for sref in stream_refs:
            txn.set(sref, {"token_id": None, "lease_id": None}, merge=True)
        return lease

    return await fstore.run_transaction(_txn)


# ---------------------------------------------------------------------------
# Token validation / guild membership (server-26#204 item 5)
# ---------------------------------------------------------------------------

DISCORD_API = "https://discord.com/api/v10"


class TokenRejected(Exception):
    """Discord answered 401: the string is not a working bot token."""


async def inspect_token(token: str) -> dict:
    """Ask Discord who this token is and which servers it is in.

    Returns ``{bot_user_id, bot_name, application_id, guild_ids, valid,
    validated_at, check_error}``. Raises TokenRejected on a 401. Any other
    failure (Discord down, rate limit) is reported in ``check_error`` with
    ``valid``/``guild_ids`` = None — unknown, not bad: a token is never
    rejected because Discord had a bad minute.
    """
    import httpx
    headers = {"Authorization": f"Bot {token}"}
    out: dict = {"valid": None, "guild_ids": None, "check_error": None,
                 "validated_at": _now().isoformat()}
    try:
        async with httpx.AsyncClient(timeout=10, base_url=DISCORD_API, headers=headers) as client:
            me = await client.get("/users/@me")
            if me.status_code == 401:
                raise TokenRejected("Discord rejected this bot token (401).")
            me.raise_for_status()
            user = me.json()
            out["bot_user_id"] = str(user.get("id"))
            out["bot_name"] = user.get("username")
            app = await client.get("/oauth2/applications/@me")
            if app.status_code == 200:
                out["application_id"] = str(app.json().get("id"))
            guilds = await client.get("/users/@me/guilds", params={"limit": 200})
            guilds.raise_for_status()
            out["guild_ids"] = [str(g["id"]) for g in guilds.json()]
            out["valid"] = True
    except TokenRejected:
        raise
    except Exception as e:  # network / 5xx / 429 — unknown, not invalid
        out["check_error"] = str(e)[:200]
    return out


async def refresh_membership(token_ids: list[str]) -> int:
    """Re-read guild membership for these tokens. Returns how many changed."""
    changed = 0
    for tok_id in token_ids:
        tok = await fstore.doc_get(TOKENS, tok_id)
        if not tok or not tok.get("token"):
            continue
        try:
            info = await inspect_token(tok["token"])
        except TokenRejected as e:
            await fstore.doc_update(TOKENS, tok_id, {"valid": False, "check_error": str(e),
                                                     "validated_at": _now().isoformat()})
            changed += 1
            continue
        updates = {k: v for k, v in info.items() if v is not None or k == "check_error"}
        if updates.get("guild_ids") != tok.get("guild_ids") or updates.get("valid") != tok.get("valid"):
            changed += 1
        await fstore.doc_update(TOKENS, tok_id, updates)
    return changed


async def leave_guild(guild_id: str, org_id: Optional[str]) -> list[dict]:
    """Make the radio bots DRB holds tokens for leave one Discord server
    (Discord REST ``DELETE /users/@me/guilds/{id}``, the bot's own token).

    Only platform-pool bots and ``org_id``'s own BYO bots — never another
    org's branded bot. Bots whose cached membership says they are not in the
    server are skipped; unknown membership is tried (404 = not a member).
    The DRB *server* bot (the discord-bot service) is not in this pool, so it
    stays in the server until someone kicks it. Returns one row per bot tried:
    ``{token_id, name, left, error}``. Never raises for a Discord failure."""
    import httpx
    guild_id = str(guild_id)
    out: list[dict] = []
    for tok in await fstore.collection_list(TOKENS):
        tok_id = tok.get("token_id")
        if not tok_id or not tok.get("token"):
            continue
        if tok.get("org_id") is not None and tok.get("org_id") != org_id:
            continue
        members = tok.get("guild_ids")
        if members is not None and guild_id not in members:
            continue
        row = {"token_id": tok_id, "name": tok.get("name") or tok_id, "left": False, "error": None}
        try:
            async with httpx.AsyncClient(timeout=10, base_url=DISCORD_API,
                                         headers={"Authorization": f"Bot {tok['token']}"}) as client:
                r = await client.delete(f"/users/@me/guilds/{guild_id}")
            if r.status_code in (200, 204, 404):     # 404: Unknown Guild — not a member
                row["left"] = True
            else:
                row["error"] = f"Discord answered {r.status_code}"
        except Exception as e:  # network — report, do not fail the removal
            row["error"] = str(e)[:200]
        if row["left"] and members is not None:
            await fstore.doc_update(TOKENS, tok_id, {"guild_ids": [g for g in members if g != guild_id]})
        out.append(row)
    return out