import asyncio as aio
import gzip
import json
import multiprocessing as mp
import queue

import aiohttp
import bson
import humanize

from ldnode import logger
from ldnode.loader.persistence.light_cos import LightCOSClient
from ldnode.runner.config import Settings
from ldnode.runner.events import TaskAnswer
from ldnode.runner.schemas import Worker
from ldnode.runner.task_manager import TaskManager

from .utils import codec_options


class TaskAnswerCompletionLockState:
    """任务应答完成锁"""

    def __init__(self):
        self.waiting_answers: dict[str, int] = {}
        self.completed_answers: dict[str, tuple[TaskAnswer, aiohttp.FormData, int]] = {}
        self.locks: dict[str, aio.Lock] = {}

    async def pre_register(self, task_id: str) -> None:
        """非 completed 帧在 COS 上传前预先占位，防止 completed 帧在 COS 上传期间抢先发送"""
        if task_id not in self.locks:
            self.locks[task_id] = aio.Lock()
        async with self.locks[task_id]:
            self.waiting_answers[task_id] = self.waiting_answers.get(task_id, 0) + 1

    async def is_completed(
        self, answer: TaskAnswer, form_data: aiohttp.FormData, payload_size: int
    ) -> bool:
        # pre_register 已在 COS 上传前完成计数，此处非 completed 帧直接发送
        if answer["task_id"] not in self.locks:
            self.locks[answer["task_id"]] = aio.Lock()
        async with self.locks[answer["task_id"]]:
            if answer["completed"]:
                if self.waiting_answers.get(answer["task_id"], 0) <= 0:
                    return True
                else:
                    self.completed_answers[answer["task_id"]] = (
                        answer,
                        form_data,
                        payload_size,
                    )
                    return False
            else:
                # waiting_answers 已由 pre_register 计数，无需再次递增
                return True

    async def complete(
        self, answer: TaskAnswer
    ) -> tuple[TaskAnswer, aiohttp.FormData, int] | None:
        if answer["task_id"] not in self.locks:
            self.locks[answer["task_id"]] = aio.Lock()
        async with self.locks[answer["task_id"]]:
            if not answer["completed"]:
                if answer["task_id"] in self.waiting_answers:
                    self.waiting_answers[answer["task_id"]] -= 1
                if self.waiting_answers[answer["task_id"]] <= 0:
                    del self.waiting_answers[answer["task_id"]]
                    if answer["task_id"] in self.completed_answers:
                        return self.completed_answers.pop(answer["task_id"])
            del self.locks[answer["task_id"]]


