#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：resource_monitor.py
模块描述：执行机资源监控器 —— 采集压测执行机自身的 CPU/内存/网络 I/O 指标，
         用于判断压测结果是否受执行机资源瓶颈限制。
         psutil 可用时优先使用；缺失时 Linux 读 /proc，Windows 走 ctypes 标准库兜底。

作者：czj
创建日期：2026-08-21
最后修改：2026-08-21
"""

import logging
import os
import sys
import time
from typing import Optional, Tuple

logger = logging.getLogger(__name__)

# psutil 为可选依赖，优先使用；缺失时走标准库兜底
try:
    import psutil
    _HAS_PSUTIL = True
except ImportError:
    psutil = None
    _HAS_PSUTIL = False


class ResourceMonitor:
    """
    执行机资源监控器

    每个采样周期调用一次 sample()，返回 CPU/内存/网络指标字典。
    网络速率通过相邻两次采样累计计数器的差值计算（首次采样速率为 0，跳过差值）。
    """

    def __init__(self, sampling_interval: float = 1.0):
        self._interval = sampling_interval
        self._last_net: Optional[Tuple[float, float]] = None  # (累计发送B, 累计接收B)
        self._last_ts: Optional[float] = None                 # 上次采样时间戳
        # Windows 兜底：CPU 差值基准（首次调用 cpu_times 建基线，第二次才有差值）
        self._win_cpu_last: Optional[Tuple[float, float, float, float]] = None
        # Linux /proc 兜底：CPU 差值基准
        self._linux_cpu_last: Optional[Tuple[float, float]] = None  # (idle, total)

    # ==================== 对外接口 ====================

    def sample(self) -> dict:
        """
        采集一次资源指标

        Returns:
            dict: {
                "cpuPercent": float,       # CPU 使用率 %
                "memPercent": float,       # 内存使用率 %
                "memUsedMb": float,        # 已用内存 MB
                "memTotalMb": float,       # 总内存 MB
                "netSentKbS": float,       # 发送速率 KB/s
                "netRecvKbS": float,       # 接收速率 KB/s
                "netSentTotalMb": float,   # 累计发送 MB
                "netRecvTotalMb": float,   # 累计接收 MB
                "backend": str             # 实际使用的采集后端（psutil/proc/winapi/none）
            }
            采集失败时各值为 0.0，不抛异常（资源监控不能影响压测主流程）
        """
        now = time.time()
        result = {
            "cpuPercent": 0.0,
            "memPercent": 0.0,
            "memUsedMb": 0.0,
            "memTotalMb": 0.0,
            "netSentKbS": 0.0,
            "netRecvKbS": 0.0,
            "netSentTotalMb": 0.0,
            "netRecvTotalMb": 0.0,
            "backend": "none",
        }

        try:
            if _HAS_PSUTIL:
                self._sample_psutil(result, now)
                result["backend"] = "psutil"
            elif sys.platform == "win32":
                self._sample_windows(result, now)
                result["backend"] = "winapi"
            else:
                self._sample_linux(result, now)
                result["backend"] = "proc"
        except Exception as e:
            # 兜底实现自身异常也不能影响压测，仅记录
            logger.debug(f"资源采样失败: {e}")

        self._last_ts = now
        return result

    def dispose(self) -> None:
        """释放资源（当前无外部句柄，保留接口以对齐 PRD 设计）"""
        self._last_net = None
        self._last_ts = None
        self._win_cpu_last = None
        self._linux_cpu_last = None

    # ==================== psutil 后端 ====================

    def _sample_psutil(self, result: dict, now: float) -> None:
        """psutil 采集（跨平台首选，开销 < 5ms）"""
        # CPU：interval=None 立即返回自上次调用以来的差值（首次调用返回 0）
        result["cpuPercent"] = float(psutil.cpu_percent(interval=None))

        mem = psutil.virtual_memory()
        result["memPercent"] = float(mem.percent)
        result["memUsedMb"] = round(mem.used / (1024 * 1024), 1)
        result["memTotalMb"] = round(mem.total / (1024 * 1024), 1)

        net = psutil.net_io_counters()
        sent_b, recv_b = float(net.bytes_sent), float(net.bytes_recv)
        result["netSentTotalMb"] = round(sent_b / (1024 * 1024), 2)
        result["netRecvTotalMb"] = round(recv_b / (1024 * 1024), 2)
        self._apply_net_rate(result, sent_b, recv_b, now)

    # ==================== Linux /proc 兜底 ====================

    def _sample_linux(self, result: dict, now: float) -> None:
        """Linux 标准库兜底：/proc/stat + /proc/meminfo + /proc/net/dev"""
        result["cpuPercent"] = self._linux_cpu()

        mem_percent, mem_used_mb, mem_total_mb = self._linux_mem()
        result["memPercent"] = mem_percent
        result["memUsedMb"] = mem_used_mb
        result["memTotalMb"] = mem_total_mb

        sent_b, recv_b = self._linux_net()
        result["netSentTotalMb"] = round(sent_b / (1024 * 1024), 2)
        result["netRecvTotalMb"] = round(recv_b / (1024 * 1024), 2)
        self._apply_net_rate(result, sent_b, recv_b, now)

    def _linux_cpu(self) -> float:
        """/proc/stat 首行差值计算 CPU 使用率（只读前 2 行降低开销）"""
        with open("/proc/stat", "r") as f:
            fields = f.readline().split()
        # cpu  user nice system idle iowait irq softirq steal ...
        values = [float(x) for x in fields[1:]]
        idle = values[3] + (values[4] if len(values) > 4 else 0.0)  # idle + iowait
        total = sum(values)
        if self._linux_cpu_last is None:
            self._linux_cpu_last = (idle, total)
            return 0.0
        last_idle, last_total = self._linux_cpu_last
        self._linux_cpu_last = (idle, total)
        d_total = total - last_total
        d_idle = idle - last_idle
        if d_total <= 0:
            return 0.0
        return round(max(0.0, min(100.0, (1.0 - d_idle / d_total) * 100.0)), 1)

    def _linux_mem(self) -> Tuple[float, float, float]:
        """/proc/meminfo 读取内存（MemTotal/MemFree/MemAvailable/Buffers/Cached）"""
        info: dict = {}
        # 关键字段只在前几行，逐行读到齐即止
        with open("/proc/meminfo", "r") as f:
            for line in f:
                parts = line.split(":", 1)
                if len(parts) != 2:
                    continue
                info[parts[0]] = float(parts[1].strip().split()[0])  # kB
                if len(info) >= 5:
                    break
        total_kb = info.get("MemTotal", 0.0)
        if total_kb <= 0:
            return 0.0, 0.0, 0.0
        free_kb = info.get("MemFree", 0.0)
        available_kb = info.get("MemAvailable", free_kb)
        buffers_kb = info.get("Buffers", 0.0)
        cached_kb = info.get("Cached", 0.0)
        used_kb = total_kb - free_kb - buffers_kb - cached_kb
        percent = round(used_kb / total_kb * 100.0, 1) if total_kb > 0 else 0.0
        # 注：percent 优先用 available 口径（与 psutil 一致）
        if total_kb > 0 and available_kb > 0:
            percent = round((total_kb - available_kb) / total_kb * 100.0, 1)
        return percent, round(used_kb / 1024.0, 1), round(total_kb / 1024.0, 1)

    def _linux_net(self) -> Tuple[float, float]:
        """/proc/net/dev 累计所有物理网卡的收发字节数"""
        sent_b = 0.0
        recv_b = 0.0
        with open("/proc/net/dev", "r") as f:
            for i, line in enumerate(f):
                if i < 2:  # 前两行是表头
                    continue
                if ":" not in line:
                    continue
                _iface, data = line.split(":", 1)
                fields = data.split()
                if len(fields) < 10:
                    continue
                recv_b += float(fields[0])   # receive bytes
                sent_b += float(fields[8])   # transmit bytes
        return sent_b, recv_b

    # ==================== Windows ctypes 兜底 ====================

    def _sample_windows(self, result: dict, now: float) -> None:
        """Windows 标准库兜底：GetSystemTimes + GlobalMemoryStatusEx + GetIfTable2"""
        result["cpuPercent"] = self._win_cpu()

        mem_percent, mem_used_mb, mem_total_mb = self._win_mem()
        result["memPercent"] = mem_percent
        result["memUsedMb"] = mem_used_mb
        result["memTotalMb"] = mem_total_mb

        sent_b, recv_b = self._win_net()
        result["netSentTotalMb"] = round(sent_b / (1024 * 1024), 2)
        result["netRecvTotalMb"] = round(recv_b / (1024 * 1024), 2)
        self._apply_net_rate(result, sent_b, recv_b, now)

    def _win_cpu(self) -> float:
        """ctypes 调 GetSystemTimes 计算 CPU 使用率（idle/kernel/user 差值）"""
        import ctypes
        from ctypes import wintypes

        class FILETIME(ctypes.Structure):
            _fields_ = [("dwLowDateTime", wintypes.DWORD), ("dwHighDateTime", wintypes.DWORD)]

        idle, kernel, user = FILETIME(), FILETIME(), FILETIME()
        if not ctypes.windll.kernel32.GetSystemTimes(
            ctypes.byref(idle), ctypes.byref(kernel), ctypes.byref(user)
        ):
            return 0.0

        def to_100ns(ft: FILETIME) -> int:
            return (ft.dwHighDateTime << 32) | ft.dwLowDateTime

        idle_t = to_100ns(idle)
        kernel_t = to_100ns(kernel)  # kernel 含 idle，需减去
        user_t = to_100ns(user)
        current = (idle_t, kernel_t, user_t, kernel_t + user_t)
        if self._win_cpu_last is None:
            self._win_cpu_last = current
            return 0.0
        last = self._win_cpu_last
        self._win_cpu_last = current
        d_total = current[3] - last[3]
        d_idle = current[0] - last[0]
        if d_total <= 0:
            return 0.0
        return round(max(0.0, min(100.0, (1.0 - d_idle / d_total) * 100.0)), 1)

    def _win_mem(self) -> Tuple[float, float, float]:
        """ctypes 调 GlobalMemoryStatusEx 读取内存"""
        import ctypes
        from ctypes import wintypes

        class MEMORYSTATUSEX(ctypes.Structure):
            _fields_ = [
                ("dwLength", wintypes.DWORD),
                ("dwMemoryLoad", wintypes.DWORD),
                ("ullTotalPhys", ctypes.c_ulonglong),
                ("ullAvailPhys", ctypes.c_ulonglong),
                ("ullTotalPageFile", ctypes.c_ulonglong),
                ("ullAvailPageFile", ctypes.c_ulonglong),
                ("ullTotalVirtual", ctypes.c_ulonglong),
                ("ullAvailVirtual", ctypes.c_ulonglong),
                ("ullAvailExtendedVirtual", ctypes.c_ulonglong),
            ]

        stat = MEMORYSTATUSEX()
        stat.dwLength = ctypes.sizeof(MEMORYSTATUSEX)
        if not ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(stat)):
            return 0.0, 0.0, 0.0
        total_mb = stat.ullTotalPhys / (1024 * 1024)
        used_mb = (stat.ullTotalPhys - stat.ullAvailPhys) / (1024 * 1024)
        return float(stat.dwMemoryLoad), round(used_mb, 1), round(total_mb, 1)

    def _win_net(self) -> Tuple[float, float]:
        """ctypes 调 GetIfEntry2 (GetIfTable2) 汇总所有网卡收发字节（Vista+）"""
        import ctypes
        from ctypes import wintypes

        try:
            iptc = ctypes.windll.iphlpapi
        except AttributeError:
            return 0.0, 0.0

        # 简化实现：用 GetIfTable（XP+ 兼容，结构体固定 32 位长度）
        class MIB_IFROW(ctypes.Structure):
            _fields_ = [
                ("wszName", ctypes.c_wchar * 256),
                ("dwIndex", wintypes.DWORD),
                ("dwType", wintypes.DWORD),
                ("dwMtu", wintypes.DWORD),
                ("dwSpeed", wintypes.DWORD),
                ("dwPhysAddrLen", wintypes.DWORD),
                ("bPhysAddr", ctypes.c_byte * 8),
                ("dwAdminStatus", wintypes.DWORD),
                ("dwOperStatus", wintypes.DWORD),
                ("dwLastChange", wintypes.DWORD),
                ("dwInOctets", wintypes.DWORD),
                ("dwInUcastPkts", wintypes.DWORD),
                ("dwInNUcastPkts", wintypes.DWORD),
                ("dwInDiscards", wintypes.DWORD),
                ("dwInErrors", wintypes.DWORD),
                ("dwInUnknownProtos", wintypes.DWORD),
                ("dwOutOctets", wintypes.DWORD),
                ("dwOutUcastPkts", wintypes.DWORD),
                ("dwOutNUcastPkts", wintypes.DWORD),
                ("dwOutDiscards", wintypes.DWORD),
                ("dwOutErrors", wintypes.DWORD),
                ("dwOutQLen", wintypes.DWORD),
                ("dwDescrLen", wintypes.DWORD),
                ("bDescr", ctypes.c_byte * 256),
            ]

        # 先查询表大小
        size = wintypes.DWORD(0)
        if iptc.GetIfTable(None, ctypes.byref(size), 0) != 0 and size.value == 0:
            return 0.0, 0.0

        num_entries = size.value // ctypes.sizeof(MIB_IFROW)
        if num_entries == 0:
            return 0.0, 0.0
        # 上限保护（异常环境防御）
        num_entries = min(num_entries, 256)

        table = (MIB_IFROW * num_entries)()
        table_size = wintypes.DWORD(ctypes.sizeof(table))
        if iptc.GetIfTable(ctypes.byref(table), ctypes.byref(table_size), 0) != 0:
            return 0.0, 0.0

        sent_b = 0.0
        recv_b = 0.0
        for row in table:
            # dwOperStatus: 6 = IF_OPER_STATUS_OPERATIONAL/UP
            if row.dwOperStatus >= 6 or row.dwInOctets > 0 or row.dwOutOctets > 0:
                recv_b += float(row.dwInOctets)
                sent_b += float(row.dwOutOctets)
        return sent_b, recv_b

    # ==================== 通用工具 ====================

    def _apply_net_rate(self, result: dict, sent_b: float, recv_b: float, now: float) -> None:
        """
        计算网络速率（KB/s）：相邻采样差值 / 时间间隔

        首次采样无差值基准，速率记 0（跳过第 1 次差值，PRD 风险缓解措施）
        """
        if self._last_net is None or self._last_ts is None:
            self._last_net = (sent_b, recv_b)
            return
        last_sent, last_recv = self._last_net
        dt = now - self._last_ts
        if dt <= 0:
            dt = self._interval
        sent_delta = max(0.0, sent_b - last_sent)
        recv_delta = max(0.0, recv_b - last_recv)
        result["netSentKbS"] = round(sent_delta / 1024.0 / dt, 1)
        result["netRecvKbS"] = round(recv_delta / 1024.0 / dt, 1)
        self._last_net = (sent_b, recv_b)


# ==================== 模块自测 ====================

if __name__ == "__main__":
    # 手动验证：python -m app.executors.resource_monitor
    monitor = ResourceMonitor(sampling_interval=1.0)
    print(f"psutil 可用: {_HAS_PSUTIL}")
    for i in range(5):
        time.sleep(1.0)
        data = monitor.sample()
        print(
            f"[{i + 1}] cpu={data['cpuPercent']}% mem={data['memPercent']}% "
            f"({data['memUsedMb']}/{data['memTotalMb']}MB) "
            f"sent={data['netSentKbS']}KB/s recv={data['netRecvKbS']}KB/s "
            f"backend={data['backend']}"
        )
    monitor.dispose()
