import asyncio
import logging
import threading
import time
from urllib.parse import urlencode

import socketio
from nanoid import generate

from ..config import settings
from ..types import APIFuture, PendingRequest

logger = logging.getLogger(__name__)


class SocketClient:
    def __init__(
        self,
        run_id: str,
        user_id: str,
        username: str,
    ):
        """
        初始化 SocketIO 客户端

        Parameters
        ----------
        run_id: 运行 ID，随机生成
        user_id: 用户唯一 ID
        username: 用户名
        """
        self.run_id = run_id
        self.user_id = user_id
        self.username = username

        self.sio = socketio.Client()
        self.timeout = settings.ws_timeout

        self._pending: dict[str, PendingRequest] = {}
        self._pending_lock = threading.Lock()

        self._setup_event()

    def connect(self):
        params = {
            "runId": self.run_id,
            "apiKey": settings.api_key,
            "sdkVersion": settings.version,
        }
        query_string = urlencode(params)
        self.sio.connect(f"{settings.ws_url}?{query_string}")

    def _setup_event(self):
        """
        初始化监听事件，包括如下事件：
        - connect: 连接成功
        - disconnect: 连接断开
        - connect_error: 连接错误
        - system: 系统通知消息
        - request: 请求响应
        """

        def connect():
            """
            连接到服务器
            """
            logger.info("🚀 连接到服务器，已准备就绪")

        self.sio.on("connect", connect)

        def disconnect():
            """
            与服务器连接断开
            """
            logger.warning("😳 与服务器连接断开，任务停止；请检查你的网络")

        self.sio.on("disconnect", disconnect)

        def connect_error(data):
            """
            遇到服务器连接错误
            """
            logger.error(f"❌ 连接错误: {data}")

        self.sio.on("connect_error", connect_error)

        def handle_system_message(data):
            """
            接收系统响应信息
            """
            logger.debug(f"🫣 当前系统信息: {data}")

        self.sio.on("system", handle_system_message)

        def on_submitted(data):
            """
            接收提交事件的响应结果
            """
            event_id = data.get("id")

            if event_id is None:
                logger.error("❌ 收到的事件数据缺少 id 字段")
                return

            with self._pending_lock:
                pending = self._pending.get(event_id)

            if pending is None:
                logger.error(f"❌ 收到未知事件 ID 的响应: {event_id}")
                return

            pending.submitted_payload = data
            pending.submitted_event.set()
            logger.debug(f"✅ 事件 {event_id} 已提交")

        self.sio.on("submitted", on_submitted)

        def on_request(data):
            """
            接收请求后的响应结果
            """
            # 保存接收到的数据
            event_id = data.get("id")

            if event_id is None:
                logger.error("❌ 收到的响应数据缺少 id 字段")
                return

            with self._pending_lock:
                pending = self._pending.pop(event_id, None)

            if pending is None:
                logger.error(f"❌ 收到未知事件 ID 的响应: {event_id}")
                return

            if not pending.result_future.done():
                pending.result_future.set_result(data)
                logger.debug(f"✅ 事件 {event_id} 已完成")

        self.sio.on("request", on_request)

    def disconnect(self):
        """
        主动断开连接
        """
        self.sio.disconnect()

    def submit_event(
        self,
        task_type: str,
        task_id: str,
        runner_id: str,
        event_type: str,
        payload: dict,
        index: int = 0,
        step: int = 0,
        input_tokens: int | None = 0,
    ):
        """
        发送事件

        {
          "id": "271chndjSQna", // 事件唯一 id，12 位随机字符，由 sdk 随机生成
          "type": "forward", // 事件类型
          "metadata": {
            "task_type": "TRAIN", // 任务类型
            "user_id": "pXJFEM", // 用户 ID，6 位随机字符，由 sdk 生成
            "username": "jackal", // 用户名，由用户定义，sdk 存储
            "task_id": "pfis7LIqmP", // 任务 ID，10 位随机字符，由 sdk 随机生成
            "run_id": "IsE90fZsKoqipqyzAAAB", // 运行 ID，16 位随机字符，由 sdk 随机生成
            "runner_id": "d9pmrwsle3", // 运行器 ID
            "index": 1, // 本次任务中当前事件对应的索引，从 0 开始累加
            "step": 1, // 训练步数
            "version": "1.0", // sdk 的版本号
            "timestamp": 1765261896, // 事件产生的时间戳
            "input_tokens": 0, // 输入 token 量
          },
          "payload": {} // 携带的输入负载，不同事件对应的负载不同
        }

        Parameters
        ----------
        task_type: 任务类型，可选 TRAIN | SAMPLE
        task_id: 任务 id，唯一标识
        runner_id: 分配的运行器 id
        event_type: 事件类型
        payload: 事件的具体内容
        index: 当前的索引
        step: 当前对应的训练步数
        input_tokens: 输入 token 数量

        Returns
        -------
        返回 Future 对象
        """
        event_id = generate(size=12)
        pending = PendingRequest(
            event_id=event_id,
        )

        with self._pending_lock:
            self._pending[event_id] = pending

        try:
            self.sio.emit(
                "request",
                {
                    "id": event_id,
                    "type": event_type,
                    "metadata": {
                        "task_type": task_type,
                        "user_id": self.user_id,
                        "username": self.username,
                        "task_id": task_id,
                        "run_id": self.run_id,
                        "runner_id": runner_id,
                        "index": index,
                        "step": step,
                        "version": settings.version,
                        "timestamp": int(time.time()),
                        "input_tokens": input_tokens,
                    },
                    "payload": payload,
                },
            )
        except Exception as exc:
            with self._pending_lock:
                self._pending.pop(event_id, None)
            raise exc

        return event_id, pending

    def wait_submitted(
        self,
        pending: PendingRequest,
        timeout: float | None = None,
    ) -> bool:
        return pending.submitted_event.wait(timeout=timeout)

    def wait_result(
        self,
        pending: PendingRequest,
        timeout: float | None = None,
    ) -> dict | None:
        return pending.result_future.result(timeout=timeout)

    def clean_up_pending_by_id(
        self,
        event_id: str,
        exc: Exception | None = None,
    ):
        with self._pending_lock:
            pending = self._pending.pop(event_id, None)

        if pending is None:
            return

        if exc is not None and not pending.result_future.done():
            pending.result_future.set_exception(exc)

    def for_task(self, task_type: str, task_id: str, runner_id: str):
        """
        为任务创建一个实例对象
        """
        return ForTask(self, task_type, task_id, runner_id)


