Blame for mediaflow_proxy/utils/hls_prebuffer.py

"""
HLS Pre-buffer system with priority-based sequential prefetching.

This module provides a smart prebuffering system that:
- Prioritizes player-requested segments (downloaded immediately)
- Prefetches remaining segments sequentially in background
- Supports multiple users watching the same channel (shared prefetcher)
- Cleans up inactive prefetchers automatically

Architecture:
1. When playlist is fetched, register_playlist() creates a PlaylistPrefetcher
2. PlaylistPrefetcher runs a background loop: priority queue -> sequential prefetch
3. When player requests a segment, request_segment() adds it to priority queue
4. Prefetcher downloads priority segment first, then continues sequential
"""

import asyncio
import logging
import time
from typing import Dict, Optional, List
from urllib.parse import urljoin

from mediaflow_proxy.utils.base_prebuffer import BasePrebuffer
from mediaflow_proxy.utils.cache_utils import get_cached_segment
from mediaflow_proxy.configs import settings

logger = logging.getLogger(__name__)


class PlaylistPrefetcher:
    """
    Manages prefetching for a single playlist with priority support.

    Key design for live streams with changing tokens:
    - Does NOT start prefetching immediately on registration
    - Only starts prefetching AFTER player requests a segment
    - This ensures we prefetch from the CURRENT playlist, not stale ones

    The prefetcher runs a background loop that:
    1. Waits for player to request a segment (priority)
    2. Downloads the priority segment first
    3. Then prefetches subsequent segments sequentially
    4. Stops when cancelled or all segments are prefetched
    """

    def __init__(
        self,
        playlist_url: str,
        segment_urls: List[str],
        headers: Dict[str, str],
        prebuffer: "HLSPreBuffer",
        prefetch_limit: int = 5,
    ):
        """
        Initialize a playlist prefetcher.

        Args:
            playlist_url: URL of the HLS playlist
            segment_urls: Ordered list of segment URLs from the playlist
            headers: Headers to use for requests
            prebuffer: Parent HLSPreBuffer instance for download methods
            prefetch_limit: Maximum number of segments to prefetch ahead of player position
        """
        self.playlist_url = playlist_url
        self.segment_urls = segment_urls
        self.headers = headers
        self.prebuffer = prebuffer
        self.prefetch_limit = prefetch_limit

        self.last_access = time.time()
        self.current_index = 0  # Next segment to prefetch sequentially
        self.player_index = 0  # Last segment index requested by player
        self.priority_event = asyncio.Event()  # Signals priority segment available
        self.priority_url: Optional[str] = None  # Current priority segment
        self.cancelled = False
        self._task: Optional[asyncio.Task] = None
        self._lock = asyncio.Lock()  # Protects priority_url

        # Track which segments are already cached or being downloaded
        self.downloading: set = set()

        # Track if prefetching has been activated by a player request
        self.activated = False

    def start(self) -> None:
        """Start the prefetch background task."""
        if self._task is None or self._task.done():
            self._task = asyncio.create_task(self._run())
            logger.info(f"[PlaylistPrefetcher] Started (waiting for activation): {self.playlist_url}")

    def stop(self) -> None:
        """Stop the prefetch background task."""
        self.cancelled = True
        self.priority_event.set()  # Wake up the loop
        if self._task and not self._task.done():
            self._task.cancel()
        logger.info(f"[PlaylistPrefetcher] Stopped for: {self.playlist_url}")

    def update_segments(self, segment_urls: List[str]) -> None:
        """
        Update segment URLs (called when playlist is refreshed).

        Args:
            segment_urls: New list of segment URLs
        """
        self.segment_urls = segment_urls
        self.last_access = time.time()
        logger.debug(f"[PlaylistPrefetcher] Updated segments ({len(segment_urls)}): {self.playlist_url}")

    async def request_priority(self, segment_url: str) -> None:
        """
        Player requested this segment - update indices and activate prefetching.

        The player will download this segment via get_or_download().
        The prefetcher's job is to prefetch segments AHEAD of the player,
        not to download the segment the player is already requesting.

        For VOD/movie streams: handles seek by detecting large jumps in segment
        index and resetting the prefetch window accordingly.

        Args:
            segment_url: URL of the segment the player needs
        """
        self.last_access = time.time()
        self.activated = True  # Activate prefetching

        # Update player position for prefetch limit calculation
        segment_index = self._find_segment_index(segment_url)
        if segment_index >= 0:
            old_player_index = self.player_index
            self.player_index = segment_index
            # Start prefetching from the NEXT segment (player handles current one)
            self.current_index = segment_index + 1

            # Detect seek: if player jumped more than prefetch_limit segments
            # This handles VOD seek scenarios where user jumps to different position
            jump_distance = abs(segment_index - old_player_index)
            if jump_distance > self.prefetch_limit and old_player_index >= 0:
                logger.info(
                    f"[PlaylistPrefetcher] Seek detected: jumped {jump_distance} segments "
                    f"(from {old_player_index} to {segment_index})"
                )

        # Signal the prefetch loop to wake up and start prefetching ahead
        async with self._lock:
            self.priority_url = segment_url
            self.priority_event.set()

    def _find_segment_index(self, segment_url: str) -> int:
        """Find the index of a segment URL in the list."""
        try:
            return self.segment_urls.index(segment_url)
        except ValueError:
            return -1

    async def _run(self) -> None:
        """
        Main prefetch loop.

        For live streams: waits until activated by player request before prefetching.
        Priority: Player-requested segment > Sequential prefetch
        After downloading priority segment, continue sequential from that point.

        Prefetching is LIMITED to `prefetch_limit` segments ahead of the player's
        current position to avoid downloading the entire stream.
        """
        logger.info(f"[PlaylistPrefetcher] Loop started for: {self.playlist_url}")

        while not self.cancelled:
            try:
                # Wait for activation (player request) before doing anything
                if not self.activated:
                    try:
                        await asyncio.wait_for(self.priority_event.wait(), timeout=1.0)
                    except asyncio.TimeoutError:
                        continue

                # Check for priority segment first
                async with self._lock:
                    priority_url = self.priority_url
                    self.priority_url = None
                    self.priority_event.clear()

                if priority_url:
                    # Player is already downloading this segment via get_or_download()
                    # We just need to update our indices and skip to prefetching NEXT segments
                    # This avoids duplicate download attempts and inflated cache miss stats
                    priority_index = self._find_segment_index(priority_url)
                    if priority_index >= 0:
                        self.player_index = priority_index
                        self.current_index = priority_index + 1  # Start prefetching from next segment
                        logger.info(
                            f"[PlaylistPrefetcher] Player at index {self.player_index}, "
                            f"will prefetch up to {self.prefetch_limit} segments ahead"
                        )
                    continue

                # Calculate prefetch limit based on player position
                max_prefetch_index = self.player_index + self.prefetch_limit + 1

                # No priority - prefetch next sequential segment (only if within limit)
                if (
                    self.activated
                    and self.current_index < len(self.segment_urls)
                    and self.current_index < max_prefetch_index
                ):
                    url = self.segment_urls[self.current_index]

                    # Skip if already cached or being downloaded
                    if url not in self.downloading:
                        cached = await get_cached_segment(url)
                        if not cached:
                            logger.info(
                                f"[PlaylistPrefetcher] Prefetching [{self.current_index}] "
                                f"(player at {self.player_index}, limit {self.prefetch_limit}): {url}"
                            )
                            await self._download_segment(url)
                        else:
                            logger.debug(f"[PlaylistPrefetcher] Already cached [{self.current_index}]: {url}")

                    self.current_index += 1
                else:
                    # Reached prefetch limit or end of segments - wait for player to advance
                    try:
                        await asyncio.wait_for(self.priority_event.wait(), timeout=1.0)
                    except asyncio.TimeoutError:
                        pass

            except asyncio.CancelledError:
                logger.info(f"[PlaylistPrefetcher] Loop cancelled: {self.playlist_url}")
                return
            except Exception as e:
                logger.warning(f"[PlaylistPrefetcher] Error in loop: {e}")
                await asyncio.sleep(0.5)

        logger.info(f"[PlaylistPrefetcher] Loop ended: {self.playlist_url}")

    async def _download_segment(self, url: str) -> None:
        """
        Download and cache a segment using the parent prebuffer.

        Args:
            url: URL of the segment to download
        """
        if url in self.downloading:
            return

        self.downloading.add(url)
        try:
            # Use the base prebuffer's get_or_download for cross-process coordination
            await self.prebuffer.get_or_download(url, self.headers)
        finally:
            self.downloading.discard(url)


