"""截图与 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)

    提供与 PIL.Image.Image 兼容的常用属性和方法:
      - .size -> (width, height)
      - .width / .height
      - .mode -> 'RGB'
      - .save(path) -> 保存到文件
      - bytes(img) / len(img) -> 原始 JPEG 数据
    """

    __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]:
        """从 JPEG 数据解析图片尺寸"""
        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]
            # SOF markers (SOF0-SOF15, except DHT/DAC)
            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:  # EOI
                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"]:
        """截图

        Args:
            filename: 保存路径，为 None 则仅返回图像对象
            quality: JPEG 质量 (1-100)

        Returns:
            PIL.Image.Image 对象（已安装 Pillow 时）
            Screenshot 兼容对象（未安装 Pillow 时，支持 .size/.save() 等）
        """
        data = self._client.get_bytes(f"/screen/{quality}")
        if filename:
            with open(filename, 'wb') as f:
                f.write(data)
        try:
            from PIL import Image
            return Image.open(io.BytesIO(data))
        except ImportError:
            return Screenshot(data)

    def screenshot_bytes(self, quality: int = 80) -> bytes:
        """截图，返回原始 JPEG 字节数据（无需 Pillow）

        Args:
            quality: JPEG 质量 (1-100)

        Returns:
            bytes JPEG 图片数据
        """
        return self._client.get_bytes(f"/screen/{quality}")

    def dump_hierarchy(self, compressed: bool = False) -> str:
        """获取 UI 控件树 XML

        Args:
            compressed: 是否压缩（去除空白节点），当前不影响结果

        Returns:
            XML 字符串
        """
        return self._client.get_text("/uia_dump")

    def window_size(self) -> Tuple[int, int]:
        """获取屏幕尺寸 (width, height)"""
        result = self._client.auto_api("/screen_size")
        if isinstance(result, dict):
            return (int(result.get("width", 0)), int(result.get("height", 0)))
        if isinstance(result, str) and "x" in result:
            parts = result.split("x")
            return (int(parts[0].strip()), int(parts[1].strip()))
        return (0, 0)

    @property
    def orientation(self) -> int:
        """获取屏幕方向 (0=竖屏, 1=横屏左, 2=倒屏, 3=横屏右)"""
        result = self._client.auto_api("/screen_rotation")
        return int(result) if result else 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 提示

        Args:
            text: 提示文字
        """
        self._client.auto_api("/toast", {"content": text})
