Blame for mediaflow_proxy/mpd_processor.py

import asyncio
import logging
import math
import time

from fastapi import Request, Response, HTTPException

from mediaflow_proxy.drm.decrypter import decrypt_segment, process_drm_init_segment
from mediaflow_proxy.utils.crypto_utils import encryption_handler
from mediaflow_proxy.utils.http_utils import (
    encode_mediaflow_proxy_url,
    get_original_scheme,
    ProxyRequestHeaders,
    apply_header_manipulation,
)
from mediaflow_proxy.utils.dash_prebuffer import dash_prebuffer
from mediaflow_proxy.utils.cache_utils import get_cached_processed_init, set_cached_processed_init
from mediaflow_proxy.utils.m3u8_processor import SkipSegmentFilter
from mediaflow_proxy.configs import settings

logger = logging.getLogger(__name__)


async def process_manifest(
    request: Request,
    mpd_dict: dict,
    proxy_headers: ProxyRequestHeaders,
    key_id: str = None,
    key: str = None,
    resolution: str = None,
    skip_segments: list = None,
) -> Response:
    """
    Processes the MPD manifest and converts it to an HLS manifest.

    Args:
        request (Request): The incoming HTTP request.
        mpd_dict (dict): The MPD manifest data.
        proxy_headers (ProxyRequestHeaders): The headers to include in the request.
        key_id (str, optional): The DRM key ID. Defaults to None.
        key (str, optional): The DRM key. Defaults to None.
        resolution (str, optional): Target resolution (e.g., '1080p', '720p'). Defaults to None.
        skip_segments (list, optional): List of time segments to skip. Each item should have 'start' and 'end' keys.

    Returns:
        Response: The HLS manifest as an HTTP response.
    """
    hls_content = build_hls(mpd_dict, request, key_id, key, resolution, skip_segments)

    # Start DASH pre-buffering in background if enabled
    if settings.enable_dash_prebuffer:
        # Extract headers for pre-buffering
        headers = {}
        for key, value in request.query_params.items():
            if key.startswith("h_"):
                headers[key[2:]] = value

        # Get the original MPD URL from the request
        mpd_url = request.query_params.get("d", "")
        if mpd_url:
            # Start pre-buffering in background
            asyncio.create_task(dash_prebuffer.prebuffer_dash_manifest(mpd_url, headers))

    return Response(content=hls_content, media_type="application/vnd.apple.mpegurl", headers=proxy_headers.response)


async def process_playlist(
    request: Request,
    mpd_dict: dict,
    profile_id: str,
    proxy_headers: ProxyRequestHeaders,
    skip_segments: list = None,
    start_offset: float = None,
) -> Response:
    """
    Processes the MPD manifest and converts it to an HLS playlist for a specific profile.

    Args:
        request (Request): The incoming HTTP request.
        mpd_dict (dict): The MPD manifest data.
        profile_id (str): The profile ID to generate the playlist for.
        proxy_headers (ProxyRequestHeaders): The headers to include in the request.
        skip_segments (list, optional): List of time segments to skip. Each item should have 'start' and 'end' keys.
        start_offset (float, optional): Start offset in seconds for live streams.

    Returns:
        Response: The HLS playlist as an HTTP response.

    Raises:
        HTTPException: If the profile is not found in the MPD manifest.
    """
    matching_profiles = [p for p in mpd_dict["profiles"] if p["id"] == profile_id]
    if not matching_profiles:
        raise HTTPException(status_code=404, detail="Profile not found")

    hls_content = build_hls_playlist(mpd_dict, matching_profiles, request, skip_segments, start_offset)

    # Trigger prebuffering of upcoming segments for live streams
    if settings.enable_dash_prebuffer and mpd_dict.get("isLive", False):
        # Extract headers for pre-buffering
        headers = {}
        for key, value in request.query_params.items():
            if key.startswith("h_"):
                headers[key[2:]] = value

        # Use the new prefetch method for live playlists
        asyncio.create_task(dash_prebuffer.prefetch_for_live_playlist(matching_profiles, headers))

    # Don't include propagate headers for playlists - they should only apply to segments
    response_headers = apply_header_manipulation({}, proxy_headers, include_propagate=False)
    return Response(content=hls_content, media_type="application/vnd.apple.mpegurl", headers=response_headers)


