"""High-performance Android screen capture via the scrcpy protocol.

Uses the ``scrcpy-client`` library to stream H.264/H.265 video from an
Android device over ADB.  Frames are decoded by PyAV and delivered as
NumPy arrays at up to 60 FPS — a major upgrade over the ``adb screencap``
polling approach in :mod:`scrcpy_engine` (~1-2 FPS).

Prerequisites (pip packages):
    - scrcpy-client>=0.5.0  (pulls in ``av``, ``adbutils``, ``numpy``)

The library pushes a small (~35 KB) ``scrcpy-server.jar`` to the device,
launches it via ``app_process``, and opens an ADB socket for the encoded
video stream.  No APK installation, no root.
"""

import threading
from typing import Any, Dict, List

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__)

try:
    import scrcpy
    from adbutils import adb as adb_client

    _HAS_SCRCPY_CLIENT = True
except ImportError:
    _HAS_SCRCPY_CLIENT = False

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _list_devices() -> List[Dict[str, Any]]:
    """Enumerate connected ADB devices via adbutils."""
    if not _HAS_SCRCPY_CLIENT:
        return []

    devices = []
    try:
        for dev in adb_client.device_list():
            serial = dev.serial
            # Query model name
            try:
                model = dev.prop.model or serial
            except Exception:
                model = serial
            # Query screen resolution
            width, height = 1920, 1080
            try:
                wm_output = dev.shell("wm size")
                for line in wm_output.strip().splitlines():
                    if "x" in line:
                        parts = line.split()[-1].split("x")
                        width, height = int(parts[0]), int(parts[1])
                        break
            except Exception:
                pass

            devices.append(
                {
                    "serial": serial,
                    "model": model,
                    "width": width,
                    "height": height,
                }
            )
    except Exception as e:
        logger.debug("Failed to enumerate ADB devices: %s", e)

    return devices


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


class ScrcpyClientCaptureStream(CaptureStream):
    """H.264/H.265 video stream from an Android device via scrcpy protocol.

    The ``scrcpy.Client`` runs in a background thread.  On each decoded
    frame it invokes our callback, which stores the latest frame.
    ``capture_frame()`` returns it without blocking.
    """

    def __init__(self, display_index: int, config: Dict[str, Any]):
        super().__init__(display_index, config)
        self._client: "scrcpy.Client" | None = None
        self._latest_frame: ScreenCapture | None = None
        self._frame_lock = threading.Lock()
        self._frame_event = threading.Event()
        self._client_thread: threading.Thread | None = None
        self._device_serial: str | None = None

    def initialize(self) -> None:
        if self._initialized:
            return

        devices = _list_devices()
        if not devices:
            raise RuntimeError(
                "No ADB devices found. Connect a device via USB or " "'adb connect <ip>' for WiFi."
            )
        if self.display_index >= len(devices):
            raise RuntimeError(
                f"Device index {self.display_index} out of range "
                f"(found {len(devices)} device(s))"
            )

        device = devices[self.display_index]
        self._device_serial = device["serial"]

        max_fps = self.config.get("max_fps", 30)
        max_size = self.config.get("max_size", 480)
        bitrate = self.config.get("bitrate", 2_000_000)

        logger.info(
            "scrcpy-client: connecting to %s (%s, %dx%d) — " "max_fps=%d, max_size=%d, bitrate=%d",
            self._device_serial,
            device["model"],
            device["width"],
            device["height"],
            max_fps,
            max_size,
            bitrate,
        )

        self._client = scrcpy.Client(
            device=self._device_serial,
            max_fps=max_fps,
            max_size=max_size,
            bitrate=bitrate,
        )
        self._client.add_listener(scrcpy.EVENT_FRAME, self._on_frame)

        # scrcpy.Client.start() blocks, so run in a thread
        self._client_thread = threading.Thread(
            target=self._run_client, daemon=True, name="scrcpy-client"
        )
        self._client_thread.start()

        # Wait for first frame (with timeout)
        if not self._frame_event.wait(timeout=10.0):
            logger.warning(
                "scrcpy-client: no frame received within 10s from %s",
                self._device_serial,
            )

        self._initialized = True
        logger.info("scrcpy-client: stream initialized for %s", self._device_serial)

    def _run_client(self) -> None:
        """Start the scrcpy client (blocking)."""
        try:
            self._client.start(threaded=False)
        except Exception as e:
            logger.error("scrcpy-client: error for %s: %s", self._device_serial, e)

    def _on_frame(self, frame: np.ndarray) -> None:
        """Callback invoked by scrcpy-client for each decoded frame.

        The frame arrives as a BGR numpy array; we convert to RGB to
        match the ``ScreenCapture`` convention.
        """
        if frame is None:
            return

        # BGR → RGB
        rgb = frame[:, :, ::-1].copy()
        h, w = rgb.shape[:2]

        with self._frame_lock:
            self._latest_frame = ScreenCapture(
                image=rgb,
                width=w,
                height=h,
                display_index=self.display_index,
            )
            self._frame_event.set()

    def capture_frame(self) -> ScreenCapture | None:
        if not self._initialized:
            self.initialize()

        if not self._frame_event.is_set():
            return None

        with self._frame_lock:
            frame = self._latest_frame
            self._frame_event.clear()

        return frame

    def cleanup(self) -> None:
        if self._client is not None:
            try:
                self._client.stop()
            except Exception as e:
                logger.debug("scrcpy-client: stop error: %s", e)

        if self._client_thread is not None and self._client_thread.is_alive():
            self._client_thread.join(timeout=5)

        self._client = None
        self._client_thread = None
        self._latest_frame = None
        self._frame_event.clear()
        self._initialized = False
        logger.info("scrcpy-client: stream cleaned up (%s)", self._device_serial)


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


class ScrcpyClientEngine(CaptureEngine):
    """High-performance Android capture via scrcpy H.264 streaming.

    Requires the ``scrcpy-client`` pip package (optional dependency).
    When available, this engine is preferred over :class:`ScrcpyEngine`
    (which falls back to ``adb screencap`` at ~1-2 FPS).

    Prerequisites:
        - ``pip install scrcpy-client`` (pulls in PyAV, adbutils)
        - adb on PATH (for device connection)
        - USB debugging enabled on Android device
    """

    ENGINE_TYPE = "scrcpy_client"
    ENGINE_PRIORITY = 10  # Higher than ScrcpyEngine (5)
    HAS_OWN_DISPLAYS = True

    @classmethod
    def is_available(cls) -> bool:
        return _HAS_SCRCPY_CLIENT

    @classmethod
    def get_default_config(cls) -> Dict[str, Any]:
        return {
            "max_fps": 30,
            "max_size": 480,
            "bitrate": 2_000_000,
        }

    @classmethod
    def get_available_displays(cls) -> List[DisplayInfo]:
        devices = _list_devices()
        displays = []
        for idx, device in enumerate(devices):
            displays.append(
                DisplayInfo(
                    index=idx,
                    name=f"{device['model']} ({device['serial']})",
                    width=device["width"],
                    height=device["height"],
                    x=idx * 500,
                    y=0,
                    is_primary=(idx == 0),
                    refresh_rate=60,
                )
            )
        logger.debug("scrcpy-client: detected %d Android device(s)", len(displays))
        return displays

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