"""RFB 协议底层处理类"""

import asyncio
import logging
import struct
from typing import Optional
import aiohttp
from PIL import Image

from .types import ServerInfo
from .exceptions import (
    ConnectionError,
    HandshakeError,
    AuthenticationError,
    ProtocolError,
    FramebufferError,
    ScreenshotError,
)
from .protocol import (
    RFBProtocol,
    parse_server_version,
    parse_security_types,
    parse_server_init,
    parse_framebuffer_update_header,
    create_framebuffer_update_request,
)

logger = logging.getLogger(__name__)


class RFBHandler:
    """RFB 协议底层处理器

    负责处理 RFB/VNC 协议相关的底层操作，包括：
    - WebSocket 通信
    - RFB 协议握手
    - 帧缓冲区管理
    - 数据收发
    """

    def __init__(self, timeout: float = 10.0):
        """初始化处理器

        Args:
            timeout: 操作超时时间（秒）
        """
        self.timeout = timeout

        # 连接相关
        self._session: Optional[aiohttp.ClientSession] = None
        self._websocket: Optional[aiohttp.ClientWebSocketResponse] = None

        # 服务器信息
        self._server_info: Optional[ServerInfo] = None
        self._framebuffer: Optional[bytearray] = None

    @property
    def server_info(self) -> Optional[ServerInfo]:
        """服务器信息"""
        return self._server_info

    @property
    def framebuffer(self) -> Optional[bytearray]:
        """帧缓冲区"""
        return self._framebuffer

    async def connect(self, host: str, port: int, path: str) -> None:
        """建立 WebSocket 连接

        Args:
            host: 服务器地址
            port: 服务器端口
            path: WebSocket 路径
        """
        self._session = aiohttp.ClientSession()
        url = f"ws://{host}:{port}{path}"

        logger.debug(f"正在连接到 {url}")
        self._websocket = await self._session.ws_connect(
            url,
            timeout=aiohttp.ClientWSTimeout(ws_receive=self.timeout),  # type: ignore
        )

    async def disconnect(self) -> None:
        """断开连接并清理资源"""
        if self._websocket:
            await self._websocket.close()
        if self._session:
            await self._session.close()
        self._websocket = None
        self._session = None

    async def send_binary(self, data: bytes) -> None:
        """发送二进制数据

        Args:
            data: 要发送的数据
        """
        if not self._websocket:
            raise ConnectionError("WebSocket 未连接")
        await self._websocket.send_bytes(data)
        logger.debug(f"发送数据: {len(data)} 字节")

    async def receive_binary(self) -> bytes:
        """接收二进制数据

        Returns:
            接收到的数据
        """
        if not self._websocket:
            raise ConnectionError("WebSocket 未连接")

        try:
            msg = await asyncio.wait_for(
                self._websocket.receive(), timeout=self.timeout
            )

            if msg.type == aiohttp.WSMsgType.BINARY:
                return msg.data
            elif msg.type == aiohttp.WSMsgType.TEXT:
                return msg.data.encode("utf-8")
            elif msg.type == aiohttp.WSMsgType.ERROR:
                raise ConnectionError(f"WebSocket 错误: {self._websocket.exception()}")
            elif msg.type == aiohttp.WSMsgType.CLOSE:
                raise ConnectionError("WebSocket 连接已关闭")
            else:
                raise ConnectionError(f"未知消息类型: {msg.type}")

        except asyncio.TimeoutError:
            raise ConnectionError("接收数据超时")

    async def perform_handshake(self) -> None:
        """执行 RFB 协议握手"""
        try:
            # 1. 协议版本协商
            version_data = await self.receive_binary()
            server_version = parse_server_version(version_data)
            logger.info(f"服务器版本: {server_version}")

            await self.send_binary(RFBProtocol.VERSION)

            # 2. 安全类型协商
            security_data = await self.receive_binary()
            security_types = parse_security_types(security_data)
            logger.info(f"可用安全类型: {security_types}")

            if RFBProtocol.SECURITY_NONE not in security_types:
                raise AuthenticationError(f"服务器不支持无认证模式: {security_types}")

            await self.send_binary(bytes([RFBProtocol.SECURITY_NONE]))

            # 3. 安全结果
            result_data = await self.receive_binary()
            if len(result_data) < 4:
                raise HandshakeError("安全结果数据不足")

            result = struct.unpack(">I", result_data[:4])[0]
            if result != 0:
                raise AuthenticationError(f"认证失败: {result}")

            # 4. 客户端初始化（共享桌面）
            await self.send_binary(bytes([1]))

            # 5. 服务器初始化
            server_init_data = await self.receive_binary()
            self._server_info = parse_server_init(server_init_data)

            logger.info(
                f"服务器初始化完成: {self._server_info.width}x{self._server_info.height}"
            )

        except Exception as e:
            if isinstance(e, (HandshakeError, AuthenticationError, ProtocolError)):
                raise
            raise HandshakeError(f"握手失败: {e}") from e

    async def receive_framebuffer_update(self) -> None:
        """接收帧缓冲区更新"""
        if not self._server_info:
            raise FramebufferError("服务器信息未初始化")

        try:
            # 接收更新消息
            data = await self.receive_binary()
            if not data or data[0] != RFBProtocol.SERVER_FRAMEBUFFER_UPDATE:
                raise FramebufferError("不是帧缓冲区更新消息")

            rect_count, rectangles = parse_framebuffer_update_header(data)

            if not rectangles:
                # 如果没有完整的矩形头部，尝试继续接收
                logger.debug("矩形头部不完整，尝试继续接收数据")
                additional_data = await self.receive_binary()
                data += additional_data
                rect_count, rectangles = parse_framebuffer_update_header(data)

            if not rectangles:
                raise FramebufferError("没有矩形数据")

            # 处理第一个矩形
            rect = rectangles[0]
            if rect.encoding != RFBProtocol.ENCODING_RAW:
                raise FramebufferError(f"不支持的编码: {rect.encoding}")

            # 接收像素数据，包含已经在data中的部分
            pixels_data = await self._receive_pixel_data(rect, data[16:])  # 跳过头部

            # 更新帧缓冲区
            self._update_framebuffer(rect, pixels_data)

        except Exception as e:
            if isinstance(e, FramebufferError):
                raise
            raise FramebufferError(f"接收帧缓冲区更新失败: {e}") from e

    async def _receive_pixel_data(self, rect, initial_data: bytes = b"") -> bytes:
        """接收像素数据"""
        bytes_per_pixel = 4  # 32位 BGRA
        expected_size = rect.width * rect.height * bytes_per_pixel

        pixel_data = initial_data
        attempts = 0
        max_attempts = 200  # 增加尝试次数

        while len(pixel_data) < expected_size and attempts < max_attempts:
            try:
                additional = await self.receive_binary()
                if additional:
                    pixel_data += additional
                    logger.debug(
                        f"接收像素数据: {len(pixel_data)}/{expected_size} 字节"
                    )
                else:
                    await asyncio.sleep(0.01)
                attempts += 1
            except Exception:
                break

        if len(pixel_data) < expected_size:
            logger.warning(
                f"像素数据不足: {len(pixel_data)}/{expected_size}，使用可用数据"
            )
            # 不抛出异常，而是使用可用的数据
            return pixel_data

        return pixel_data[:expected_size]

    def _update_framebuffer(self, rect, pixel_data: bytes) -> None:
        """更新帧缓冲区"""
        if not self._server_info:
            return

        if self._framebuffer is None:
            total_size = self._server_info.width * self._server_info.height * 4
            self._framebuffer = bytearray(total_size)

        # 将像素数据复制到帧缓冲区
        for row in range(rect.height):
            src_offset = row * rect.width * 4
            dst_offset = ((rect.y + row) * self._server_info.width + rect.x) * 4
            end_offset = dst_offset + rect.width * 4

            self._framebuffer[dst_offset:end_offset] = pixel_data[
                src_offset : src_offset + rect.width * 4
            ]

    def framebuffer_to_image(self) -> Image.Image:
        """将帧缓冲区转换为 PIL 图像

        Returns:
            PIL 图像对象
        """
        if not self._framebuffer or not self._server_info:
            raise ScreenshotError("帧缓冲区或服务器信息为空")

        # 转换 BGRA 为 RGB
        rgb_data = bytearray()
        for i in range(0, len(self._framebuffer), 4):
            if i + 3 < len(self._framebuffer):
                b = self._framebuffer[i]
                g = self._framebuffer[i + 1]
                r = self._framebuffer[i + 2]
                # 跳过 alpha 通道
                rgb_data.extend([r, g, b])

        # 创建 PIL 图像
        image = Image.frombytes(
            "RGB", (self._server_info.width, self._server_info.height), bytes(rgb_data)
        )

        return image

    async def request_screen_update(self, incremental: bool = False) -> None:
        """请求屏幕更新

        Args:
            incremental: 是否为增量更新
        """
        if not self._server_info:
            raise ProtocolError("服务器信息未初始化")

        message = create_framebuffer_update_request(
            incremental=incremental,
            x=0,
            y=0,
            width=self._server_info.width,
            height=self._server_info.height,
        )
        await self.send_binary(message)
        logger.debug(f"请求屏幕更新: 增量={incremental}")
