"""H.264 → aiortc VideoStreamTrack adapter.

Reads H.264 NAL units produced by ScrcpyAdapter and feeds them as
av.VideoFrame objects to an aiortc VideoStreamTrack for WebRTC streaming.
"""

from __future__ import annotations

import asyncio
import fractions
import logging
import time
from typing import Any

import av
from aiortc import MediaStreamTrack

logger = logging.getLogger(__name__)

CLOCK_RATE = 90_000
TIME_BASE = fractions.Fraction(1, CLOCK_RATE)


class ScrcpyVideoTrack(MediaStreamTrack):
    """aiortc video track fed by H.264 NAL units from scrcpy."""

    kind = "video"

    def __init__(self, target_fps: int = 30) -> None:
        super().__init__()
        self._queue: asyncio.Queue[bytes] = asyncio.Queue(maxsize=120)
        self._codec: av.CodecContext | None = None
        self._target_fps = target_fps
        self._frame_interval = 1.0 / target_fps
        self._start_time: float | None = None
        self._frame_count = 0

    def push_nal(self, data: bytes) -> None:
        """Push raw H.264 NAL unit data (called by ScrcpyAdapter reader)."""
        try:
            self._queue.put_nowait(data)
        except asyncio.QueueFull:
            try:
                self._queue.get_nowait()
            except asyncio.QueueEmpty:
                pass
            try:
                self._queue.put_nowait(data)
            except asyncio.QueueFull:
                pass

    async def recv(self) -> av.VideoFrame:
        """Called by aiortc to get the next video frame."""
        if self._start_time is None:
            self._start_time = time.monotonic()

        while True:
            try:
                data = await asyncio.wait_for(self._queue.get(), timeout=2.0)
            except asyncio.TimeoutError:
                return self._black_frame()
            frames = self._decode_h264(data)
            if frames:
                frame = frames[-1]
                frame.pts = self._frame_count
                frame.time_base = TIME_BASE
                self._frame_count += 1
                return frame

    def _black_frame(self) -> av.VideoFrame:
        """Return a minimal black frame when no data arrives in time."""
        import numpy as np

        arr = np.zeros((480, 640, 3), dtype=np.uint8)
        frame = av.VideoFrame.from_ndarray(arr, format="rgb24")
        frame.pts = self._frame_count
        frame.time_base = TIME_BASE
        self._frame_count += 1
        return frame

    def _decode_h264(self, data: bytes) -> list[av.VideoFrame]:
        if self._codec is None:
            self._codec = av.CodecContext.create("h264", "r")
        try:
            packet = av.Packet(data)
            return self._codec.decode(packet)
        except av.error.InvalidDataError:
            logger.debug("Invalid H.264 data, skipping packet")
            return []

    def stop(self) -> None:
        super().stop()
        if self._codec:
            self._codec = None
