KiraTV/backend/core/vod_tracker.py

"""VodTracker: lightweight session tracking for VOD and Series streams."""
import asyncio
import logging
import time
import uuid
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

STALE_TIMEOUT = 10  # seconds of no new bytes before a session is considered dead


@dataclass
class VodSession:
    session_id: str
    user_id: str
    username: str
    channel_id: int
    channel_name: str
    provider_account_id: int
    provider_name: str
    stream_type: str  # "movie" or "series"
    stream_url: str
    started_at: float = field(default_factory=time.time)
    last_active: float = field(default_factory=time.time)
    bytes_pumped: int = 0
    bps_down: float = 0.0  # download speed from provider
    bps_up: float = 0.0    # upload speed to client (same as down for VOD, 1 client)
    running: bool = True
    _stop: asyncio.Event = field(default_factory=asyncio.Event)
    # Set by _proxy_vod so kill() can immediately abort the upstream connection
    _http_session: object = field(default=None, repr=False)

    @property
    def should_stop(self) -> bool:
        return self._stop.is_set()

    def kill(self) -> None:
        self._stop.set()
        # Close the upstream aiohttp session immediately so iter_chunked() raises
        # and the generator's finally block runs without waiting for the next chunk.
        http = self._http_session
        if http is not None:
            try:
                http.connector.close()
            except Exception:
                pass

    def to_dict(self) -> dict:
        return {
            "stream_type": self.stream_type,
            "channel_id": self.channel_id,
            "channel_name": self.channel_name,
            "provider_account_id": self.provider_account_id,
            "provider_name": self.provider_name,
            "client_count": 1,
            "clients": [{"client_id": self.session_id, "username": self.username}],
            "started_at": self.started_at,
            "bytes_pumped": self.bytes_pumped,
            "bps_down": self.bps_down,
            "bps_up": self.bps_up,
            "running": self.running,
        }


class VodTracker:
    def __init__(self) -> None:
        self._sessions: dict[str, VodSession] = {}
        self._user_channel_map: dict[tuple[str, int], str] = {}
        self._cleanup_task: asyncio.Task | None = None

    def _ensure_cleanup_running(self) -> None:
        """Start the background staleness sweep if not already running."""
        if self._cleanup_task is None or self._cleanup_task.done():
            try:
                loop = asyncio.get_running_loop()
                self._cleanup_task = loop.create_task(self._cleanup_loop(), name="vod-cleanup")
            except RuntimeError:
                pass  # no event loop yet — task will start on first session

    async def _cleanup_loop(self) -> None:
        """Every 5 s, remove sessions that have received no new bytes for STALE_TIMEOUT s."""
        while True:
            await asyncio.sleep(5)
            now = time.time()
            stale = [
                s for s in list(self._sessions.values())
                if now - s.last_active > STALE_TIMEOUT
            ]
            for s in stale:
                logger.info(
                    f"Removing stale VOD session {s.session_id} "
                    f"(user={s.user_id}, ch={s.channel_id}, "
                    f"idle={now - s.last_active:.0f}s)"
                )
                s.kill()
                self._sessions.pop(s.session_id, None)
                key = (s.user_id, s.channel_id)
                if self._user_channel_map.get(key) == s.session_id:
                    self._user_channel_map.pop(key, None)

    def start(
        self,
        user_id: str,
        username: str,
        channel_id: int,
        channel_name: str,
        provider_account_id: int,
        provider_name: str,
        stream_type: str,
        stream_url: str,
        user_max_connections: int = 1,
    ) -> VodSession:
        self._ensure_cleanup_running()

        # Kill ALL existing VOD sessions for this user — one active VOD per user.
        # Handles both probe-zombies (same channel) and movie switching (different channel).
        for s in list(self._sessions.values()):
            if s.user_id == user_id:
                self._sessions.pop(s.session_id, None)
                old_key = (s.user_id, s.channel_id)
                if self._user_channel_map.get(old_key) == s.session_id:
                    self._user_channel_map.pop(old_key, None)
                s.kill()
                logger.info(
                    f"Killed VOD session {s.session_id} (ch={s.channel_id}) "
                    f"for user={user_id} — new session starting"
                )

        session_id = f"vod-{uuid.uuid4().hex[:8]}"
        session = VodSession(
            session_id=session_id,
            user_id=user_id,
            username=username,
            channel_id=channel_id,
            channel_name=channel_name,
            provider_account_id=provider_account_id,
            provider_name=provider_name,
            stream_type=stream_type,
            stream_url=stream_url,
        )
        self._sessions[session_id] = session
        self._user_channel_map[(user_id, channel_id)] = session_id
        return session

    def end(self, session_id: str) -> None:
        session = self._sessions.pop(session_id, None)
        if session:
            key = (session.user_id, session.channel_id)
            if self._user_channel_map.get(key) == session_id:
                self._user_channel_map.pop(key, None)

    def kill(self, session_id: str) -> bool:
        s = self._sessions.get(session_id)
        if s:
            s.kill()
            return True
        return False

    def get_all(self) -> list[dict]:
        return [s.to_dict() for s in self._sessions.values()]


vod_tracker = VodTracker()