"""Android MediaProjection capture engine.

Receives screen frames pushed from Kotlin (via Chaquopy) through a
module-level frame queue.  The Kotlin layer captures the screen using
the ``MediaProjection`` API and calls :func:`push_frame` with raw RGBA
bytes for each frame.

This engine is only available when running inside an Android app that
has set up the frame queue.
"""

import queue
from typing import Any, Dict, List, Optional

import numpy as np

from ledgrab.core.capture_engines.base import (
    CaptureEngine,
    CaptureStream,
    DisplayInfo,
    ScreenCapture,
)
from ledgrab.utils import get_logger

logger = get_logger(__name__)

# ---------------------------------------------------------------------------
# Frame queue — the bridge between Kotlin and Python
# ---------------------------------------------------------------------------

_frame_queue: queue.Queue["ScreenCapture"] = queue.Queue(maxsize=2)
_display_info: DisplayInfo | None = None
_active = False
_frames_received = 0
_frames_consumed = 0
# MediaProjection only fires onImageAvailable when the screen changes.
# Cache the last frame so preview stays usable on static screens.
_last_frame: Optional["ScreenCapture"] = None


def configure(width: int, height: int) -> None:
    """Set display dimensions.  Called from Kotlin before server start.

    Drains any stale frames from a previous capture session so the
    first frame after restart is actually current.
    """
    global _display_info, _active, _last_frame, _frames_received
    # Drain the queue — frames from a previous capture session are stale
    while not _frame_queue.empty():
        try:
            _frame_queue.get_nowait()
        except queue.Empty:
            break
    _last_frame = None
    _frames_received = 0
    _display_info = DisplayInfo(
        index=0,
        name="Android TV Screen",
        width=width,
        height=height,
        x=0,
        y=0,
        is_primary=True,
        refresh_rate=60,
    )
    _active = True
    logger.info("MediaProjection engine configured: %dx%d", width, height)


def push_frame(rgba_bytes: bytes, width: int, height: int) -> None:
    """Push a captured frame from Kotlin into the Python pipeline.

    Called from Kotlin's ``PythonBridge.pushFrame()`` on the capture
    thread.  The RGBA byte buffer is converted to an RGB NumPy array
    and placed on the queue.  If the queue is full (Python consumer is
    slow), the oldest frame is dropped.

    Args:
        rgba_bytes: Raw RGBA pixel data (width * height * 4 bytes).
        width: Frame width in pixels.
        height: Frame height in pixels.
    """
    global _frames_received
    _frames_received += 1
    if _frames_received == 1 or _frames_received % 100 == 0:
        logger.info("MediaProjection: received %d frames", _frames_received)

    # RGBA → RGB (drop alpha channel)
    rgba = np.frombuffer(rgba_bytes, dtype=np.uint8).reshape((height, width, 4))
    rgb = rgba[:, :, :3].copy()

    frame = ScreenCapture(
        image=rgb,
        width=width,
        height=height,
        display_index=0,
    )

    global _last_frame
    _last_frame = frame

    # Drop oldest frame if queue is full (non-blocking)
    try:
        _frame_queue.put_nowait(frame)
    except queue.Full:
        try:
            _frame_queue.get_nowait()
        except queue.Empty:
            pass
        try:
            _frame_queue.put_nowait(frame)
        except queue.Full:
            pass


def shutdown() -> None:
    """Deactivate the engine.  Called when the Android app stops."""
    global _active
    _active = False


# ---------------------------------------------------------------------------
# CaptureStream
# ---------------------------------------------------------------------------


class MediaProjectionCaptureStream(CaptureStream):
    """Reads frames pushed by Kotlin from the module-level queue."""

    def __init__(self, display_index: int, config: Dict[str, Any]):
        super().__init__(display_index, config)

    def initialize(self) -> None:
        if self._initialized:
            return
        if not _active:
            raise RuntimeError(
                "MediaProjection engine not configured. "
                "This engine is only available inside the Android app."
            )
        self._initialized = True
        logger.info("MediaProjection capture stream initialized")

    def capture_frame(self) -> ScreenCapture | None:
        if not self._initialized:
            self.initialize()
        # Prefer fresh frames from the queue; fall back to the last
        # received frame when the screen is static (MediaProjection
        # only emits frames on actual content changes).
        try:
            return _frame_queue.get(timeout=0.1)
        except queue.Empty:
            return _last_frame

    def cleanup(self) -> None:
        # Drain the queue
        while not _frame_queue.empty():
            try:
                _frame_queue.get_nowait()
            except queue.Empty:
                break
        self._initialized = False
        logger.info("MediaProjection capture stream cleaned up")


# ---------------------------------------------------------------------------
# CaptureEngine
# ---------------------------------------------------------------------------


class MediaProjectionEngine(CaptureEngine):
    """Android MediaProjection capture engine.

    Only available when running inside the LedGrab Android app.
    The Kotlin layer calls :func:`configure` at startup and
    :func:`push_frame` for each captured frame.
    """

    ENGINE_TYPE = "mediaprojection"
    ENGINE_PRIORITY = 100  # Highest priority on Android
    HAS_OWN_DISPLAYS = True

    @classmethod
    def is_available(cls) -> bool:
        return _active and _display_info is not None

    @classmethod
    def get_default_config(cls) -> Dict[str, Any]:
        return {}

    @classmethod
    def get_available_displays(cls) -> List[DisplayInfo]:
        if _display_info is not None:
            return [_display_info]
        return []

    @classmethod
    def create_stream(
        cls, display_index: int, config: Dict[str, Any]
    ) -> MediaProjectionCaptureStream:
        return MediaProjectionCaptureStream(display_index, config)
