"""资源探针：本地引擎推理时采样峰值内存/CPU；模型大小读磁盘。

注意：psutil 测的是整个 Python 进程 RSS（含解释器与已加载库），用于本地模型横向对比是
合理代理量；报告会标注硬件环境，且只在本地引擎启用，云端 gummy 不测（填 N/A）。
"""

from __future__ import annotations

import os
import threading
import time

import psutil


class ResourceProbe:
    """上下文管理器：with 块内后台采样进程 RSS / CPU%。"""

    def __init__(self, interval: float = 0.05):
        self.interval = interval
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self._proc = psutil.Process(os.getpid())
        self.peak_rss = 0
        self._cpu: list[float] = []

    def __enter__(self) -> "ResourceProbe":
        self._proc.cpu_percent(None)  # 预热：首次调用返回 0，丢弃
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        return self

    def _run(self) -> None:
        while not self._stop.is_set():
            try:
                self.peak_rss = max(self.peak_rss, self._proc.memory_info().rss)
                c = self._proc.cpu_percent(None)
                if c > 0:
                    self._cpu.append(c)
            except Exception:
                pass
            time.sleep(self.interval)

    def __exit__(self, *_exc) -> None:
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=1.0)

    @property
    def peak_rss_mb(self) -> float:
        return self.peak_rss / 1024 / 1024

    @property
    def avg_cpu(self) -> float:
        return sum(self._cpu) / len(self._cpu) if self._cpu else 0.0


def dir_size_mb(path: str) -> float:
    """目录下所有文件体积之和（MB）；路径不存在返回 0。"""
    total = 0
    if not path or not os.path.exists(path):
        return 0.0
    if os.path.isfile(path):
        return os.path.getsize(path) / 1024 / 1024
    for root, _dirs, files in os.walk(path):
        for f in files:
            try:
                total += os.path.getsize(os.path.join(root, f))
            except OSError:
                pass
    return total / 1024 / 1024


def hardware_info() -> dict:
    """报告里标注的硬件环境。"""
    import platform

    vm = psutil.virtual_memory()
    return {
        "platform": platform.platform(),
        "machine": platform.machine(),
        "cpu_count": psutil.cpu_count(logical=True),
        "cpu_count_physical": psutil.cpu_count(logical=False),
        "total_mem_gb": round(vm.total / 1024**3, 1),
    }