module/system_monitor.py

import os
import threading
import time
import warnings
import psutil
import torch
from datetime import datetime

from module.config import LOGS_DIR


def _snapshot():
    """采集当前时刻的CPU、内存和GPU占用率"""
    stats = {
        'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3],
        'cpu_percent': psutil.cpu_percent(interval=0),
        'memory_percent': psutil.virtual_memory().percent,
        'memory_used_gb': round(psutil.virtual_memory().used / (1024**3), 2),
        'memory_total_gb': round(psutil.virtual_memory().total / (1024**3), 2),
    }
    if torch.cuda.is_available():
        stats['gpu_memory_used_gb'] = round(torch.cuda.memory_allocated() / (1024**3), 2)
        stats['gpu_memory_reserved_gb'] = round(torch.cuda.memory_reserved() / (1024**3), 2)
        stats['gpu_memory_total_gb'] = round(torch.cuda.get_device_properties(0).total_memory / (1024**3), 2)
        stats['gpu_memory_percent'] = round(
            stats['gpu_memory_used_gb'] / stats['gpu_memory_total_gb'] * 100, 1
        )
        stats['gpu_util_percent'] = _get_gpu_util()
    else:
        stats['gpu_memory_percent'] = 0
        stats['gpu_memory_used_gb'] = 0
        stats['gpu_memory_reserved_gb'] = 0
        stats['gpu_memory_total_gb'] = 0
        stats['gpu_util_percent'] = 0
    return stats


def _get_gpu_util():
    """获取GPU计算利用率（非显存占用），单例模式避免反复init/shutdown"""
    if not _get_gpu_util._initialized:
        try:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                import pynvml as nvml_mod
            nvml_mod.nvmlInit()
            _get_gpu_util._nvml = nvml_mod
            _get_gpu_util._handle = nvml_mod.nvmlDeviceGetHandleByIndex(0)
            _get_gpu_util._initialized = True
        except Exception:
            return -1
    try:
        nvml_mod = _get_gpu_util._nvml
        handle = _get_gpu_util._handle
        util = nvml_mod.nvmlDeviceGetUtilizationRates(handle)
        return util.gpu
    except Exception:
        return -1

_get_gpu_util._initialized = False
_get_gpu_util._nvml = None
_get_gpu_util._handle = None


class ResourceMonitor:
    """后台资源监控器：在检测期间持续采样，追踪峰值"""

    def __init__(self, sample_interval=0.5):
        """初始化监控器"""
        self.sample_interval = sample_interval
        self._thread = None
        self._stop_event = threading.Event()
        self._samples = []
        self._lock = threading.Lock()
        self._peak = None  # 检测期间的峰值
        self._monitoring = False

    def start_monitoring(self):
        """开始后台采样（检测开始前调用）"""
        self._samples = []
        self._peak = None
        self._stop_event.clear()
        self._monitoring = True
        self._thread = threading.Thread(target=self._sample_loop, daemon=True)
        self._thread.start()

    def stop_monitoring(self):
        """停止后台采样（检测结束后调用），返回检测期间的峰值统计"""
        self._stop_event.set()
        if self._thread:
            self._thread.join(timeout=3)
        self._monitoring = False

        with self._lock:
            if not self._samples:
                return None
            peak = self._compute_peak(self._samples)
            self._peak = peak
            return peak

    @property
    def is_monitoring(self):
        return self._monitoring

    def _sample_loop(self):
        """后台采样循环"""
        # 先做一次有interval的CPU采样来初始化psutil的基准值
        psutil.cpu_percent(interval=0)
        time.sleep(0.1)

        while not self._stop_event.is_set():
            sample = _snapshot()
            with self._lock:
                self._samples.append(sample)
            self._stop_event.wait(self.sample_interval)

    def _compute_peak(self, samples):
        """从采样列表中计算峰值统计"""
        if not samples:
            return None

        peak = {
            'timestamp_start': samples[0]['timestamp'],
            'timestamp_end': samples[-1]['timestamp'],
            'sample_count': len(samples),
            'duration_sec': round(len(samples) * self.sample_interval, 1),
            'cpu_peak': max(s['cpu_percent'] for s in samples),
            'cpu_avg': round(sum(s['cpu_percent'] for s in samples) / len(samples), 1),
            'memory_peak_percent': max(s['memory_percent'] for s in samples),
            'memory_peak_gb': max(s['memory_used_gb'] for s in samples),
            'memory_avg_percent': round(sum(s['memory_percent'] for s in samples) / len(samples), 1),
            'gpu_memory_peak_gb': max(s['gpu_memory_used_gb'] for s in samples),
            'gpu_memory_peak_percent': max(s['gpu_memory_percent'] for s in samples),
            'gpu_memory_avg_gb': round(sum(s['gpu_memory_used_gb'] for s in samples) / len(samples), 2),
            'gpu_util_peak': max(s['gpu_util_percent'] for s in samples),
            'gpu_util_avg': round(
                sum(s['gpu_util_percent'] for s in samples if s['gpu_util_percent'] >= 0)
                / max(1, sum(1 for s in samples if s['gpu_util_percent'] >= 0)),
                1
            ),
        }
        return peak