async def process_segment(
    init_content: bytes,
    segment_content: bytes,
    mimetype: str,
    proxy_headers: ProxyRequestHeaders,
    key_id: str = None,
    key: str = None,
    use_map: bool = False,
) -> Response:
    """
    Processes and decrypts a media segment.

    Args:
        init_content (bytes): The initialization segment content.
        segment_content (bytes): The media segment content.
        mimetype (str): The MIME type of the segment.
        proxy_headers (ProxyRequestHeaders): The headers to include in the request.
        key_id (str, optional): The DRM key ID. Defaults to None.
        key (str, optional): The DRM key. Defaults to None.
        use_map (bool, optional): If True, init segment is served separately via EXT-X-MAP,
            so don't concatenate init with segment. Defaults to False.

    Returns:
        Response: The decrypted segment as an HTTP response.
    """
    if key_id and key:
        # For DRM protected content
        now = time.time()
        decrypted_content = decrypt_segment(init_content, segment_content, key_id, key, include_init=not use_map)
        logger.info(f"Decryption of {mimetype} segment took {time.time() - now:.4f} seconds")
    else:
        # For non-DRM protected content
        if use_map:
            # Init is served separately via EXT-X-MAP
            decrypted_content = segment_content
        else:
            # Concatenate init and segment content
            decrypted_content = init_content + segment_content

    response_headers = apply_header_manipulation({}, proxy_headers)
    return Response(content=decrypted_content, media_type=mimetype, headers=response_headers)


async def process_init_segment(
    init_content: bytes,
    mimetype: str,
    proxy_headers: ProxyRequestHeaders,
    key_id: str = None,
    key: str = None,
    init_url: str = None,
) -> Response:
    """
    Processes an initialization segment for EXT-X-MAP.

    Args:
        init_content (bytes): The initialization segment content.
        mimetype (str): The MIME type of the segment.
        proxy_headers (ProxyRequestHeaders): The headers to include in the request.
        key_id (str, optional): The DRM key ID. Defaults to None.
        key (str, optional): The DRM key. Defaults to None.
        init_url (str, optional): The init URL for caching. Defaults to None.

    Returns:
        Response: The processed init segment as an HTTP response.
    """
    if key_id and key:
        # Check if we have a cached processed version
        if init_url:
            cached_processed = await get_cached_processed_init(init_url, key_id)
            if cached_processed:
                logger.debug(f"Using cached processed init segment for {init_url}")
                response_headers = apply_header_manipulation({}, proxy_headers)
                return Response(content=cached_processed, media_type=mimetype, headers=response_headers)

        # For DRM protected content, we need to process the init segment
        # to remove encryption-related boxes but keep the moov structure
        now = time.time()
        processed_content = process_drm_init_segment(init_content, key_id, key)
        logger.info(f"Processing of {mimetype} init segment took {time.time() - now:.4f} seconds")

        # Cache the processed init segment
        if init_url:
            await set_cached_processed_init(init_url, key_id, processed_content, ttl=3600)
    else:
        # For non-DRM protected content, just return the init segment as-is
        processed_content = init_content

    response_headers = apply_header_manipulation({}, proxy_headers)
    return Response(content=processed_content, media_type=mimetype, headers=response_headers)


def build_hls(
    mpd_dict: dict,
    request: Request,
    key_id: str = None,
    key: str = None,
    resolution: str = None,
    skip_segments: list = None,
) -> str:
    """
    Builds an HLS manifest from the MPD manifest.

    Args:
        mpd_dict (dict): The MPD manifest data.
        request (Request): The incoming HTTP request.
        key_id (str, optional): The DRM key ID. Defaults to None.
        key (str, optional): The DRM key. Defaults to None.
        resolution (str, optional): Target resolution (e.g., '1080p', '720p'). Defaults to None.
        skip_segments (list, optional): List of time segments to skip. Each item should have 'start' and 'end' keys.

    Returns:
        str: The HLS manifest as a string.
    """
    hls = ["#EXTM3U", "#EXT-X-VERSION:6"]
    query_params = dict(request.query_params)

    # Preserve skip parameter in query params so it propagates to playlists
    if skip_segments:
        # Convert back to compact format for URL
        skip_str = ",".join(f"{s['start']}-{s['end']}" for s in skip_segments)
        query_params["skip"] = skip_str
    has_encrypted = query_params.pop("has_encrypted", False)

    video_profiles = {}
    audio_profiles = {}

    # Get the base URL for the playlist_endpoint endpoint
    proxy_url = request.url_for("playlist_endpoint")
    proxy_url = str(proxy_url.replace(scheme=get_original_scheme(request)))

    for profile in mpd_dict["profiles"]:
        query_params.update({"profile_id": profile["id"], "key_id": key_id or "", "key": key or ""})
        playlist_url = encode_mediaflow_proxy_url(
            proxy_url,
            query_params=query_params,
            encryption_handler=encryption_handler if has_encrypted else None,
        )

        if "video" in profile["mimeType"]:
            video_profiles[profile["id"]] = (profile, playlist_url)
        elif "audio" in profile["mimeType"]:
            audio_profiles[profile["id"]] = (profile, playlist_url)

    # Filter video profiles by resolution if specified
    if resolution and video_profiles:
        video_profiles = _filter_video_profiles_by_resolution(video_profiles, resolution)

    # Add audio streams
    for i, (profile, playlist_url) in enumerate(audio_profiles.values()):
        is_default = "YES" if i == 0 else "NO"  # Set the first audio track as default
        hls.append(
            f'#EXT-X-MEDIA:TYPE=AUDIO,GROUP-ID="audio",NAME="{profile["id"]}",DEFAULT={is_default},AUTOSELECT={is_default},LANGUAGE="{profile.get("lang", "und")}",URI="{playlist_url}"'
        )

    # Add video streams
    for profile, playlist_url in video_profiles.values():
        # Only add AUDIO attribute if there are audio profiles available
        audio_attr = ',AUDIO="audio"' if audio_profiles else ""
        hls.append(
            f'#EXT-X-STREAM-INF:BANDWIDTH={profile["bandwidth"]},RESOLUTION={profile["width"]}x{profile["height"]},CODECS="{profile["codecs"]}",FRAME-RATE={profile["frameRate"]}{audio_attr}'
        )
        hls.append(playlist_url)

    return "\n".join(hls)


