# coding: utf-8

from typing import List
from dataclasses import dataclass, field
import time

from .psutil_helper import get_cpu_percent, get_memory_used


@dataclass
class Arguments:
    host: str = "0.0.0.0"
    port: int = 8000
    max_num: int = 3000
    interval: float = 1.0

    def to_dict(self):
        return {
            "host": self.host,
            "port": self.port,
            "max_num": self.max_num,
            "interval": self.interval,
        }

@dataclass
class SystemInfo:
    cpu_percent: float
    memory_used: float
    timestamp: float
    gpu: List["GPUInfo"] = field(default_factory=list)

    def to_dict(self):
        return {
            "cpu_percent": self.cpu_percent,
            "memory_used": self.memory_used,
            "timestamp": self.timestamp,
            "gpu": [g.to_dict() for g in self.gpu],
        }

    @classmethod
    def from_dict(cls, data: dict):
        return cls(
            cpu_percent=data["cpu_percent"],
            memory_used=data["memory_used"],
            timestamp=data["timestamp"],
            gpu=[GPUInfo.from_dict(g) for g in data.get("gpu", [])],
        )
    
@dataclass
class CPUStatistic:
    min: float = 0.0    
    max: float = 0.0
    avg: float = 0.0
    def to_dict(self):
        return {
            "min": self.min,
            "max": self.max,
            "avg": self.avg,
        }
    @classmethod
    def from_dict(cls, data: dict):
        return cls(
            min=data["min"],
            max=data["max"],
            avg=data["avg"],
        )

@dataclass
class MemoryStatistic:
    min: float = 0.0
    max: float = 0.0
    avg: float = 0.0
    def to_dict(self):
        return {
            "min": self.min,
            "max": self.max,
            "avg": self.avg,
        }
    @classmethod
    def from_dict(cls, data: dict):
        return cls(
            min=data["min"],
            max=data["max"],
            avg=data["avg"],
        )


@dataclass
class GPUStatistic:
    min: float = 0.0
    max: float = 0.0
    avg: float = 0.0

    def to_dict(self):
        return {
            "min": self.min,
            "max": self.max,
            "avg": self.avg,
        }

    @classmethod
    def from_dict(cls, data: dict):
        return cls(
            min=data["min"],
            max=data["max"],
            avg=data["avg"],
        )


@dataclass
class GPUInfo:
    """单块 GPU 的占用率信息"""
    index: int = 0
    name: str = ""
    gpu_util: int = 0
    mem_usage: float = 0.0
    mem_util: int = 0
    mem_used_mb: float = 0.0
    mem_total_mb: float = 0.0

    def to_dict(self):
        return {
            "index": self.index,
            "name": self.name,
            "gpu_util": self.gpu_util,
            "mem_usage": self.mem_usage,
            "mem_util": self.mem_util,
            "mem_used_mb": self.mem_used_mb,
            "mem_total_mb": self.mem_total_mb,
        }

    @classmethod
    def from_dict(cls, data: dict):
        return cls(
            index=data.get("index", 0),
            name=data.get("name", ""),
            gpu_util=data.get("gpu_util", 0),
            mem_usage=data.get("mem_usage", 0.0),
            mem_util=data.get("mem_util", 0),
            mem_used_mb=data.get("mem_used_mb", 0.0),
            mem_total_mb=data.get("mem_total_mb", 0.0),
        )


@dataclass
class Statistic:
    cpu: CPUStatistic
    memory: MemoryStatistic
    data_list: List[SystemInfo]
    total_length: int
    gpu_statistic: GPUStatistic
    gpu: List[GPUInfo] = field(default_factory=list)

    def to_dict(self):
        return {
            "cpu": self.cpu.to_dict(),
            "memory": self.memory.to_dict(),
            "data_list": [system_info.to_dict() for system_info in self.data_list],
            "total_length": self.total_length,
            "gpu_statistic": self.gpu_statistic.to_dict(),
            "gpu": [g.to_dict() for g in self.gpu],
        }

    @classmethod
    def from_dict(cls, data: dict):
        return cls(
            cpu=CPUStatistic.from_dict(data["cpu"]),
            memory=MemoryStatistic.from_dict(data["memory"]),
            data_list=[SystemInfo.from_dict(system_info) for system_info in data["data_list"]],
            total_length=data["total_length"],
            gpu_statistic=GPUStatistic.from_dict(
                data.get("gpu_statistic", {"min": 0.0, "max": 0.0, "avg": 0.0})
            ),
            gpu=[GPUInfo.from_dict(g) for g in data.get("gpu", [])],
        )

class SystemInfoCollector:
    def __init__(self, max_system_info_list_length: int = 3000):
        self.system_info_list: List[SystemInfo] = []
        self.max_system_info_list_length = max_system_info_list_length

    @property
    def total_length(self) -> int:
        return len(self.system_info_list)

    @property
    def latest_system_info(self) -> SystemInfo:
        return self.system_info_list[-1] if self.system_info_list else None

    def collect(self) -> SystemInfo:
        # 延迟导入，避免 schema 与 gpu_helper 的循环导入问题
        from .gpu_helper import get_gpu_info

        system_info = SystemInfo(
            get_cpu_percent(), get_memory_used(), time.time(), get_gpu_info())
        self.system_info_list.append(system_info)
        if len(self.system_info_list) > self.max_system_info_list_length:
            self.system_info_list.pop(0)
        return system_info

    def get_system_info_list(self) -> List[SystemInfo]:
        return self.system_info_list

    def reset_system_info_list(self) -> None:
        self.system_info_list.clear()

    def statistic(self) -> dict:
        # 计算GPU显存占用，合并多张显卡的显存用量
        gpu_mem_used_mb_list = [
            sum(g.mem_used_mb for g in system_info.gpu)
            for system_info in self.system_info_list
            if system_info.gpu
        ]
        if self.total_length == 0:
            statistic = Statistic(
                cpu=CPUStatistic(),
                memory=MemoryStatistic(),
                data_list=[],
                total_length=0,
                gpu_statistic=GPUStatistic(),
                gpu=[],
            )
        else:
            statistic = Statistic(
                cpu=CPUStatistic(
                    min=min(system_info.cpu_percent for system_info in self.system_info_list),
                    max=max(system_info.cpu_percent for system_info in self.system_info_list),
                    avg=sum(system_info.cpu_percent for system_info in self.system_info_list) / len(self.system_info_list),
                ),
                memory=MemoryStatistic(
                    min=min(system_info.memory_used for system_info in self.system_info_list),
                    max=max(system_info.memory_used for system_info in self.system_info_list),
                    avg=sum(system_info.memory_used for system_info in self.system_info_list) / len(self.system_info_list),
                ),
                gpu=GPUStatistic(
                    min=min(gpu_mem_used_mb_list) if gpu_mem_used_mb_list else 0.0,
                    max=max(gpu_mem_used_mb_list) if gpu_mem_used_mb_list else 0.0,
                    avg=sum(gpu_mem_used_mb_list) / len(gpu_mem_used_mb_list) if gpu_mem_used_mb_list else 0.0,
                ),
                data_list=self.system_info_list,
                total_length=self.total_length,
            )
        return statistic.to_dict()