"""MSS-based screen capture engine (cross-platform)."""

from typing import Any, Dict, List, Optional

import mss
import numpy as np

try:
    import cv2

    _HAS_CV2 = True
except ImportError:
    _HAS_CV2 = False

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

logger = get_logger(__name__)

# Rotating RGB output pool: keeps prior frame references stable for any
# consumer still reading them while a new frame is written.
_RGB_POOL_SIZE = 3
# Number of bytes from .raw to hash for change detection (cheap pre-check
# that avoids the full BGRA→RGB conversion when the screen is idle).
_CHANGE_DETECT_BYTES = 256


class MSSCaptureStream(CaptureStream):
    """MSS capture stream for a specific display."""

    def __init__(self, display_index: int, config: Dict[str, Any]):
        super().__init__(display_index, config)
        self._sct = None
        # Pre-allocated RGB destination pool — avoids per-frame allocation.
        self._rgb_pool: list = [None] * _RGB_POOL_SIZE
        self._rgb_idx: int = 0
        self._rgb_shape: tuple = (0, 0)
        # Cheap hash of the previous .raw bytes, for change detection.
        self._prev_hash: Optional[int] = None

    def initialize(self) -> None:
        try:
            self._sct = mss.mss()
            self._initialized = True
            logger.info(f"MSS capture stream initialized (display={self.display_index})")
        except Exception as e:
            raise RuntimeError(f"Failed to initialize MSS: {e}")

    def cleanup(self) -> None:
        if self._sct:
            self._sct.close()
            self._sct = None
        self._initialized = False
        self._prev_hash = None
        logger.info(f"MSS capture stream cleaned up (display={self.display_index})")

    def capture_frame(self) -> Optional[ScreenCapture]:
        if not self._initialized:
            self.initialize()

        try:
            # mss monitors[0] is the combined screen, monitors[1+] are individual displays
            monitor_index = self.display_index + 1

            if monitor_index >= len(self._sct.monitors):
                raise ValueError(
                    f"Invalid display index {self.display_index}. "
                    f"Available displays: 0-{len(self._sct.monitors) - 2}"
                )

            monitor = self._sct.monitors[monitor_index]
            screenshot = self._sct.grab(monitor)

            # Cheap change detection: hash a small slice of the raw BGRA
            # buffer. ~256 bytes is enough to differentiate any cursor/pixel
            # change. Skips the BGRA→RGB conversion when nothing changed
            # (common on idle desktops). DXcam/BetterCam return None in this
            # case natively; mss does not, so we add it here.
            raw = screenshot.raw
            sample = bytes(raw[:_CHANGE_DETECT_BYTES])
            cur_hash = hash(sample)
            if cur_hash == self._prev_hash:
                return None
            self._prev_hash = cur_hash

            height = screenshot.height
            width = screenshot.width

            # Reshape .raw (BGRA) — zero-copy view over the screenshot's buffer.
            # ``screenshot.rgb`` (used previously) is a pure-Python BGRA→RGB
            # rebuild costing ~6 MB/frame at 1080p in the slowest possible
            # way. cv2.cvtColor is SIMD and writes directly into our pool.
            bgra = np.frombuffer(raw, dtype=np.uint8).reshape(height, width, 4)

            if self._rgb_shape != (height, width):
                for i in range(_RGB_POOL_SIZE):
                    self._rgb_pool[i] = np.empty((height, width, 3), dtype=np.uint8)
                self._rgb_shape = (height, width)

            dst = self._rgb_pool[self._rgb_idx]
            self._rgb_idx = (self._rgb_idx + 1) % _RGB_POOL_SIZE

            if _HAS_CV2:
                cv2.cvtColor(bgra, cv2.COLOR_BGRA2RGB, dst=dst)
            else:
                dst[..., 0] = bgra[..., 2]
                dst[..., 1] = bgra[..., 1]
                dst[..., 2] = bgra[..., 0]

            logger.debug(
                "MSS captured frame",
                display=self.display_index,
                w=monitor["width"],
                h=monitor["height"],
            )

            return ScreenCapture(
                image=dst,
                width=monitor["width"],
                height=monitor["height"],
                display_index=self.display_index,
            )

        except ValueError:
            raise
        except Exception as e:
            logger.error(f"Failed to capture display {self.display_index} with MSS: {e}")
            raise RuntimeError(f"Screen capture failed: {e}")


class MSSEngine(CaptureEngine):
    """MSS-based screen capture engine.

    Uses the mss library for cross-platform screen capture support.
    Works on Windows, macOS, and Linux.
    """

    ENGINE_TYPE = "mss"
    ENGINE_PRIORITY = 1

    @classmethod
    def is_available(cls) -> bool:
        try:
            import importlib.util

            return importlib.util.find_spec("mss") is not None
        except (ImportError, ModuleNotFoundError, ValueError):
            return False

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

    @classmethod
    def get_available_displays(cls) -> List[DisplayInfo]:
        try:
            monitor_names = get_monitor_names()
            refresh_rates = get_monitor_refresh_rates()

            displays = []
            with mss.mss() as sct:
                for idx, monitor in enumerate(sct.monitors[1:], start=0):
                    friendly_name = monitor_names.get(idx, f"Display {idx}")
                    refresh_rate = refresh_rates.get(idx, 60)

                    displays.append(
                        DisplayInfo(
                            index=idx,
                            name=friendly_name,
                            width=monitor["width"],
                            height=monitor["height"],
                            x=monitor["left"],
                            y=monitor["top"],
                            is_primary=(idx == 0),
                            refresh_rate=refresh_rate,
                        )
                    )

            logger.debug(f"MSS detected {len(displays)} display(s)")
            return displays

        except Exception as e:
            logger.error(f"Failed to detect displays with MSS: {e}")
            raise RuntimeError(f"Failed to detect displays: {e}")

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