# -*- coding: utf-8 -*-
"""屏幕捕获：为「实时屏幕镜像」和「单次截屏」提供 JPEG 帧。

只使用 Pillow（已随程序打包），不引入额外依赖。
"""
from __future__ import annotations

import io
import sys
import threading
import time

IS_WINDOWS = sys.platform.startswith("win")


class ScreenError(RuntimeError):
    pass


def _grab_error_text(error: Exception) -> str:
    """把抓屏失败翻译成人话。

    锁屏是最常见的失败原因（离开电脑就会锁），
    而且 Windows 锁屏时**不给抓屏**，报的是 OSError: screen grab failed。
    直接透传这句用户看不懂，所以这里给出可操作的提示。
    """
    text = str(error).lower()
    if "screen grab failed" in text or "grab" in text:
        return ("抓屏失败：电脑现在锁屏了（或没有桌面会话）。"
                "请先解锁屏幕再试。")
    return f"抓屏失败: {error}"


class ScreenCapturer:
    """抓屏并编码为 JPEG。线程安全（内部串行化）。"""

    def __init__(self, max_width: int = 960, quality: int = 72, monitor: int = 1,
                 dedup: bool = True):
        self.max_width = int(max_width)
        self.quality = int(quality)
        self.monitor = int(monitor)
        # 画面没变就不重复编码（见 grab_jpeg）
        self.dedup = bool(dedup)
        self._lock = threading.RLock()
        self._last_size = (0, 0)
        self._last_digest = None

    # ------------------------------------------------------------------

    def available(self) -> bool:
        try:
            import importlib
            importlib.import_module("PIL.ImageGrab")
            return True
        except Exception:
            return False

    def grab_jpeg(self, force: bool = False) -> tuple[bytes, int, int]:
        """抓一帧，返回 (jpeg 字节, 宽, 高)。"""
        if not IS_WINDOWS:
            # macOS 也可用 ImageGrab；其他平台给出明确错误
            pass
        try:
            from PIL import Image, ImageGrab
        except ImportError as exc:                        # pragma: no cover
            raise ScreenError("缺少 Pillow，无法抓屏") from exc

        with self._lock:
            try:
                image = ImageGrab.grab()
            except Exception as exc:                      # noqa: BLE001
                raise ScreenError(_grab_error_text(exc)) from exc

            if image.mode not in ("RGB", "L"):
                image = image.convert("RGB")

            width, height = image.size
            if self.max_width and width > self.max_width:
                ratio = self.max_width / float(width)
                new_size = (self.max_width, max(1, int(height * ratio)))
                image = image.resize(new_size, Image.BILINEAR)
                width, height = image.size

            # 画面没变就不编码。
            # 用 32x18 的小缩略图做指纹：算一次只要 0.1ms 左右，
            # 但能省掉 17ms 的 JPEG 编码和之后的传输。
            # 手表看静态桌面是最常见的情况，这一条省得最多。
            digest = None
            if self.dedup and not force:
                try:
                    thumb = image.resize((32, 18), Image.BILINEAR)
                    digest = hash(thumb.tobytes())
                except Exception:                         # noqa: BLE001
                    digest = None
                if digest is not None and digest == self._last_digest:
                    self._last_size = (width, height)
                    return b"", width, height
                self._last_digest = digest

            buffer = io.BytesIO()
            try:
                image.save(buffer, format="JPEG", quality=self.quality,
                           optimize=False, subsampling=2)
            except Exception as exc:                      # noqa: BLE001
                raise ScreenError(f"JPEG 编码失败: {exc}") from exc

            self._last_size = (width, height)
            return buffer.getvalue(), width, height

    @property
    def last_size(self) -> tuple[int, int]:
        return self._last_size

    def screen_size(self) -> tuple[int, int]:
        """原始屏幕分辨率（用于界面显示）。"""
        try:
            from PIL import ImageGrab
            image = ImageGrab.grab()
            return image.size
        except Exception:
            return (0, 0)


