"""Android root-only screen capture engine (screenrecord + MediaCodec).

Mirrors :mod:`mediaprojection_engine` but represents a different source:
Kotlin's :class:`RootScreenrecord` spawns ``su -c screenrecord --output-format=h264``
and pipes the bitstream through a MediaCodec decoder into an ImageReader
Surface. Frames reach this module via :func:`push_frame`, delivered by
:meth:`PythonBridge.pushRootFrame`.

Why a separate engine and not a flag on MediaProjection: the root path
avoids both the one-time consent dialog *and* the persistent capture
indicator that Android 14+ draws while MediaProjection is active, so
it's meaningfully different from the user's perspective. Keeping it as
a distinct engine also lets the dashboard show which backend is live
and lets us give it a higher priority so the factory selects it
automatically when available.
"""

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: Optional[DisplayInfo] = None
_active = False
_frames_received = 0
# screenrecord emits a full bitstream every frame (keyframes aside), so
# the decoder output is always fresh — no need for a "last frame" cache
# like MediaProjection. We still keep one just in case the consumer polls
# faster than the decoder produces.
_last_frame: Optional["ScreenCapture"] = None


def configure(width: int, height: int) -> None:
    """Set display dimensions. Called from Kotlin before server start."""
    global _display_info, _active, _last_frame, _frames_received
    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 (root)",
        width=width,
        height=height,
        x=0,
        y=0,
        is_primary=True,
        refresh_rate=30,
    )
    _active = True
    logger.info("Root screenrecord engine configured: %dx%d", width, height)


def push_frame(rgba_bytes: bytes, width: int, height: int) -> None:
    """Push a decoded frame from Kotlin into the Python pipeline."""
    global _frames_received, _last_frame
    _frames_received += 1
    if _frames_received == 1 or _frames_received % 100 == 0:
        logger.info("Root screenrecord: received %d frames", _frames_received)

    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)
    _last_frame = frame

    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 / CaptureEngine
# ---------------------------------------------------------------------------


class RootScreenrecordCaptureStream(CaptureStream):
    """Reads frames pushed by Kotlin's RootScreenrecord from the module 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(
                "Root screenrecord engine not configured. "
                "This engine is only available on rooted Android devices."
            )
        self._initialized = True
        logger.info("Root screenrecord capture stream initialized")

    def capture_frame(self) -> Optional[ScreenCapture]:
        if not self._initialized:
            self.initialize()
        try:
            return _frame_queue.get(timeout=0.1)
        except queue.Empty:
            return _last_frame

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


class RootScreenrecordEngine(CaptureEngine):
    """Root-only Android capture engine. Preferred over MediaProjection."""

    ENGINE_TYPE = "root_screenrecord"
    # Higher than MediaProjection (100) — factory picks this when available.
    ENGINE_PRIORITY = 110
    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]
    ) -> RootScreenrecordCaptureStream:
        return RootScreenrecordCaptureStream(display_index, config)