class ForTask:
    def __init__(
        self,
        client: SocketClient,
        task_type: str,
        task_id: str,
        runner_id: str,
    ):
        """
        初始化一个任务对象，用于记录任务类型和任务 id

        Parameters
        ----------
        client: SocketClient 实例
        task_type: 任务类型，可选 TRAIN | SAMPLE
        task_id: 任务唯一 id
        runner_id: 运行器 id
        """
        self.task_type = task_type
        self.task_id = task_id
        self.client = client
        self.runner_id = runner_id

        # 初始化索引，从 0 开始累加
        self.index = 0
        self.step = 0

    def _prepare_event(
        self, event_type: str, payload: dict, step: int | None, input_tokens: int
    ):
        """处理发送前的数据准备和索引记录"""
        if step is None:
            step = self.step
        logger.debug(f"👾 发起事件请求 {event_type}, 当前索引为 {self.index}")
        return self.client.submit_event(
            self.task_type,
            self.task_id,
            self.runner_id,
            event_type,
            payload,
            index=self.index,
            step=step,
            input_tokens=input_tokens,
        )

    def _post_event_update(self, event_type: str, event_id: str, pending, ok: bool):
        """处理发送后的异常捕获和状态更新"""
        if not ok:
            self.client.clean_up_pending_by_id(
                event_id, exc=TimeoutError("提交事件超时")
            )
            raise TimeoutError("提交事件超时")

        self.index += 1
        if event_type == "forward-backward":
            self.step += 1
        return APIFuture(pending.result_future)

    def send_event(
        self, event_type: str, payload: dict, step=None, input_tokens: int = 0
    ) -> APIFuture:
        event_id, pending = self._prepare_event(event_type, payload, step, input_tokens)
        ok = self.client.wait_submitted(pending, timeout=self.client.timeout)
        return self._post_event_update(event_type, event_id, pending, ok)

    async def send_event_async(
        self,
        event_type: str,
        payload: dict,
        step: int | None = None,
        input_tokens: int = 0,
    ) -> APIFuture:
        event_id, pending = self._prepare_event(event_type, payload, step, input_tokens)
        loop = asyncio.get_running_loop()
        ok = await loop.run_in_executor(
            None, self.client.wait_submitted, pending, self.client.timeout
        )
        return self._post_event_update(event_type, event_id, pending, ok)