def _filter_video_profiles_by_resolution(video_profiles: dict, target_resolution: str) -> dict:
    """
    Filter video profiles to select the one matching the target resolution.
    Falls back to closest lower resolution if exact match not found.

    Args:
        video_profiles: Dictionary of profile_id -> (profile, playlist_url).
        target_resolution: Target resolution string (e.g., '1080p', '720p').

    Returns:
        Filtered dictionary with only the selected profile.
    """
    # Parse target height from "1080p" -> 1080
    target_height = int(target_resolution.rstrip("p"))

    # Convert to list and sort by height descending
    profiles_list = [
        (profile_id, profile, playlist_url)
        for profile_id, (profile, playlist_url) in video_profiles.items()
        if profile.get("height", 0) > 0
    ]

    if not profiles_list:
        logger.warning("No video profiles with valid height found, returning all profiles")
        return video_profiles

    sorted_profiles = sorted(profiles_list, key=lambda x: x[1]["height"], reverse=True)

    # Find exact match or closest lower
    selected = None
    for profile_id, profile, playlist_url in sorted_profiles:
        if profile["height"] <= target_height:
            selected = (profile_id, profile, playlist_url)
            break

    # If all profiles are higher than target, use lowest available
    if selected is None:
        selected = sorted_profiles[-1]

    profile_id, profile, playlist_url = selected
    logger.info(
        f"Selected MPD video profile with resolution {profile['width']}x{profile['height']} for target {target_resolution}"
    )

    return {profile_id: (profile, playlist_url)}


