import asyncio
import time
from typing import Any, Optional

from ldnode import logger
from ldnode.runner.config import Settings
from ldnode.runner.events import (
    CalculatorLoadUpdated,
    Events,
    LockResource,
    WorkerStateChanged,
    on,
)
from ldnode.runner.schemas import BaseContext, Service

from .operators import (
    load_ping_info,
    restart,
    upgrade,
)


class MonitorService(Service):
    """监控服务"""

    def __init__(self, settings: Settings, context: BaseContext):
        super().__init__(settings, context)
        self.ping_running = False
        self.ping_task: Optional[asyncio.Task] = None
        self.end_flag = asyncio.Event()
        self.calculator_worker: set[str] = set()
        self.calculator_load: float = 0
        self.calculator_free: int = 0
        self.upgrade_lock = asyncio.Lock()

    async def next_ping(self) -> None:
        """下一个心跳"""
        if self.ping_running:
            await self.context.submit_event(
                "Main",
                Events.PING,
                {
                    "calculator_worker_num": len(self.calculator_worker),
                    "calculator_load": self.calculator_load,
                    "calculator_free": self.calculator_free,
                    **load_ping_info(),
                },
            )
            logger.monitor.info(f"发送心跳")
            await asyncio.sleep(self.settings.PING_INTERVAL)
            self.ping_task = asyncio.create_task(self.next_ping())
        else:
            self.ping_task = None

    @on(Events.ONLINE)
    async def handle_online(self, _source: str, _data: Any) -> None:
        """处理上线事件"""
        self.ping_running = True
        if not self.ping_task:
            self.ping_task = asyncio.create_task(self.next_ping())

    @on(Events.RESTART)
    async def handle_restart(self, _source: str, _data: Any) -> None:
        """处理重启事件"""
        restart()

    @on(Events.UPGRADE)
    async def handle_upgrade(self, _source: str, data: Any) -> None:
        """处理升级事件"""
        if self.upgrade_lock.locked():
            logger.monitor.warning("升级任务已在执行，忽略本次请求")
            return
        async with self.upgrade_lock:
            version = data.get("version") if isinstance(data, dict) else None
            logger.monitor.info(f"开始升级 ldnode, version={version or 'latest'}")
            ok, output = await asyncio.to_thread(upgrade, version)
            if ok:
                logger.monitor.info("升级成功，准备重启 ldnode")
                if output:
                    logger.monitor.info(f"升级输出(尾部): {output}")
                restart()
            else:
                logger.monitor.error("升级失败，保持当前进程继续运行")
                if output:
                    logger.monitor.error(f"升级输出(尾部): {output}")

    @on(Events.WORKER_START)
    async def handle_worker_start(self, _source: str, data: WorkerStateChanged) -> None:
        """处理工作进程启动事件"""
        if "calculator" in data["worker_id"]:
            self.calculator_worker.add(data["worker_id"])

    @on(Events.WORKER_EXIT)
    async def handle_worker_exit(self, _source: str, data: WorkerStateChanged) -> None:
        """处理工作进程退出事件"""
        if "calculator" in data["worker_id"]:
            self.calculator_worker.remove(data["worker_id"])

    @on(Events.UPDATE_CALCULATOR_LOAD)
    async def handle_caculator_load_updated(
        self, _source: str, data: CalculatorLoadUpdated
    ) -> None:
        """处理计算负载更新事件"""
        self.calculator_load = round(data["load"], 2)

    @on(Events.LOCK_RESOURCE)
    async def handle_lock_resource(self, _source: str, _data: LockResource) -> None:
        """处理锁定资源事件"""
        self.calculator_free -= 1

    @on(Events.RELEASE_RESOURCE)
    async def handle_release_resource(self, _source: str, _data: LockResource) -> None:
        """处理释放资源事件"""
        self.calculator_free += 1

    async def start(self) -> None:
        """启动服务"""
        await self.end_flag.wait()
        logger.monitor.info("监控服务已退出~")

    async def shutdown(self) -> None:
        """关闭服务"""
        self.ping_running = False
        self.end_flag.set()
