"""资源配额管理 — 基于 psutil 的真实资源检查"""

from __future__ import annotations

import logging
import shutil
from dataclasses import dataclass, field

import psutil

from tool_agent.config import settings

logger = logging.getLogger(__name__)


@dataclass
class ResourceUsage:
    cpu_cores: float = 0.0
    memory_mb: float = 0.0
    gpu_memory_mb: float = 0.0


@dataclass
class SystemInfo:
    """系统资源概览"""
    cpu_count: int = 0
    cpu_percent: float = 0.0
    memory_total_mb: float = 0.0
    memory_available_mb: float = 0.0
    memory_percent: float = 0.0
    disk_total_gb: float = 0.0
    disk_free_gb: float = 0.0


class ResourceManager:
    """管理工具的 CPU/内存/显存配额，基于 psutil 做真实资源检查"""

    def __init__(self) -> None:
        self._allocated: dict[str, ResourceUsage] = {}

    @staticmethod
    def get_system_info() -> SystemInfo:
        """获取当前系统资源概览"""
        mem = psutil.virtual_memory()
        disk = shutil.disk_usage("/")
        return SystemInfo(
            cpu_count=psutil.cpu_count() or 1,
            cpu_percent=psutil.cpu_percent(interval=0.1),
            memory_total_mb=mem.total / (1024 * 1024),
            memory_available_mb=mem.available / (1024 * 1024),
            memory_percent=mem.percent,
            disk_total_gb=disk.total / (1024 ** 3),
            disk_free_gb=disk.free / (1024 ** 3),
        )

    def get_allocated_total(self) -> ResourceUsage:
        """已分配资源总计"""
        total = ResourceUsage()
        for usage in self._allocated.values():
            total.cpu_cores += usage.cpu_cores
            total.memory_mb += usage.memory_mb
            total.gpu_memory_mb += usage.gpu_memory_mb
        return total

    def check_available(self, cpu_cores: float, memory_mb: int, gpu: bool = False) -> bool:
        """检查是否有足够资源启动新工具

        比较已分配资源 + 新请求 与 系统可用资源。
        预留 20% 系统内存作为安全余量。
        """
        sys_info = self.get_system_info()
        allocated = self.get_allocated_total()

        # CPU 检查：已分配 + 新请求 不超过总核数
        if allocated.cpu_cores + cpu_cores > sys_info.cpu_count:
            logger.warning(
                f"CPU insufficient: allocated={allocated.cpu_cores}, "
                f"requested={cpu_cores}, total={sys_info.cpu_count}"
            )
            return False

        # 内存检查：已分配 + 新请求 不超过总内存的 80%
        safe_memory = sys_info.memory_total_mb * 0.8
        if allocated.memory_mb + memory_mb > safe_memory:
            logger.warning(
                f"Memory insufficient: allocated={allocated.memory_mb}MB, "
                f"requested={memory_mb}MB, safe_limit={safe_memory:.0f}MB"
            )
            return False

        return True

    def allocate(self, tool_id: str, cpu_cores: float, memory_mb: float, gpu_memory_mb: float = 0.0) -> None:
        """登记资源分配"""
        self._allocated[tool_id] = ResourceUsage(
            cpu_cores=cpu_cores,
            memory_mb=memory_mb,
            gpu_memory_mb=gpu_memory_mb,
        )
        logger.info(f"Allocated for '{tool_id}': cpu={cpu_cores}, mem={memory_mb}MB")

    def release(self, tool_id: str) -> None:
        """释放资源"""
        if tool_id in self._allocated:
            usage = self._allocated.pop(tool_id)
            logger.info(f"Released for '{tool_id}': cpu={usage.cpu_cores}, mem={usage.memory_mb}MB")

    def list_allocations(self) -> dict[str, ResourceUsage]:
        return dict(self._allocated)
