"""性能分析器 - 独立组件，与 FactorEngine 解耦

支持上下文管理器用法：
    with profiler.measure("factor_name"):
        df = df.with_columns(expr)  # 触发真实计算
"""

import os
import time
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Dict, Iterator, List, Optional, Tuple


@dataclass
class ProfileRecord:
    """性能记录数据结构"""

    time_ms: float
    # 预留扩展字段
    # memory_mb: float = 0.0
    # call_count: int = 1


class FactorTimeoutError(TimeoutError):
    """因子计算超时异常"""

    def __init__(self, factor_name: str, timeout_seconds: float):
        self.factor_name = factor_name
        self.timeout_seconds = timeout_seconds
        super().__init__(
            f"因子 '{factor_name}' 计算超时（超过 {timeout_seconds:.0f} 秒）"
        )


class PerformanceProfiler:
    """性能分析器 - 独立组件，与 FactorEngine 解耦

    支持上下文管理器用法：
        with profiler.measure("factor_name"):
            df = df.with_columns(expr)  # 触发真实计算
    """

    # 默认超时时间：3分钟（180秒）
    DEFAULT_TIMEOUT_SECONDS = 180

    def __init__(
        self,
        enabled: bool = False,
        timeout_seconds: Optional[float] = None,
        immediate_print: bool = False,
    ):
        self.enabled = enabled
        self.timeout_seconds = timeout_seconds or self.DEFAULT_TIMEOUT_SECONDS
        self.immediate_print = immediate_print
        self.records: Dict[str, ProfileRecord] = {}
        self._context_stack: List[str] = []  # 支持嵌套计时

    @contextmanager
    def measure(self, name: str) -> Iterator[None]:
        """上下文管理器：安全计时，自动处理异常，支持超时检测

        Args:
            name: 计时任务名称（通常是因子名称）

        Yields:
            None

        Raises:
            FactorTimeoutError: 当计算时间超过超时阈值时
        """
        if not self.enabled:
            yield
            return

        start_time = time.perf_counter()
        self._context_stack.append(name)

        try:
            yield
        finally:
            elapsed = time.perf_counter() - start_time
            self._context_stack.pop()

            # 更加简洁和安全的更新方式
            if name not in self.records:
                self.records[name] = ProfileRecord(time_ms=0.0)

            self.records[name].time_ms += elapsed

            # 实时打印当前因子性能信息
            if self.immediate_print:
                self._print_single_factor(name, elapsed)

            # 检查是否超时
            if elapsed > self.timeout_seconds:
                self._handle_timeout(name, elapsed)

    def _print_single_factor(self, name: str, elapsed_seconds: float) -> None:
        """打印单个因子的性能信息

        Args:
            name: 因子名称
            elapsed_seconds: 耗时（秒）
        """
        elapsed_ms = elapsed_seconds * 1000
        print(f"[Performance] {name}: {elapsed_ms:.2f}ms")

    def _handle_timeout(self, name: str, elapsed_seconds: float) -> None:
        """处理超时情况

        在 debug 模式下，如果因子计算超过阈值，强制中止进程。

        Args:
            name: 因子名称
            elapsed_seconds: 实际耗时（秒）
        """
        error_msg = (
            f"[ERROR] 因子 '{name}' 计算超时！"
            f" 耗时: {elapsed_seconds:.1f}s，阈值: {self.timeout_seconds:.0f}s"
        )
        print(error_msg, flush=True)
        print(f"[ERROR] 强制中止进程", flush=True)

        # 强制退出进程
        os._exit(1)

    def get_report(self) -> Dict[str, Any]:
        """生成性能报告

        Returns:
            包含性能统计数据的字典：
            - enabled: 是否启用了性能监控
            - total_time_ms: 总计算时间（毫秒）
            - factor_count: 因子数量
            - avg_time_ms: 平均耗时（毫秒）
            - slowest_factors: 最慢的5个因子列表
            - all_records: 所有因子的详细记录
        """
        if not self.records:
            return {}

        times = [r.time_ms for r in self.records.values()]
        sorted_records = sorted(
            self.records.items(), key=lambda x: x[1].time_ms, reverse=True
        )

        return {
            "enabled": self.enabled,
            "total_time_ms": sum(times) * 1000,
            "factor_count": len(self.records),
            "avg_time_ms": (sum(times) / len(times)) * 1000,
            "slowest_factors": [
                (name, record.time_ms * 1000) for name, record in sorted_records[:5]
            ],
            "all_records": {
                name: {"time_ms": record.time_ms * 1000}
                for name, record in self.records.items()
            },
        }

    def print_report(self) -> None:
        """打印格式化的性能报告"""
        if not self.enabled:
            print("[Performance] 性能监控未启用")
            return

        report = self.get_report()
        if not report:
            print("[Performance] 无性能数据")
            return

        print("\n" + "=" * 60)
        print("Factor Engine 性能报告")
        print("=" * 60)
        print(f"总计算时间: {report['total_time_ms']:.2f}ms")
        print(f"因子数量: {report['factor_count']}")
        print(f"平均耗时: {report['avg_time_ms']:.2f}ms")
        print(f"\n最慢的 5 个因子:")
        for i, (name, time_ms) in enumerate(report["slowest_factors"], 1):
            print(f"  {i}. {name}: {time_ms:.2f}ms")
        print("=" * 60)

    def reset(self) -> None:
        """重置性能数据"""
        self.records.clear()
        self._context_stack.clear()

    def get_factor_time(self, name: str) -> Optional[float]:
        """获取指定因子的计算时间

        Args:
            name: 因子名称

        Returns:
            计算时间（秒），如果未找到则返回 None
        """
        if name in self.records:
            return self.records[name].time_ms
        return None