def build_hls_playlist(
    mpd_dict: dict, profiles: list[dict], request: Request, skip_segments: list = None, start_offset: float = None
) -> str:
    """
    Builds an HLS playlist from the MPD manifest for specific profiles.

    Args:
        mpd_dict (dict): The MPD manifest data.
        profiles (list[dict]): The profiles to include in the playlist.
        request (Request): The incoming HTTP request.
        skip_segments (list, optional): List of time segments to skip. Each item should have 'start' and 'end' keys.
        start_offset (float, optional): Start offset in seconds for live streams. Defaults to settings.livestream_start_offset for live.

    Returns:
        str: The HLS playlist as a string.
    """
    hls = ["#EXTM3U", "#EXT-X-VERSION:6"]

    added_segments = 0
    skipped_segments = 0
    is_live = mpd_dict.get("isLive", False)

    # Inject EXT-X-START for live streams (enables prebuffering by starting behind live edge)
    # User-provided start_offset always takes precedence; otherwise use default for live streams only
    effective_start_offset = (
        start_offset if start_offset is not None else (settings.livestream_start_offset if is_live else None)
    )
    if effective_start_offset is not None:
        hls.append(f"#EXT-X-START:TIME-OFFSET={effective_start_offset:.1f},PRECISE=YES")

    # Initialize skip filter if skip_segments provided
    skip_filter = SkipSegmentFilter(skip_segments) if skip_segments else None

    # Use EXT-X-MAP for live streams to avoid duplicate moov atoms
    use_map = is_live

    proxy_url = request.url_for("segment_endpoint")
    proxy_url = str(proxy_url.replace(scheme=get_original_scheme(request)))

    # Get init endpoint URL for EXT-X-MAP
    init_proxy_url = request.url_for("init_endpoint")
    init_proxy_url = str(init_proxy_url.replace(scheme=get_original_scheme(request)))

    for index, profile in enumerate(profiles):
        segments = profile["segments"]
        if not segments:
            logger.warning(f"No segments found for profile {profile['id']}")
            continue

        if is_live:
            depth = max(settings.mpd_live_playlist_depth, 1)
            trimmed_segments = segments[-depth:]
        else:
            trimmed_segments = segments

        # Add headers for only the first profile
        if index == 0:
            first_segment = trimmed_segments[0]
            extinf_values = [f["extinf"] for f in trimmed_segments if "extinf" in f]
            target_duration = math.ceil(max(extinf_values)) if extinf_values else 3

            # Align HLS media sequence with MPD-provided numbering when available
            mpd_start_number = profile.get("segment_template_start_number")
            sequence = first_segment.get("number")

            if sequence is None:
                # Fallback to MPD template start number
                if mpd_start_number is not None:
                    sequence = mpd_start_number
                else:
                    # As a last resort, derive from timeline information
                    time_val = first_segment.get("time")
                    duration_val = first_segment.get("duration_mpd_timescale")
                    if time_val is not None and duration_val and duration_val > 0:
                        sequence = math.floor(time_val / duration_val)
                    else:
                        sequence = 1

            hls.extend(
                [
                    f"#EXT-X-TARGETDURATION:{target_duration}",
                    f"#EXT-X-MEDIA-SEQUENCE:{sequence}",
                ]
            )
            # For live streams, don't set PLAYLIST-TYPE to allow sliding window
            if not is_live:
                hls.append("#EXT-X-PLAYLIST-TYPE:VOD")

        init_url = profile["initUrl"]
        # For SegmentBase profiles, we may have byte range for initialization segment
        init_range = profile.get("initRange")

        query_params = dict(request.query_params)
        query_params.pop("profile_id", None)
        query_params.pop("d", None)
        has_encrypted = query_params.pop("has_encrypted", False)

        # Add EXT-X-MAP for init segment (for live streams or when beneficial)
        if use_map:
            init_query_params = {
                "init_url": init_url,
                "mime_type": profile["mimeType"],
                "is_live": "true" if is_live else "false",
            }
            if init_range:
                init_query_params["init_range"] = init_range
            # Add key parameters
            if query_params.get("key_id"):
                init_query_params["key_id"] = query_params["key_id"]
            if query_params.get("key"):
                init_query_params["key"] = query_params["key"]
            # Add api_password for authentication
            if query_params.get("api_password"):
                init_query_params["api_password"] = query_params["api_password"]

            init_map_url = encode_mediaflow_proxy_url(
                init_proxy_url,
                query_params=init_query_params,
                encryption_handler=encryption_handler if has_encrypted else None,
            )
            hls.append(f'#EXT-X-MAP:URI="{init_map_url}"')

        need_discontinuity = False
        for segment in trimmed_segments:
            duration = segment["extinf"]

            # Check if this segment should be skipped
            if skip_filter:
                if skip_filter.should_skip_segment(duration):
                    skip_filter.advance_time(duration)
                    skipped_segments += 1
                    need_discontinuity = True
                    continue
                skip_filter.advance_time(duration)

            # Add discontinuity marker after skipped segments
            if need_discontinuity:
                hls.append("#EXT-X-DISCONTINUITY")
                need_discontinuity = False

            program_date_time = segment.get("program_date_time")
            if program_date_time:
                hls.append(f"#EXT-X-PROGRAM-DATE-TIME:{program_date_time}")
            hls.append(f"#EXTINF:{duration:.3f},")

            segment_query_params = {
                "init_url": init_url,
                "segment_url": segment["media"],
                "mime_type": profile["mimeType"],
                "is_live": "true" if is_live else "false",
            }

            # Add use_map flag so segment endpoint knows not to include init
            if use_map:
                segment_query_params["use_map"] = "true"

            # Add byte range parameters for SegmentBase
            if init_range:
                segment_query_params["init_range"] = init_range
            # Segment may also have its own range (for SegmentBase)
            if "initRange" in segment:
                segment_query_params["init_range"] = segment["initRange"]

            query_params.update(segment_query_params)
            hls.append(
                encode_mediaflow_proxy_url(
                    proxy_url,
                    query_params=query_params,
                    encryption_handler=encryption_handler if has_encrypted else None,
                )
            )
            added_segments += 1

    if not mpd_dict["isLive"]:
        hls.append("#EXT-X-ENDLIST")

    if skip_filter and skipped_segments > 0:
        logger.info(f"Added {added_segments} segments to HLS playlist (skipped {skipped_segments} segments)")
    else:
        logger.info(f"Added {added_segments} segments to HLS playlist")
    return "\n".join(hls)