server-26/drb-c2-core/app/routers/tokens.py

import uuid
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import Optional
from datetime import datetime, timezone
from app.internal import firestore as fstore

router = APIRouter(prefix="/tokens", tags=["tokens"])


class TokenCreate(BaseModel):
    name: str        # friendly label e.g. "DRB Bot 1"
    token: str       # the actual Discord bot token


# ---------------------------------------------------------------------------
# CRUD
# ---------------------------------------------------------------------------

@router.get("")
async def list_tokens():
    """List all tokens. The actual token string is masked for safety."""
    tokens = await fstore.collection_list("bot_tokens")
    return [
        {**t, "token": t["token"][:10] + "…" + t["token"][-4:]}
        for t in tokens
    ]


@router.post("", status_code=201)
async def add_token(body: TokenCreate):
    token_id = str(uuid.uuid4())
    doc = {
        "token_id": token_id,
        "name": body.name,
        "token": body.token,
        "in_use": False,
        "assigned_node_id": None,
        "assigned_at": None,
    }
    await fstore.doc_set("bot_tokens", token_id, doc, merge=False)
    return {"token_id": token_id, "name": body.name}


@router.post("/flush", status_code=200)
async def flush_tokens():
    """Force-release all in-use tokens (admin utility — use when tokens get orphaned)."""
    def _find():
        from app.internal.firestore import db
        return [d for d in db.collection("bot_tokens").where("in_use", "==", True).stream()]

    import asyncio
    results = await asyncio.to_thread(_find)
    for doc in results:
        await fstore.doc_update("bot_tokens", doc.id, {
            "in_use": False,
            "assigned_node_id": None,
            "assigned_at": None,
        })
    return {"released": len(results)}


@router.put("/{token_id}/prefer/{system_id}", status_code=200)
async def set_preferred_system(token_id: str, system_id: str):
    """
    Mark this token as the preferred bot for a system.
    When a discord_join is issued for any node in that system, this token
    is tried first before falling back to the general pool.
    Pass system_id="_none" to clear the preference.
    """
    existing = await fstore.doc_get("bot_tokens", token_id)
    if not existing:
        raise HTTPException(404, "Token not found.")
    if system_id == "_none":
        # Clear any existing preference on the system that pointed to this token.
        system_doc = await fstore.doc_get("systems", existing.get("preferred_for_system_id", ""))
        if system_doc:
            await fstore.doc_set("systems", existing["preferred_for_system_id"], {"preferred_token_id": None})
        await fstore.doc_set("bot_tokens", token_id, {"preferred_for_system_id": None})
        return {"ok": True, "preferred_for_system_id": None}

    system_doc = await fstore.doc_get("systems", system_id)
    if not system_doc:
        raise HTTPException(404, "System not found.")
    # Set preference on both sides for easy lookup in either direction.
    await fstore.doc_set("systems", system_id, {"preferred_token_id": token_id})
    await fstore.doc_set("bot_tokens", token_id, {"preferred_for_system_id": system_id})
    return {"ok": True, "preferred_for_system_id": system_id}


@router.delete("/{token_id}", status_code=204)
async def delete_token(token_id: str):
    existing = await fstore.doc_get("bot_tokens", token_id)
    if not existing:
        raise HTTPException(404, "Token not found.")
    if existing.get("in_use"):
        raise HTTPException(409, "Token is currently in use by a node.")
    await fstore.doc_delete("bot_tokens", token_id)


# ---------------------------------------------------------------------------
# Internal helpers — used by the nodes router, not exposed via HTTP
# ---------------------------------------------------------------------------

async def assign_token(node_id: str, preferred_token_id: Optional[str] = None) -> Optional[str]:
    """
    Find a free token, mark it as in-use, return the token string.
    If preferred_token_id is given, try that token first (only if it's free).
    Returns None if no tokens are available.
    """
    def _find_free(preferred: Optional[str]):
        from app.internal.firestore import db
        if preferred:
            snap = db.collection("bot_tokens").document(preferred).get()
            if snap.exists and not snap.to_dict().get("in_use"):
                return [snap]
        docs = db.collection("bot_tokens").where("in_use", "==", False).limit(1).stream()
        return [d for d in docs]

    import asyncio
    results = await asyncio.to_thread(_find_free, preferred_token_id)
    if not results:
        return None

    doc = results[0]
    token_id = doc.id
    token_value = doc.to_dict()["token"]

    await fstore.doc_update("bot_tokens", token_id, {
        "in_use": True,
        "assigned_node_id": node_id,
        "assigned_at": datetime.now(timezone.utc),
    })
    return token_value


async def release_token(node_id: str) -> None:
    """Free whichever token is currently assigned to this node."""
    def _find_assigned():
        from app.internal.firestore import db
        return [
            d for d in db.collection("bot_tokens")
            .where("assigned_node_id", "==", node_id)
            .stream()
        ]

    import asyncio
    results = await asyncio.to_thread(_find_assigned)
    for doc in results:
        await fstore.doc_update("bot_tokens", doc.id, {
            "in_use": False,
            "assigned_node_id": None,
            "assigned_at": None,
        })