KiraTV/backend/api/admin/monitoring.py

"""Admin API: real-time stream monitoring via REST and WebSocket."""
import asyncio
import json
import logging

import redis.asyncio as aioredis
from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from ...config import settings
from ...core.pool import pool, MONITORING_KEY
from ...core.vod_tracker import vod_tracker
from ...database import get_db
from ...models.jellyfin import JellyfinConfig
from ...models.provider import ProviderAccount
from ..auth import get_current_admin

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/monitoring", tags=["monitoring"])


def _all_streams() -> list[dict]:
    """Combined live + VOD/series sessions."""
    return pool.get_active_streams() + vod_tracker.get_all()


@router.get("/streams")
async def get_active_streams(_=Depends(get_current_admin)):
    return _all_streams()


@router.get("/slots")
async def get_provider_slots(db: AsyncSession = Depends(get_db), _=Depends(get_current_admin)):
    # Exclude Jellyfin dummy providers — they are local, have no slot limits
    jf_subq = (
        select(JellyfinConfig.provider_account_id)
        .where(JellyfinConfig.provider_account_id.isnot(None))
        .scalar_subquery()
    )
    result = await db.execute(
        select(ProviderAccount).where(
            ProviderAccount.is_active == True,  # noqa: E712
            ProviderAccount.id.notin_(jf_subq),
        )
    )
    accounts = result.scalars().all()

    active_streams = _all_streams()
    busy_accounts = {s["provider_account_id"] for s in active_streams}

    return [
        {
            "id": acc.id,
            "name": acc.name,
            "max_connections": acc.max_connections,
            "status": "busy" if acc.id in busy_accounts else "free",
            "channel_id": next(
                (s["channel_id"] for s in active_streams if s["provider_account_id"] == acc.id), None
            ),
        }
        for acc in accounts
    ]


class KillRequest(BaseModel):
    client_id: str


@router.post("/kill")
async def kill_connection(body: KillRequest, _=Depends(get_current_admin)):
    """Force-disconnect a client. Works for both live streams and VOD/series sessions."""
    cid = body.client_id
    if cid.startswith("vod-"):
        killed = vod_tracker.kill(cid)
    else:
        killed = await pool.kill_client(cid)
    if not killed:
        raise HTTPException(404, f"Client '{cid}' not found in any active stream")
    return {"ok": True, "client_id": cid}


@router.websocket("/ws")
async def monitoring_websocket(websocket: WebSocket):
    await websocket.accept()
    redis_client = aioredis.from_url(settings.redis_url, decode_responses=True)
    pubsub = redis_client.pubsub()

    try:
        await pubsub.subscribe(MONITORING_KEY)

        # Send initial combined state
        await websocket.send_text(json.dumps({"type": "state", "data": _all_streams()}))

        async def redis_reader():
            """Forward live-stream pool events from Redis pub/sub."""
            async for message in pubsub.listen():
                if message["type"] == "message":
                    try:
                        # Pool publishes only live streams; merge with VOD here
                        await websocket.send_text(json.dumps({"type": "update", "data": _all_streams()}))
                    except Exception:
                        return

        async def broadcaster():
            """Push full state every 2 s for VOD bytes/bps updates."""
            while True:
                await asyncio.sleep(2)
                try:
                    await websocket.send_text(json.dumps({"type": "update", "data": _all_streams()}))
                except Exception:
                    return

        reader_task = asyncio.create_task(redis_reader())
        broadcaster_task = asyncio.create_task(broadcaster())

        try:
            await asyncio.gather(reader_task, broadcaster_task)
        except WebSocketDisconnect:
            pass
        finally:
            reader_task.cancel()
            broadcaster_task.cancel()

    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"Monitoring WS error: {e}")
    finally:
        await pubsub.unsubscribe(MONITORING_KEY)
        await redis_client.aclose()