class PostWorker(Worker):
    """后端工作进程"""

    ANSWER_URI: str = "/v1/calculation/node-answer"

    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.http_session: aiohttp.ClientSession | None = None
        self.cos_client: LightCOSClient | None = None
        self.task_answer_completion_lock_state = TaskAnswerCompletionLockState()

    async def post_task_answer(
        self,
        data: TaskAnswer,
        form_data: aiohttp.FormData,
        payload_size: int,
        retry_count: int = 0,
    ) -> None:
        """发送任务应答"""
        if not self.http_session:
            return
        is_retry = False
        try:
            if retry_count > 3:
                logger.connector.error(
                    f"上报({data['method']})任务({data['task_id']})数据失败: 重试次数超过3次"
                )
                return
            response = await self.http_session.post(self.ANSWER_URI, data=form_data)
            match response.status:
                case 200:
                    logger.connector.info(
                        f"上报({data['method']})任务({data['task_id']})数据成功 大小: {humanize.naturalsize(payload_size)}"
                    )
                    if (
                        completed_answer
                        := await self.task_answer_completion_lock_state.complete(data)
                    ):
                        _data, _form_data, _payload_size = completed_answer
                        await self.post_task_answer(_data, _form_data, _payload_size)
                    return
                case 401:
                    logger.connector.error(
                        f"上报({data['method']})任务({data['task_id']})数据失败: 认证失败"
                    )
                    return
                case _:
                    logger.connector.warning(
                        f"上报({data['method']})任务({data['task_id']})数据失败: HTTP Error {response.status}, 重试 {retry_count + 1} / 3"
                    )
                    is_retry = True
        except aiohttp.ConnectionTimeoutError:
            logger.connector.warning(
                f"上报({data['method']})任务({data['task_id']})数据失败: 连接超时, 重试 {retry_count + 1} / 3"
            )
            is_retry = True
        except aiohttp.ServerConnectionError:
            logger.connector.warning(
                f"上报({data['method']})任务({data['task_id']})数据失败: 服务端连接失败, 疑似服务端不在线, 重试 {retry_count + 1} / 3"
            )
            is_retry = True
        except Exception as e:
            logger.connector.warning(
                f"上报({data['method']})任务({data['task_id']})数据失败: {e}, 重试 {retry_count + 1} / 3"
            )
            is_retry = True
        finally:
            if is_retry:
                await aio.sleep(retry_count)
                await self.tasks.create_task(
                    self.post_task_answer(
                        data, form_data, payload_size, retry_count + 1
                    )
                )

    async def post_to_cos(self, key: str, content: bytes | dict) -> None:
        """上传文件到COS"""
        if self.cos_client:
            if type(content) is bytes:
                bt_content = content
            else:
                bt_content = json.dumps(content).encode()
            await aio.to_thread(
                self.cos_client.upload_file, key=key, content=bt_content
            )

    async def post_answer(self, data: TaskAnswer) -> None:
        """发送任务应答"""
        if not self.settings or not self.cos_client:
            return
        # 非 completed 帧先占位：防止 completed 帧在 COS 上传期间因计数归零而抢先发送
        if not data["completed"]:
            await self.task_answer_completion_lock_state.pre_register(data["task_id"])
        if storage_key := data["result"].get("storage_key"):
            # 上传数据到COS
            await self.post_to_cos(storage_key, data["result"]["data"])
            # 并设置返回格式
            data_result = {
                "data_type": data["result"]["data_type"],
                "data": {
                    "storage_key": storage_key,
                },
            }
        else:
            data_result = data["result"]

        result: bytes = bson.encode(data_result, codec_options=codec_options)
        result = gzip.compress(result, compresslevel=6)
        form_data = aiohttp.FormData()
        form_data.add_field("result", result, content_type="application/gzip")
        form_data.add_field("task_id", data["task_id"])
        form_data.add_field("completed", str(data["completed"]))
        form_data.add_field("method", data["method"])
        form_data.add_field("access_id", self.settings.ENDPOINT_ACCESS_ID)
        form_data.add_field("access_token", self.settings.ENDPOINT_ACCESS_TOKEN)
        if await self.task_answer_completion_lock_state.is_completed(
            data, form_data, len(result)
        ):
            await self.post_task_answer(data, form_data, len(result))

    async def _run(self) -> None:
        """运行工作进程"""
        if not self.settings or not self.inq:
            return
        self.http_session = aiohttp.ClientSession(
            base_url=self.settings.ENDPOINT_URL.replace("ws://", "http://").replace(
                "wss://", "https://"
            )
        )
        self.cos_client = LightCOSClient(
            region=self.settings.RESULT_COS_REGION,
            secret_id=self.settings.RESULT_COS_SECRET_ID,
            secret_key=self.settings.RESULT_COS_SECRET_KEY,
            bucket=self.settings.RESULT_COS_BUCKET,
        )
        while self.running:
            try:
                data = await aio.to_thread(self.inq.get, timeout=1)
                if data is None:
                    self.running = False
                    break
                self.tasks.create_task(self.post_answer(data))
            except queue.Empty:
                continue
            except Exception as e:
                logger.connector.error(f"监听输入队列失败: {e}")
                continue
        await self.http_session.close()

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