class FrameStreamer:
    """按固定帧率在后台线程推帧，供服务端广播。

    画面静止时不会傻发同样的画面（用 32x18 缩略图指纹判断），
    但每隔 HEARTBEAT 秒会强制发一帧，让客户端知道流还活着。
    """

    # 画面没变时，最多隔这么久也要发一帧（心跳）
    HEARTBEAT = 1.5

    def __init__(self, capturer: ScreenCapturer, on_frame, on_stop=None):
        self.capturer = capturer
        self.on_frame = on_frame            # on_frame(jpeg: bytes, w: int, h: int)
        self.on_stop = on_stop
        self.fps = 3.0
        self._thread: threading.Thread | None = None
        self._stop = threading.Event()
        self.frames_sent = 0
        self.last_error = ""
        self._started_at = 0.0
        # 上次真正发出帧的时间（心跳判断用）
        self._last_sent = 0.0

    @property
    def running(self) -> bool:
        return self._thread is not None and self._thread.is_alive()

    def start(self, fps: float = 3.0) -> None:
        if self.running:
            self.fps = max(0.5, min(15.0, float(fps)))
            return
        self.fps = max(0.5, min(30.0, float(fps)))
        self._stop.clear()
        self.frames_sent = 0
        self.last_error = ""
        self._started_at = time.time()
        self._thread = threading.Thread(target=self._loop, name="screen-stream",
                                        daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        thread = self._thread
        if thread is not None and thread.is_alive():
            thread.join(timeout=3.0)
        self._thread = None
        if self.on_stop:
            try:
                self.on_stop()
            except Exception:
                pass

    def _loop(self) -> None:
        interval = 1.0 / self.fps
        # 画面静止时把采样间隔逐级拉长（最多 4 倍），省 CPU。
        # 抓屏本身就要 30ms 左右，静止时没必要一直满速抓。
        idle = 1.0
        while not self._stop.is_set():
            started = time.time()
            try:
                jpeg, width, height = self.capturer.grab_jpeg()
            except ScreenError as exc:
                self.last_error = str(exc)
                time.sleep(0.5)
                continue

            if not jpeg:
                # 画面没变（capturer 做了去重）。
                # 但不能一直不发 —— 客户端会以为流断了。
                # 所以每隔 HEARTBEAT 秒强制发一帧当心跳。
                now = time.time()
                if now - self._last_sent >= self.HEARTBEAT:
                    try:
                        jpeg, width, height = self.capturer.grab_jpeg(
                            force=True)
                    except ScreenError as exc:
                        self.last_error = str(exc)
                        time.sleep(0.5)
                        continue
                    if jpeg:
                        self.frames_sent += 1
                        self._last_sent = now
                        idle = 1.0
                        try:
                            if self.on_frame(jpeg, width, height) is False:
                                break
                        except Exception as exc:          # noqa: BLE001
                            self.last_error = str(exc)
                            break
                        elapsed = time.time() - started
                        wait = interval - elapsed
                        if wait > 0:
                            self._stop.wait(wait)
                        continue
                idle = min(4.0, idle * 1.6)
                self._stop.wait(interval * idle)
                continue
            self._last_sent = time.time()

            idle = 1.0
            self.frames_sent += 1
            try:
                if self.on_frame(jpeg, width, height) is False:
                    break                        # 没有观看者了，自动结束
            except Exception as exc:              # noqa: BLE001
                self.last_error = str(exc)
                break
            elapsed = time.time() - started
            wait = interval - elapsed
            if wait > 0:
                self._stop.wait(wait)

    @property
    def uptime(self) -> float:
        return time.time() - self._started_at if self._started_at else 0.0

    @property
    def actual_fps(self) -> float:
        uptime = self.uptime
        return (self.frames_sent / uptime) if uptime > 0.2 else 0.0