import asyncio as aio
import multiprocessing as mp
import queue
from typing import Any, Optional

import bson
import websockets
from pydantic import BaseModel

from ldnode import logger
from ldnode.calculator.handlers import handler_registry
from ldnode.runner.config import Settings
from ldnode.runner.events import Events
from ldnode.runner.schemas import Worker
from ldnode.runner.task_manager import TaskManager

from .schemas.protocol import (
    EndpointAction,
    FinishOrCancelTaskModel,
    LockResourceModel,
    Method,
    NodeAction,
    OnlineModel,
    PingModel,
    PushTaskModel,
    ReadLogModel,
    ReportDatasetStatusModel,
    SyncDataLogModel,
    UpgradeModel,
)
from .schemas.transfer import Decoder, Encoder
from .utils import codec_options


class WSWorker(Worker):
    """WebSocket 工作进程"""

    CONNECT_URI: str = "/v1/calculation/node-connections"

    def __init__(self, worker_name: str):
        super().__init__(worker_name)
        self.settings: Settings | None = None
        self.inq: mp.Queue | None = None
        self.outq: mp.Queue | None = None
        self.tasks: TaskManager = TaskManager()
        self.encoder: Encoder = Encoder()
        self.decoder: Decoder = Decoder()
        self.websocket: websockets.ClientConnection | None = None
        self.waiting_send_chunks: list[bytes] = []

    async def submit_event(self, source: str, event: str, data: Any):
        """提交事件"""
        await aio.to_thread(self.outq.put, (source, event, data))

    async def handle_event(self, action: NodeAction, data: Any):
        """处理事件"""
        match action:
            case NodeAction.PING:
                await self.send_msg(NodeAction.PING, PingModel(**data))
            case NodeAction.SYNC_DATA_LOG:
                await self.send_msg(NodeAction.SYNC_DATA_LOG, SyncDataLogModel(**data))
            case NodeAction.FINISHED_TASK:
                await self.send_msg(
                    NodeAction.FINISHED_TASK, FinishOrCancelTaskModel(**data)
                )
            case NodeAction.CANCELLED_TASK:
                await self.send_msg(
                    NodeAction.CANCELLED_TASK, FinishOrCancelTaskModel(**data)
                )
            case NodeAction.REPORT_DATASET_STATUS:
                await self.send_msg(
                    NodeAction.REPORT_DATASET_STATUS, ReportDatasetStatusModel(**data)
                )

    async def route_message_to_event(
        self, action: EndpointAction, msg: Optional[Any] = None
    ):
        """路由消息到事件"""
        match action:
            case EndpointAction.PUSH_TASK:
                task_msg: PushTaskModel = msg
                logger.connector.info(
                    f"收到计算任务: {task_msg.method}({task_msg.task_id})"
                )
                await self.submit_event(
                    "Main", Events.TASK_SUBMITTED, task_msg.model_dump()
                )
            case EndpointAction.RESTART:
                logger.connector.info("收到重启请求")
                await self.submit_event("Main", Events.RESTART, None)
            case EndpointAction.UPGRADE:
                upgrade_msg: UpgradeModel = msg
                logger.connector.info(f"收到更新请求: version={upgrade_msg.version}")
                await self.submit_event(
                    "Main", Events.UPGRADE, upgrade_msg.model_dump()
                )
            case EndpointAction.OFFLINE:
                logger.connector.info("收到下线请求")
                await self.submit_event("Main", Events.EXIT, None)
            case EndpointAction.READ_LOGGERS:
                logger.connector.info("收到读取日志器请求")
                await self.submit_event("Main", Events.READ_LOGGERS, None)
            case EndpointAction.READ_LOG:
                read_log_msg: ReadLogModel = msg
                logger.connector.info(
                    f"收到读取日志请求: {read_log_msg.logger}({read_log_msg.offset}, {read_log_msg.limit})"
                )
                await self.submit_event(
                    "Main", Events.READ_LOG, read_log_msg.model_dump()
                )
            case EndpointAction.PONG:
                logger.connector.info("收到心跳数据")
                await self.submit_event("Main", Events.PONG, None)
            case EndpointAction.CANCEL_TASK:
                cancel_task_msg: FinishOrCancelTaskModel = msg
                logger.connector.info(
                    f"收到取消计算任务请求: {cancel_task_msg.method}({cancel_task_msg.task_id})"
                )
                await self.submit_event(
                    "Main",
                    Events.TASK_CANCEL,
                    {
                        "task_id": cancel_task_msg.task_id,
                        "method": cancel_task_msg.method,
                    },
                )

    async def handle_message(self, msg_data: bytes):
        """处理消息"""
        try:
            action = EndpointAction(msg_data[0])
            if action.model:
                msg_dict = bson.decode(msg_data[1:])
                msg = action.model.model_validate(msg_dict)
            else:
                msg = None
            self.tasks.create_task(self.route_message_to_event(action, msg))
        except ValueError:
            logger.connector.warning(f"未知动作: {msg[0]}")
        except Exception as e:
            logger.connector.warning(f"处理消息失败: {e}")

    async def _listen_input(self):
        """监听输入队列"""
        while self.running:
            try:
                data = await aio.to_thread(self.inq.get, timeout=1)
                if data is None:
                    self.running = False
                    break
                action, msg = data
                self.tasks.create_task(self.handle_event(action, msg))
            except queue.Empty:
                continue
            except Exception as e:
                logger.connector.error(f"监听输入队列失败: {e}")
                continue

    async def send_chunk(self, chunk: bytes) -> None:
        """发送块"""
        if self.websocket:
            try:
                await self.websocket.send(chunk)
                # logger.connector.info(f"发送块成功: {humanize.naturalsize(len(chunk))}")
            except websockets.exceptions.ConnectionClosed as e:
                # 连接已关闭，放入等待队列，重连后由 prepare_connection() 统一补发
                logger.connector.error(f"发送块失败（连接已关闭），将在重连后重发: {e}")
                self.waiting_send_chunks.append(chunk)
            except Exception as e:
                # 其他临时异常（如写缓冲满），重新入队重试
                logger.connector.error(f"发送块失败，重试: {e}")
                self.tasks.create_task(self.send_chunk(chunk))
        else:
            self.waiting_send_chunks.append(chunk)

    async def send_msg(self, action: NodeAction, data: BaseModel | None) -> None:
        """发送消息"""
        byte_data = (
            bson.encode(data.model_dump(), codec_options=codec_options) if data else b""
        )
        action_value: int = action.value
        payload = action_value.to_bytes(1, "big") + byte_data
        chunks = self.encoder.encode(payload)
        for chunk in chunks:
            await self.send_chunk(chunk)

    async def prepare_connection(self) -> None:
        """准备连接"""
        # 发送认证消息
        await self.send_msg(
            NodeAction.ONLINE,
            OnlineModel(
                access_id=self.settings.ENDPOINT_ACCESS_ID,
                access_token=self.settings.ENDPOINT_ACCESS_TOKEN,
                platform_tag=self.settings.PLATFORM_TAG,
                node_version=self.settings.NODE_VERSION,
                cpu_resource=self.settings.CPU_RESOURCE,
                gpu_resource=0,
                enabled_methods=[
                    Method(
                        name=handler.name,
                        cpu_usage=handler.cpu_usage,
                        gpu_usage=handler.gpu_usage,
                    )
                    for handler in handler_registry.all()
                ],
            ),
        )
        logger.connector.info(f"发送上线消息: {self.settings.ENDPOINT_ACCESS_ID}")
        await self.submit_event("Main", Events.ONLINE, None)
        # 判断是否还有等待发送的chunk
        if self.waiting_send_chunks:
            # 发送等待发送的chunk
            for chunk in self.waiting_send_chunks:
                await self.send_chunk(chunk)
            # 清空等待发送的chunk
            self.waiting_send_chunks = []

    # 应用层心跳超时：超过此时长未收到服务端 PONG，判定连接僵死并主动断开触发重连
    # 设为心跳间隔（5s）的 6 倍，容忍短暂网络抖动
    PONG_TIMEOUT: float = 30.0

    async def connect_if_disconnected(self) -> None:
        """如果连接断开，则重新连接"""
        connect_uri = self.settings.ENDPOINT_URL + self.CONNECT_URI
        async with websockets.connect(
            connect_uri, ping_interval=None, ping_timeout=None
        ) as websocket:
            logger.connector.info(f"连接成功: {connect_uri}")
            # 创建最新的websocket客户端
            self.websocket = websocket
            # 准备连接
            self.tasks.create_task(self.prepare_connection())
            # 接收数据：加超时检测服务端静默掉线
            # 服务端正常运行时每 5s 回一个 PONG，超过 PONG_TIMEOUT 未收到任何数据
            # 说明连接已僵死（TCP 半开），抛 TimeoutError 触发 start_connection 重连
            while self.running:
                try:
                    data = await aio.wait_for(
                        self.websocket.recv(decode=False),
                        timeout=self.PONG_TIMEOUT,
                    )
                except aio.TimeoutError:
                    logger.connector.warning(
                        f"超过 {self.PONG_TIMEOUT}s 未收到服务端数据，判定连接僵死，主动断开"
                    )
                    raise websockets.exceptions.ConnectionClosedError(None, None)
                # 解码数据
                ok, decoded_data = self.decoder.decode(data)
                if ok:
                    await self.handle_message(decoded_data)
        self.websocket = None

    async def start_connection(self) -> None:
        """启动工作进程"""
        while self.running:
            try:
                await self.connect_if_disconnected()
            except Exception as e:
                self.websocket = None
                await self.submit_event("Main", Events.OFFLINE, None)
                logger.connector.warning(f"websocket 连接异常: {e}")
                await aio.sleep(1)
            if self.running:
                logger.connector.info(
                    f"重连 {self.settings.ENDPOINT_URL + self.CONNECT_URI}..."
                )
        logger.connector.info("连接服务已退出~")

    async def _run(self):
        """运行工作进程"""
        self.tasks.create_task(self.start_connection())
        await self._listen_input()

    def run(self, settings: Settings, input: mp.Queue, output: mp.Queue):
        """运行工作进程"""
        self.settings = settings
        self.inq = input
        self.outq = output
        aio.run(self._run())