class HLSPreBuffer(BasePrebuffer):
    """
    Pre-buffer system for HLS streams with priority-based prefetching.

    Features:
    - Priority queue: Player-requested segments downloaded first
    - Sequential prefetch: Background prefetch of remaining segments
    - Multi-user support: Multiple users share same prefetcher
    - Automatic cleanup: Inactive prefetchers removed after timeout
    """

    def __init__(
        self,
        max_cache_size: Optional[int] = None,
        prebuffer_segments: Optional[int] = None,
    ):
        """
        Initialize the HLS pre-buffer system.

        Args:
            max_cache_size: Maximum number of segments to cache (uses config if None)
            prebuffer_segments: Number of segments to pre-buffer ahead (uses config if None)
        """
        super().__init__(
            max_cache_size=max_cache_size or settings.hls_prebuffer_cache_size,
            prebuffer_segments=prebuffer_segments or settings.hls_prebuffer_segments,
            max_memory_percent=settings.hls_prebuffer_max_memory_percent,
            emergency_threshold=settings.hls_prebuffer_emergency_threshold,
            segment_ttl=settings.hls_segment_cache_ttl,
        )

        self.inactivity_timeout = settings.hls_prebuffer_inactivity_timeout

        # Active prefetchers: playlist_url -> PlaylistPrefetcher
        self.active_prefetchers: Dict[str, PlaylistPrefetcher] = {}

        # Reverse mapping: segment URL -> playlist_url
        self.segment_to_playlist: Dict[str, str] = {}

        # Lock for prefetcher management
        self._prefetcher_lock = asyncio.Lock()

        # Cleanup task
        self._cleanup_task: Optional[asyncio.Task] = None
        self._cleanup_interval = 30  # Check every 30 seconds

    def log_stats(self) -> None:
        """Log current prebuffer statistics with HLS-specific info."""
        stats = self.stats.to_dict()
        stats["active_prefetchers"] = len(self.active_prefetchers)
        logger.info(f"HLS Prebuffer Stats: {stats}")

    def _extract_segment_urls(self, playlist_content: str, base_url: str) -> List[str]:
        """
        Extract segment URLs from HLS playlist content.

        Args:
            playlist_content: Content of the HLS playlist
            base_url: Base URL for resolving relative URLs

        Returns:
            List of segment URLs
        """
        segment_urls = []
        lines = playlist_content.split("\n")

        for line in lines:
            line = line.strip()
            if line and not line.startswith("#"):
                # Absolute URL
                if line.startswith("http://") or line.startswith("https://"):
                    segment_urls.append(line)
                else:
                    # Relative URL - resolve against base
                    segment_url = urljoin(base_url, line)
                    segment_urls.append(segment_url)

        return segment_urls

    def _is_master_playlist(self, playlist_content: str) -> bool:
        """Check if this is a master playlist (contains variant streams)."""
        return "#EXT-X-STREAM-INF" in playlist_content

    async def register_playlist(
        self,
        playlist_url: str,
        segment_urls: List[str],
        headers: Dict[str, str],
    ) -> None:
        """
        Register a playlist for prefetching.

        Creates a new PlaylistPrefetcher or updates existing one.
        Called by M3U8 processor when a playlist is fetched.

        Args:
            playlist_url: URL of the HLS playlist
            segment_urls: Ordered list of segment URLs from the playlist
            headers: Headers to use for requests
        """
        if not segment_urls:
            logger.debug(f"[register_playlist] No segments, skipping: {playlist_url}")
            return

        async with self._prefetcher_lock:
            # Update reverse mapping
            for url in segment_urls:
                self.segment_to_playlist[url] = playlist_url

            if playlist_url in self.active_prefetchers:
                # Update existing prefetcher
                prefetcher = self.active_prefetchers[playlist_url]
                prefetcher.update_segments(segment_urls)
                prefetcher.headers = headers
                logger.info(f"[register_playlist] Updated existing prefetcher: {playlist_url}")
            else:
                # Create new prefetcher with configured prefetch limit
                prefetcher = PlaylistPrefetcher(
                    playlist_url=playlist_url,
                    segment_urls=segment_urls,
                    headers=headers,
                    prebuffer=self,
                    prefetch_limit=settings.hls_prebuffer_segments,
                )
                self.active_prefetchers[playlist_url] = prefetcher
                prefetcher.start()
                logger.info(
                    f"[register_playlist] Created new prefetcher ({len(segment_urls)} segments, "
                    f"prefetch_limit={settings.hls_prebuffer_segments}): {playlist_url}"
                )

            # Ensure cleanup task is running
            self._ensure_cleanup_task()

    async def request_segment(self, segment_url: str) -> None:
        """
        Player requested a segment - set as priority for prefetching.

        Finds the prefetcher for this segment and adds it to priority queue.
        Called by the segment endpoint when a segment is requested.

        Args:
            segment_url: URL of the segment the player needs
        """
        playlist_url = self.segment_to_playlist.get(segment_url)
        if not playlist_url:
            logger.debug(f"[request_segment] No prefetcher found for: {segment_url}")
            return

        prefetcher = self.active_prefetchers.get(playlist_url)
        if prefetcher:
            await prefetcher.request_priority(segment_url)
        else:
            logger.debug(f"[request_segment] Prefetcher not active for: {playlist_url}")

    def _ensure_cleanup_task(self) -> None:
        """Ensure the cleanup task is running."""
        if self._cleanup_task is None or self._cleanup_task.done():
            self._cleanup_task = asyncio.create_task(self._cleanup_loop())

    async def _cleanup_loop(self) -> None:
        """Periodically clean up inactive prefetchers."""
        while True:
            try:
                await asyncio.sleep(self._cleanup_interval)
                await self._cleanup_inactive_prefetchers()
            except asyncio.CancelledError:
                return
            except Exception as e:
                logger.warning(f"[cleanup_loop] Error: {e}")

    async def _cleanup_inactive_prefetchers(self) -> None:
        """Remove prefetchers that haven't been accessed recently."""
        now = time.time()
        to_remove = []

        async with self._prefetcher_lock:
            for playlist_url, prefetcher in self.active_prefetchers.items():
                inactive_time = now - prefetcher.last_access
                if inactive_time > self.inactivity_timeout:
                    to_remove.append(playlist_url)
                    logger.info(f"[cleanup] Removing inactive prefetcher ({inactive_time:.0f}s): {playlist_url}")

            for playlist_url in to_remove:
                prefetcher = self.active_prefetchers.pop(playlist_url, None)
                if prefetcher:
                    prefetcher.stop()
                    # Clean up reverse mapping
                    for url in prefetcher.segment_urls:
                        self.segment_to_playlist.pop(url, None)

        if to_remove:
            logger.info(f"[cleanup] Removed {len(to_remove)} inactive prefetchers")

    def get_stats(self) -> dict:
        """Get current prebuffer statistics."""
        stats = self.stats.to_dict()
        stats["active_prefetchers"] = len(self.active_prefetchers)
        return stats

    def clear_cache(self) -> None:
        """Clear all prebuffer state and log final stats."""
        self.log_stats()

        # Stop all prefetchers
        for prefetcher in self.active_prefetchers.values():
            prefetcher.stop()

        self.active_prefetchers.clear()
        self.segment_to_playlist.clear()
        self.stats.reset()

        logger.info("HLS pre-buffer state cleared")

    async def close(self) -> None:
        """Close the pre-buffer system."""
        self.clear_cache()
        if self._cleanup_task:
            self._cleanup_task.cancel()


# Global HLS pre-buffer instance
hls_prebuffer = HLSPreBuffer()