def _compute_delta(before, after, label=""):
    """计算两个快照之间的变化量"""
    delta = {
        'label': label,
        'cpu_delta': round(after['cpu_percent'] - before['cpu_percent'], 1),
        'memory_delta_gb': round(after['memory_used_gb'] - before['memory_used_gb'], 2),
        'memory_delta_percent': round(after['memory_percent'] - before['memory_percent'], 1),
        'gpu_memory_delta_gb': round(after['gpu_memory_used_gb'] - before['gpu_memory_used_gb'], 2),
        'gpu_memory_delta_percent': round(after['gpu_memory_percent'] - before['gpu_memory_percent'], 1),
    }
    return delta


def _compute_delta_pre_peak(pre_stats, peak, label=""):
    """计算检测前快照与检测中峰值之间的变化量（展示检测的实际资源需求）"""
    delta = {
        'label': label,
        'cpu_delta': round(peak['cpu_peak'] - pre_stats['cpu_percent'], 1),
        'memory_delta_gb': round(peak['memory_peak_gb'] - pre_stats['memory_used_gb'], 2),
        'memory_delta_percent': round(peak['memory_peak_percent'] - pre_stats['memory_percent'], 1),
        'gpu_memory_delta_gb': round(peak['gpu_memory_peak_gb'] - pre_stats['gpu_memory_used_gb'], 2),
        'gpu_memory_delta_percent': round(peak['gpu_memory_peak_percent'] - pre_stats['gpu_memory_percent'], 1),
    }
    return delta


def format_stats_line(stats, label=""):
    """将系统状态格式化为一行日志文本"""
    parts = [
        f"[{stats['timestamp']}]",
        f"CPU: {stats['cpu_percent']}%",
        f"内存: {stats['memory_used_gb']}/{stats['memory_total_gb']}GB ({stats['memory_percent']}%)",
        f"GPU显存: {stats['gpu_memory_used_gb']}/{stats['gpu_memory_total_gb']}GB ({stats['gpu_memory_percent']}%)",
    ]
    if label:
        parts.insert(1, f"[{label}]")
    return " | ".join(parts)


def format_peak_line(peak, label=""):
    """将峰值统计格式化为日志文本"""
    gpu_util_str = f"{peak['gpu_util_peak']}%" if peak['gpu_util_peak'] >= 0 else "N/A"
    gpu_util_avg_str = f"{peak['gpu_util_avg']}%" if peak['gpu_util_avg'] >= 0 else "N/A"
    lines = [
        f"  采样数: {peak['sample_count']}, 时长: {peak['duration_sec']}s",
        f"  CPU  峰值: {peak['cpu_peak']}%  均值: {peak['cpu_avg']}%",
        f"  内存 峰值: {peak['memory_peak_gb']}GB ({peak['memory_peak_percent']}%)  均值: {peak['memory_avg_percent']}%",
        f"  显存 峰值: {peak['gpu_memory_peak_gb']}GB ({peak['gpu_memory_peak_percent']}%)  均值: {peak['gpu_memory_avg_gb']}GB",
        f"  GPU利用率 峰值: {gpu_util_str}  均值: {gpu_util_avg_str}",
    ]
    if label:
        lines.insert(0, f"  [{label}]")
    return "\n".join(lines)


