"""
通用性能监控器 (Performance Monitor)
====================================

此模块作为`langchain_agentx` 的通用能力使用，不依赖本工程其它业务代码。

职责：
- 记录工具调用和 LLM 调用的时序与耗时
- 计算聚合指标（调用链、总延迟、上下文 token 估算等）
- 生成调试 / benchmark 级别的性能报告
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional


class OutputMode(str, Enum):
    """输出模式枚举。"""

    PRODUCTION = "production"   # 生产环境：简洁输出，不展示性能指标
    DEBUG = "debug"             # 调试环境：展示工具调用链路和基础性能指标
    BENCHMARK = "benchmark"     # 性能测试：完整的性能报告


class PerformanceStrategy(str, Enum):
    """策略选择枚举。"""

    ALL_IN_ONE = "all_in_one"       # 一次性加载（小信息量）
    PROGRESSIVE = "progressive"     # 渐进式加载（大信息量）
    HYBRID = "hybrid"               # 混合策略（中等信息量）
    UNKNOWN = "unknown"             # 未知策略


@dataclass
class ToolCallRecord:
    """工具调用记录。"""

    tool_name: str
    start_time: float
    end_time: Optional[float] = None
    input_args: Optional[Dict[str, Any]] = None
    output_length: int = 0
    cost_time_ms: Optional[float] = None
    error: Optional[str] = None

    def finalize(self, end_time: float, output_length: int) -> None:
        """完成记录。"""
        self.end_time = end_time
        self.output_length = output_length
        self.cost_time_ms = round((end_time - self.start_time) * 1000, 2)


@dataclass
class LLMCallRecord:
    """LLM 调用记录。"""

    model_name: str
    start_time: float
    end_time: Optional[float] = None
    prompt_length: int = 0  # 提示词长度（字符数）
    prompt_tokens: Optional[int] = None  # 提示词 tokens（如果可用）
    output_length: int = 0
    output_tokens: Optional[int] = None  # 输出 tokens（如果可用）
    cost_time_ms: Optional[float] = None
    error: Optional[str] = None
    # 工具调用退化统计（当模型把 tool_calls 退化为 <tool_call> 文本时）
    degradation_detected: bool = False
    degraded_tool_names: List[str] = field(default_factory=list)

    def finalize(
        self,
        end_time: float,
        output_length: int,
        prompt_tokens: Optional[int] = None,
        output_tokens: Optional[int] = None,
    ) -> None:
        """完成记录。"""
        self.end_time = end_time
        self.output_length = output_length
        self.prompt_tokens = prompt_tokens
        self.output_tokens = output_tokens
        self.cost_time_ms = round((end_time - self.start_time) * 1000, 2)


@dataclass
class PerformanceMetrics:
    """聚合后的性能指标。"""

    total_tool_calls: int = 0
    total_latency_ms: float = 0.0
    tool_call_chain: List[str] = field(default_factory=list)
    unique_tools: set = field(default_factory=set)
    context_tokens: int = 0  # 上下文占用（tokens，估算）
    strategy: PerformanceStrategy = PerformanceStrategy.UNKNOWN

    # LLM 调用相关
    total_llm_calls: int = 0
    llm_call_chain: List[str] = field(default_factory=list)
    total_llm_latency_ms: float = 0.0

    # 完整调用链（包含 LLM 和 Tool）
    complete_call_chain: List[str] = field(default_factory=list)

    # 阈值
    TOOL_COUNT_THRESHOLD = 15              # 工具调用次数阈值（单次 run 内）
    TOOL_AVG_LATENCY_THRESHOLD_MS = 2000    # 单次工具调用平均耗时目标（ms），超过 2s 提示可能网络/IO 瓶颈
    LLM_AVG_LATENCY_THRESHOLD_MS = 3000     # 单次 LLM 调用平均耗时目标（ms），每个 LLM 瓶颈宜 <3s

    def has_warning(self) -> bool:
        """是否有性能告警。"""
        avg_tool = (
            self.total_latency_ms / self.total_tool_calls
            if self.total_tool_calls > 0
            else 0.0
        )
        avg_llm = (
            self.total_llm_latency_ms / self.total_llm_calls
            if self.total_llm_calls > 0
            else 0.0
        )
        return (
            self.total_tool_calls > self.TOOL_COUNT_THRESHOLD
            or avg_tool > self.TOOL_AVG_LATENCY_THRESHOLD_MS
            or avg_llm > self.LLM_AVG_LATENCY_THRESHOLD_MS
        )

    def get_warnings(self) -> List[str]:
        """获取性能告警列表。"""
        warnings: List[str] = []

        if self.total_tool_calls > self.TOOL_COUNT_THRESHOLD:
            warnings.append(
                f"⚠️ 工具调用次数过多(>{self.TOOL_COUNT_THRESHOLD}次)，建议优化策略"
            )

        if self.total_tool_calls > 0:
            avg_tool = self.total_latency_ms / self.total_tool_calls
            if avg_tool > self.TOOL_AVG_LATENCY_THRESHOLD_MS:
                warnings.append(
                    f"⚠️ 工具单次调用平均耗时过高(>{self.TOOL_AVG_LATENCY_THRESHOLD_MS}ms)，可能存在网络/IO 瓶颈"
                )

        if self.total_llm_calls > 0:
            avg_llm = self.total_llm_latency_ms / self.total_llm_calls
            if avg_llm > self.LLM_AVG_LATENCY_THRESHOLD_MS:
                warnings.append(
                    f"⚠️ 单次 LLM 调用平均耗时过高(>{self.LLM_AVG_LATENCY_THRESHOLD_MS}ms)，建议检查模型选型或上下文大小"
                )

        # 高效策略的正面反馈
        if self.total_tool_calls > 10 and self.total_latency_ms < 500:
            warnings.append("✅ 高效策略：多次调用但延迟可控")

        return warnings


class PerformanceMonitor:
    """
    性能监控器。

    职责：
    1. 收集工具 / LLM 调用数据（通常由 Callback 自动填充）
    2. 计算性能指标（调用次数、延迟、上下文占用）
    3. 检测性能问题（阈值告警）
    4. 生成不同粒度的性能报告
    """

    def __init__(self) -> None:
        self.records: List[ToolCallRecord] = []
        self.llm_records: List[LLMCallRecord] = []
        self.session_start_time: Optional[float] = None
        self.session_end_time: Optional[float] = None
        self.strategy: PerformanceStrategy = PerformanceStrategy.UNKNOWN
        self.strategy_explanation: str = ""

    def start_session(self) -> None:
        """开始监控会话。"""
        self.session_start_time = time.time()

    def end_session(self) -> None:
        """结束监控会话。"""
        self.session_end_time = time.time()

    def add_tool_call(
        self,
        tool_name: str,
        start_time: float,
        input_args: Optional[Dict[str, Any]] = None,
    ) -> ToolCallRecord:
        """添加工具调用记录（开始）。"""
        record = ToolCallRecord(
            tool_name=tool_name,
            start_time=start_time,
            input_args=input_args,
        )
        self.records.append(record)
        return record

    def add_llm_call(
        self,
        model_name: str,
        start_time: float,
        prompt_length: int = 0,
        prompt_tokens: Optional[int] = None,
    ) -> LLMCallRecord:
        """添加 LLM 调用记录（开始）。"""
        record = LLMCallRecord(
            model_name=model_name,
            start_time=start_time,
            prompt_length=prompt_length,
            prompt_tokens=prompt_tokens,
        )
        self.llm_records.append(record)
        return record

    def set_strategy(self, strategy: PerformanceStrategy, explanation: str = "") -> None:
        """设置策略类型与说明。"""
        self.strategy = strategy
        self.strategy_explanation = explanation

    def calculate_metrics(self) -> PerformanceMetrics:
        """计算聚合性能指标。"""
        metrics = PerformanceMetrics()

        # 工具调用统计
        metrics.total_tool_calls = len(self.records)
        metrics.tool_call_chain = [r.tool_name for r in self.records]
        metrics.unique_tools = set(metrics.tool_call_chain)

        total_latency = 0.0
        for record in self.records:
            if record.cost_time_ms:
                total_latency += record.cost_time_ms
        metrics.total_latency_ms = round(total_latency, 2)

        # LLM 调用统计
        metrics.total_llm_calls = len(self.llm_records)
        metrics.llm_call_chain = [r.model_name for r in self.llm_records]

        total_llm_latency = 0.0
        for record in self.llm_records:
            if record.cost_time_ms:
                total_llm_latency += record.cost_time_ms
        metrics.total_llm_latency_ms = round(total_llm_latency, 2)

        # 上下文占用估算
        if self.llm_records:
            last_llm_call = self.llm_records[-1]
            if getattr(last_llm_call, "prompt_tokens", None) is not None:
                metrics.context_tokens = last_llm_call.prompt_tokens  # type: ignore[assignment]
            else:
                prompt_length_chars = getattr(last_llm_call, "prompt_length", 0)
                if prompt_length_chars > 0:
                    metrics.context_tokens = prompt_length_chars // 4
                else:
                    total_output_length = sum(r.output_length for r in self.records)
                    metrics.context_tokens = total_output_length // 4
        else:
            total_output_length = sum(r.output_length for r in self.records)
            metrics.context_tokens = total_output_length // 4

        # 完整调用链（按时间顺序合并 LLM 和 TOOL）
        all_calls: List[tuple[str, str, float]] = []
        for record in self.records:
            all_calls.append(("tool", record.tool_name, record.start_time))
        for record in self.llm_records:
            all_calls.append(("llm", record.model_name, record.start_time))
        all_calls.sort(key=lambda x: x[2])

        metrics.complete_call_chain = [
            f"LLM:{name}" if call_type == "llm" else f"TOOL:{name}"
            for call_type, name, _ in all_calls
        ]

        metrics.strategy = self.strategy

        return metrics

    def generate_report(self, mode: OutputMode = OutputMode.DEBUG) -> str:
        """生成性能报告文本。"""
        if mode == OutputMode.PRODUCTION:
            return ""

        metrics = self.calculate_metrics()

        if mode == OutputMode.DEBUG:
            return self._generate_debug_report(metrics)
        if mode == OutputMode.BENCHMARK:
            return self._generate_benchmark_report(metrics)
        return ""

    def _generate_debug_report(self, metrics: PerformanceMetrics) -> str:
        """生成调试报告（简版）。"""
        lines: List[str] = []

        if metrics.complete_call_chain:
            complete_chain = " → ".join(metrics.complete_call_chain)
            lines.append(f"✅ 完整调用链: {complete_chain}")
        else:
            tool_chain = " → ".join(metrics.tool_call_chain)
            lines.append(f"✅ 工具调用次数: {metrics.total_tool_calls}次 ({tool_chain})")

        warnings = metrics.get_warnings()
        if warnings:
            lines.append("")
            lines.extend(warnings)

        if self.strategy != PerformanceStrategy.UNKNOWN and self.strategy_explanation:
            lines.append(f"📊 策略: {self.strategy_explanation}")

        return "\n".join(lines)

    def _generate_benchmark_report(self, metrics: PerformanceMetrics) -> str:
        """生成完整性能报告。"""
        lines: List[str] = []
        lines.append("=" * 80)
        lines.append("📈 性能报告 (Performance Report)")
        lines.append("=" * 80)
        lines.append("")

        # 完整调用链
        if metrics.complete_call_chain:
            lines.append("## 完整调用链（按时间顺序）")
            lines.append(f"- {' → '.join(metrics.complete_call_chain)}")
            lines.append("")

        # 工具调用统计
        lines.append("## 工具调用统计")
        lines.append(f"- 总次数: {metrics.total_tool_calls}")
        lines.append(f"- 唯一工具数: {len(metrics.unique_tools)}")
        lines.append(f"- 调用链路: {' → '.join(metrics.tool_call_chain)}")
        lines.append("")

        if self.records:
            lines.append("## 工具调用明细")
            for i, record in enumerate(self.records, 1):
                lines.append(f"{i}. {record.tool_name}")
                if record.cost_time_ms:
                    lines.append(f"   - 耗时: {record.cost_time_ms}ms")
                if record.output_length > 0:
                    lines.append(f"   - 输出长度: {record.output_length} 字符")
            lines.append("")

        # LLM 调用统计
        if metrics.total_llm_calls > 0:
            lines.append("## LLM 调用统计")
            lines.append(f"- 总次数: {metrics.total_llm_calls}")
            lines.append(f"- 调用链路: {' → '.join(metrics.llm_call_chain)}")
            lines.append(f"- 总延迟: {metrics.total_llm_latency_ms}ms")
            lines.append("")

            if self.llm_records:
                lines.append("## LLM 调用明细")
                for i, record in enumerate(self.llm_records, 1):
                    lines.append(f"{i}. {record.model_name}")
                    if record.cost_time_ms:
                        lines.append(f"   - 耗时: {record.cost_time_ms}ms")
                    if record.prompt_length > 0:
                        lines.append(f"   - 提示长度: {record.prompt_length} 字符")
                    if record.output_length > 0:
                        lines.append(f"   - 输出长度: {record.output_length} 字符")
                lines.append("")

        # 性能指标
        lines.append("## 性能指标")
        lines.append(f"- 工具总延迟: {metrics.total_latency_ms}ms")
        if metrics.total_tool_calls > 0:
            lines.append(
                f"- 工具平均延迟: {round(metrics.total_latency_ms / metrics.total_tool_calls, 2)}ms"
            )
        if metrics.total_llm_calls > 0:
            lines.append(f"- LLM 总延迟: {metrics.total_llm_latency_ms}ms")
            lines.append(
                f"- LLM 平均延迟: {round(metrics.total_llm_latency_ms / metrics.total_llm_calls, 2)}ms"
            )
        lines.append(f"- 上下文占用: ~{metrics.context_tokens} tokens")
        lines.append("")

        # 策略分析
        if self.strategy != PerformanceStrategy.UNKNOWN:
            lines.append("## 策略分析")
            lines.append(f"- 策略类型: {self.strategy.value}")
            if self.strategy_explanation:
                lines.append(f"- 策略说明: {self.strategy_explanation}")
            lines.append("")

        # 性能告警
        warnings = metrics.get_warnings()
        if warnings:
            lines.append("## 性能告警")
            for warning in warnings:
                lines.append(f"- {warning}")
            lines.append("")

        # 优化建议
        if metrics.has_warning():
            lines.append("## 优化建议")
            if metrics.total_tool_calls > metrics.TOOL_COUNT_THRESHOLD:
                lines.append("- 考虑使用混合策略（先摘要后定向）减少工具调用次数")
                lines.append("- 避免循环调用小工具，优先使用一次性加载")
            if metrics.total_tool_calls > 0:
                avg_tool = metrics.total_latency_ms / metrics.total_tool_calls
                if avg_tool > metrics.TOOL_AVG_LATENCY_THRESHOLD_MS:
                    lines.append("- 优化工具实现或检查网络/IO，降低单次工具平均耗时")
            if metrics.total_llm_calls > 0:
                avg_llm = metrics.total_llm_latency_ms / metrics.total_llm_calls
                if avg_llm > metrics.LLM_AVG_LATENCY_THRESHOLD_MS:
                    lines.append("- 单次 LLM 调用平均过慢，建议检查模型选型或上下文大小")
            lines.append("")

        lines.append("=" * 80)
        return "\n".join(lines)

    def get_tool_stats_summary(self) -> str:
        """获取工具统计摘要（适合作为一行文本插入到输出中）。"""
        metrics = self.calculate_metrics()
        tool_chain = " → ".join(metrics.tool_call_chain)
        return f"{metrics.total_tool_calls}次 ({tool_chain})"

    def get_degradation_stats(self) -> Dict[str, Any]:
        """获取工具调用退化统计（供日志 / 面板展示）。"""
        total_llm_calls = len(self.llm_records)
        degraded_records = [
            r for r in self.llm_records if getattr(r, "degradation_detected", False)
        ]
        degradation_count = len(degraded_records)

        tool_name_distribution: Dict[str, int] = {}
        for r in degraded_records:
            for name in getattr(r, "degraded_tool_names", []) or []:
                tool_name_distribution[name] = tool_name_distribution.get(name, 0) + 1

        avg_prompt_len = 0.0
        if degradation_count > 0:
            avg_prompt_len = (
                sum(r.prompt_length for r in degraded_records) / degradation_count
            )

        rate_percent = (degradation_count / total_llm_calls * 100.0) if total_llm_calls else 0.0

        return {
            "total_llm_calls": total_llm_calls,
            "degradation_count": degradation_count,
            "degradation_rate_percent": round(rate_percent, 2),
            "tool_name_distribution": tool_name_distribution,
            "avg_prompt_length_on_degradation": round(avg_prompt_len, 2),
        }

    def get_degradation_stats_summary(self) -> str:
        """简短退化统计摘要。"""
        stats = self.get_degradation_stats()
        return (
            f"{stats['degradation_count']}/{stats['total_llm_calls']} "
            f"({stats['degradation_rate_percent']}%)"
        )


def explain_strategy_choice(info_size_tokens: int, task_type: str) -> tuple[PerformanceStrategy, str]:
    """
    解释策略选择（复制自原实现，供上层在系统提示词中使用）。

    Args:
        info_size_tokens: 信息量（tokens）
        task_type: 任务类型 ("exploratory" | "targeted" | "global")
    """
    if info_size_tokens < 10_000:
        return (
            PerformanceStrategy.ALL_IN_ONE,
            f"文件大小 {info_size_tokens} tokens < 10K，采用一次性加载",
        )

    if info_size_tokens < 50_000:
        if task_type == "targeted":
            return (
                PerformanceStrategy.HYBRID,
                f"文件大小 {info_size_tokens} tokens，采用混合策略（先摘要后定向）",
            )
        return (
            PerformanceStrategy.ALL_IN_ONE,
            f"文件大小 {info_size_tokens} tokens，需要全局视角，采用一次性加载",
        )

    return (
        PerformanceStrategy.PROGRESSIVE,
        f"文件大小 {info_size_tokens} tokens > 50K，必须分步加载以避免触碰注意力限制",
    )


# 可选：系统提示中嵌入的性能输出模板
PERFORMANCE_OUTPUT_TEMPLATE = """
## 性能报告（输出规范）

完成任务后，请附加性能统计信息：

```
✅ 工具调用次数: {total}次 ({tool_chain})
```

例如：
```
✅ 工具调用次数: 3次 (bound_search_code → grep → read_file)
```

注意：
- 仅在调试模式下展示此信息
- 生产环境不展示性能指标
- 性能测试模式会自动生成完整报告
"""

