"""截图与 UI 树 Mixin"""

from __future__ import annotations

import io
import struct
from typing import TYPE_CHECKING, Optional, Tuple, Union

if TYPE_CHECKING:
    from PIL import Image as PILImage
    from ..device import Device


class Screenshot:
    """轻量截图对象（无需 Pillow）"""

    __slots__ = ('_data', '_width', '_height')

    def __init__(self, data: bytes):
        self._data = data
        self._width, self._height = self._parse_jpeg_size(data)

    @staticmethod
    def _parse_jpeg_size(data: bytes) -> Tuple[int, int]:
        if len(data) < 2 or data[0:2] != b'\xff\xd8':
            return (0, 0)
        i = 2
        while i < len(data) - 1:
            if data[i] != 0xFF:
                i += 1
                continue
            marker = data[i + 1]
            if marker in (0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7,
                          0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF):
                if i + 9 < len(data):
                    height = struct.unpack('>H', data[i+5:i+7])[0]
                    width = struct.unpack('>H', data[i+7:i+9])[0]
                    return (width, height)
                return (0, 0)
            elif marker == 0xD9:
                return (0, 0)
            elif marker in (0xD0, 0xD1, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0x01):
                i += 2
            else:
                if i + 3 < len(data):
                    seg_len = struct.unpack('>H', data[i+2:i+4])[0]
                    i += 2 + seg_len
                else:
                    return (0, 0)
        return (0, 0)

    @property
    def size(self) -> Tuple[int, int]:
        return (self._width, self._height)

    @property
    def width(self) -> int:
        return self._width

    @property
    def height(self) -> int:
        return self._height

    @property
    def mode(self) -> str:
        return 'RGB'

    def save(self, fp, format=None, **kwargs):
        if isinstance(fp, str):
            with open(fp, 'wb') as f:
                f.write(self._data)
        else:
            fp.write(self._data)

    def tobytes(self) -> bytes:
        return self._data

    def __bytes__(self) -> bytes:
        return self._data

    def __len__(self) -> int:
        return len(self._data)

    def __repr__(self) -> str:
        return f'<Screenshot {self._width}x{self._height} ({len(self._data)} bytes)>'


class ScreenMixin:
    """截图与屏幕操作"""

    _client: ...

    def screenshot(self, filename: Optional[str] = None, quality: int = 80) -> Union["PILImage.Image", "Screenshot"]:
        """截图"""
        data = self._client.get_bytes(f"/screen/{quality}")
        if filename:
            with open(filename, 'wb') as f:
                f.write(data)
        try:
            from PIL import Image
            try:
                return Image.open(io.BytesIO(data))
            except Exception:
                # PIL 无法识别时降级到轻量 Screenshot 实现
                return Screenshot(data)
        except ImportError:
            return Screenshot(data)

    def screenshot_bytes(self, quality: int = 80) -> bytes:
        """截图，返回原始 JPEG 字节"""
        return self._client.get_bytes(f"/screen/{quality}")

    def dump_hierarchy(self, compressed: bool = False) -> str:
        """获取 UI 控件树 XML"""
        return self._client.get_text("/uia_dump")

    def window_size(self) -> Tuple[int, int]:
        """获取屏幕尺寸 (width, height)"""
        result = self._client.auto_api("/screen_size")
        # 解包 {'data': ..., 'ok': True} 格式
        if isinstance(result, dict) and "data" in result:
            result = result["data"]
        if isinstance(result, dict):
            return (int(result.get("width", 0)), int(result.get("height", 0)))
        if isinstance(result, str):
            # "1080x2340" 格式
            s = result.strip()
            if "x" in s:
                parts = s.split("x")
                try:
                    return (int(parts[0].strip()), int(parts[1].strip()))
                except (ValueError, IndexError):
                    pass
            # "1080 2340" 格式
            parts = s.split()
            if len(parts) == 2:
                try:
                    return (int(parts[0]), int(parts[1]))
                except ValueError:
                    pass
        return (0, 0)

    @property
    def orientation(self) -> int:
        """获取屏幕方向 (0-3)"""
        try:
            result = self._client.auto_api("/screen_rotation")
            if isinstance(result, dict) and "data" in result:
                result = result["data"]
            return int(result) if result is not None else 0
        except Exception:
            return 0

    def screen_on(self) -> None:
        """点亮屏幕"""
        self._client.auto_api("/wake_up")

    def screen_off(self) -> None:
        """熄灭屏幕"""
        self._client.auto_api("/sleep")

    def freeze_rotation(self, freeze: bool = True) -> None:
        """冻结/解冻屏幕旋转"""
        self._client.auto_api("/freeze_rotation", {"freeze": str(freeze).lower()})

    def toast(self, text: str) -> None:
        """在设备屏幕上显示 Toast 提示"""
        self._client.auto_api("/toast", {"content": text})

    def wait_screen_change(self, timeout: float = 10, interval: float = 0.3, 
                           threshold: float = 0.95) -> bool:
        """等待屏幕发生变化"""
        import time
        deadline = time.time() + timeout
        # 这里实际上可能需要引擎支持比较，或者在 Python 端简单实现
        # 但既然文档里有，我们应该调用引擎 API (如果引擎有对应 endpoint)
        # 假设引擎有 /wait_screen_change
        params = {"timeout": timeout, "interval": interval, "threshold": threshold}
        result = self._client.auto_api("/wait_screen_change", params)
        if isinstance(result, bool): return result
        return str(result).lower() == "true"

    def wait_screen_stable(self, stable_duration: float = 1.0, timeout: float = 10, 
                           interval: float = 0.3, threshold: float = 0.98) -> bool:
        """等待屏幕稳定"""
        params = {"stable_duration": stable_duration, "timeout": timeout, 
                  "interval": interval, "threshold": threshold}
        result = self._client.auto_api("/wait_screen_stable", params)
        if isinstance(result, bool): return result
        return str(result).lower() == "true"
