from enum import StrEnum
from typing import Any, Awaitable, Callable, TypedDict


def on(event: str):
    """事件装饰器"""

    def decorator(
        func: Callable[[str, str, Any], Awaitable[None]],
    ) -> Callable[[str, str, Any], Awaitable[None]]:
        """事件装饰器"""
        func.__annotations__["on_ev"] = event
        return func

    return decorator


class Events(StrEnum):
    """事件"""

    RESTART = "restart"  # 重启
    UPGRADE = "upgrade"  # 升级
    EXIT = "exit"  # 退出
    READ_LOGGERS = "read_loggers"  # 读取日志器
    READ_LOG = "read_log"  # 读取日志
    ONLINE = "online"  # 上线
    OFFLINE = "offline"  # 下线
    PING = "ping"  # 心跳
    PONG = "pong"  # 回应心跳
    SHUTDOWN = "shutdown"  # 关闭

    WORKER_EXIT = "worker_exit"  # 工作进程退出
    WORKER_START = "worker_start"  # 工作进程启动

    TASK_SUBMITTED = "task_submitted"  # 任务提交
    TASK_ANSWER = "task_answer"  # 任务应答
    TASK_CANCEL = "task_cancel"  # 任务取消
    TASK_FINISHED = "task_finished"  # 任务完成
    TASK_CANCELED = "task_canceled"  # 任务取消
    LOCK_RESOURCE = "lock_resource"  # 锁定资源
    RELEASE_RESOURCE = "release_resource"  # 释放资源
    UPDATE_CALCULATOR_LOAD = "update_calculator_load"  # 更新计算负载

    DATA_SYNC = "data_sync"  # 数据同步
    DATASET_STATUS_REPORT = "dataset_status_report"  # 数据集状态上报


class CalculatorLoadUpdated(TypedDict):
    """计算负载更新"""

    load: float


class LockResource(TypedDict):
    """锁定资源"""

    cpu_resource: int  # CPU资源
    gpu_resource: int  # GPU资源


class WorkerStateChanged(TypedDict):
    """工作进程状态改变"""

    worker_id: int  # 工作进程ID
    process_id: int  # 进程ID


class DataSync(TypedDict):
    """数据同步"""

    filename: str  # 文件名
    is_ok: bool  # 是否成功
    latest_tag: str  # 最新标签
    error_message: str  # 错误信息


class TaskSubmitted(TypedDict):
    """任务提交"""

    task_id: str  # 任务ID
    method: str  # 方法
    parameters: dict[str, Any]  # 参数


class TaskAnswer(TypedDict):
    """任务应答"""

    task_id: str  # 任务ID
    method: str  # 方法
    completed: bool  # 是否完成
    result: dict  # 结果


class TaskFinished(TypedDict):
    """任务完成"""

    task_id: str  # 任务ID
    method: str  # 方法


class TaskCancel(TypedDict):
    """任务取消"""

    task_id: str  # 任务ID
    method: str  # 方法


class TaskCanceled(TypedDict):
    """任务取消"""

    task_id: str  # 任务ID
    method: str  # 方法
    is_released: bool  # 是否释放资源


class DatasetStatusReport(TypedDict):
    """数据集状态上报"""

    datasets: list[dict]  # DatasetStatusItem 的 dict 列表


class ReadLog(TypedDict):
    """读取日志"""

    logger: str  # 日志器
    offset: int  # 起始字节数
    limit: int  # 限制返回字节数


class UpgradeRequest(TypedDict):
    """升级请求"""

    version: str | None  # 目标版本，为空则升级到最新


class Ping(TypedDict):
    """心跳"""

    report_time: int  # 报告时间
    cpu_usage: float  # CPU使用率
    memory_usage: float  # 内存使用率
    disk_usage: float  # 磁盘使用率
    load_average: float  # 负载平均值
    calculator_worker_num: int  # 计算工作进程数量
    calculator_load: float  # 计算工作进程负载
    calculator_free: float  # 计算工作进程空闲数
