import asyncio
from typing import Callable

from ldnode.runner.schemas import BaseContext
from ldnode.runner.events import Events


class LoadCounter:
    """ 负载计数器 """
    DECAY_1M = 0.92 # interval: 5s
    
    def __init__(self, context: BaseContext, load_querier: Callable[[], float], running_checker: Callable[[], bool]):
        self.context = context
        self.load_1m: float = 0
        self.tick_task: asyncio.Task | None = None
        # 外部接口
        self.query_load = load_querier
        self.check_running = running_checker
        
    def _update(self, load_now: float) -> None:
        """ 更新负载 """
        self.load_1m = self.load_1m * self.DECAY_1M + load_now * (1 - self.DECAY_1M)

    def _ewma(self) -> None:
        """ 指数加权移动平均更新负载 """
        self._update(self.query_load())
        
    async def _tick_loop(self) -> None:
        """ 定时循环 """
        if self.check_running():
            self._ewma()
            await self.context.submit_event(
                "calculator", 
                Events.UPDATE_CALCULATOR_LOAD, 
                { "load": self.load_1m },
            )
            await asyncio.sleep(5)
            self.tick_task = asyncio.create_task(self._tick_loop())
            
    async def count(self) -> None:
        """ 开始计数 """
        self.tick_task = asyncio.create_task(self._tick_loop())