def format_delta_line(delta):
    """将变化量格式化为日志文本"""
    sign = lambda v: f"+{v}" if v > 0 else f"{v}"
    lines = [
        f"  [{delta['label']}]",
        f"  CPU: {sign(delta['cpu_delta'])}%  "
        f"内存: {sign(delta['memory_delta_gb'])}GB ({sign(delta['memory_delta_percent'])}%)  "
        f"显存: {sign(delta['gpu_memory_delta_gb'])}GB ({sign(delta['gpu_memory_delta_percent'])}%)",
    ]
    return "\n".join(lines)


def write_session_log(session_data):
    """将会话日志写入Logs目录下的txt文件

    session_data 结构:
        startup_stats: 启动时快照
        model_loaded_stats: 模型加载后快照
        detections: 列表，每项包含:
            pre_stats: 检测前快照
            peak: 检测期间峰值
            post_stats: 检测后快照
            delta_pre_post: 检测前后变化量
        shutdown_stats: 关闭时快照
        delta_startup_shutdown: 启动到关闭总变化量
    """
    os.makedirs(LOGS_DIR, exist_ok=True)
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    log_path = os.path.join(LOGS_DIR, f'session_{timestamp}.txt')

    with open(log_path, 'w', encoding='utf-8') as f:
        f.write("=" * 70 + "\n")
        f.write("  MASt3R 场景一致性检测系统 - 运行日志\n")
        f.write("=" * 70 + "\n\n")

        # 1. 启动时资源
        startup = session_data.get('startup_stats')
        if startup:
            f.write("--- ① 服务启动时 ---\n")
            f.write(format_stats_line(startup, "启动") + "\n\n")

        # 2. 模型加载后
        model_loaded = session_data.get('model_loaded_stats')
        if model_loaded and startup:
            f.write("--- ② 模型加载后 ---\n")
            f.write(format_stats_line(model_loaded, "加载后") + "\n")
            delta = _compute_delta(startup, model_loaded, "模型加载增量")
            f.write(format_delta_line(delta) + "\n\n")
        elif model_loaded:
            f.write("--- ② 模型加载后 ---\n")
            f.write(format_stats_line(model_loaded, "加载后") + "\n\n")

        # 3. 每次检测的详细记录
        detections = session_data.get('detections', [])
        if detections:
            f.write("--- ③ 检测过程资源变化 ---\n")
            for i, det in enumerate(detections, 1):
                f.write(f"\n  ── 检测 #{i} ──\n")

                pre = det.get('pre_stats')
                peak = det.get('peak')
                duration = det.get('duration_sec')

                if duration is not None:
                    f.write(f"  耗时: {duration:.1f}s\n")

                if pre:
                    f.write(f"  检测前: {format_stats_line(pre, '检测前').split('|', 1)[1].strip()}\n")

                if peak:
                    f.write(format_peak_line(peak, "检测中峰值") + "\n")

                if pre and peak:
                    delta = _compute_delta_pre_peak(pre, peak, f"检测#{i} 资源需求（前→峰值）")
                    f.write(format_delta_line(delta) + "\n")

            f.write("\n")

        # 4. 服务关闭时
        shutdown = session_data.get('shutdown_stats')
        if shutdown:
            f.write("--- ④ 服务关闭时 ---\n")
            f.write(format_stats_line(shutdown, "关闭") + "\n")

            if startup:
                delta = _compute_delta(startup, shutdown, "全程总变化")
                f.write(format_delta_line(delta) + "\n\n")

        f.write("=" * 70 + "\n")
        f.write(f"  日志生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write("=" * 70 + "\n")

        f.flush()
        os.fsync(f.fileno())

    return log_path