"""
上下文感知 Middleware（消息末尾注入）
==================================

提供 ContextAwareMiddleware，在 LLM 调用前：
1. 触发 ContextMonitor.detect() 检测重复调用
2. 将重复调用提示注入到 request.messages 末尾

说明：
- 工具调用最大步数已由底层主循环处理
- 仅保留重复调用检测（同工具+同参数），引导 LLM 避免无效重复
"""
from typing import Callable, List, Awaitable

from langchain.agents.middleware.types import (
    AgentMiddleware,
    ModelRequest,
    ModelResponse,
)
from langchain_core.messages import SystemMessage  # ⭐ 新增：用于消息末尾注入

from .context_monitor import ContextMonitor, Warning

import logging
logger = logging.getLogger(__name__)


class ContextAwareMiddleware(AgentMiddleware):
    """
    上下文感知 Middleware（重复调用提示）

    职责：
    - 检测重复调用（同工具+同参数）
    - 将提示注入到消息末尾，引导 LLM 避免无效重复
    """

    def __init__(self, context_monitor: ContextMonitor, model=None):
        """
        初始化 Middleware

        Args:
            context_monitor: ContextMonitor 实例（数据源）
            model: 兼容参数（历史版本用于压缩摘要）。当前版本不再依赖该参数。
        """
        self.context_monitor = context_monitor
        self.model = model
        # 注意：压缩/上下文治理已下沉到底层主循环（compaction），此处不再创建 compressor。

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelResponse:
        """
        包装 LLM 调用，注入重复调用提示（末尾注入）

        Args:
            request: LLM 请求对象
            handler: 实际的 LLM 调用处理器

        Returns:
            LLM 响应对象
        """
        try:
            warnings = self.context_monitor.detect()
        except Exception as e:
            logger.error(f"Failed to detect warnings: {e}", exc_info=True)
            warnings = []

        if not warnings:
            return handler(request)

        # 仅处理重复调用警告（duplicate_calls）
        duplicate_warnings = [w for w in warnings if w.warning_type == "duplicate_calls"]
        if not duplicate_warnings:
            return handler(request)

        self._cleanup_old_warnings(request)
        info_msg = self._create_info_message(duplicate_warnings)
        request.messages.append(SystemMessage(content=info_msg))
        logger.info("Injected duplicate-calls hint at message END")

        return handler(request)

    async def awrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
    ) -> ModelResponse:
        """
        异步包装 LLM 调用，注入重复调用提示（末尾注入）

        Args:
            request: LLM 请求对象
            handler: 实际的 LLM 调用处理器（异步）

        Returns:
            LLM 响应对象
        """
        try:
            warnings = self.context_monitor.detect()
        except Exception as e:
            logger.error(f"Failed to detect warnings: {e}", exc_info=True)
            warnings = []

        if not warnings:
            return await handler(request)

        duplicate_warnings = [w for w in warnings if w.warning_type == "duplicate_calls"]
        if not duplicate_warnings:
            return await handler(request)

        self._cleanup_old_warnings(request)
        info_msg = self._create_info_message(duplicate_warnings)
        request.messages.append(SystemMessage(content=info_msg))
        logger.info("Injected duplicate-calls hint at message END")

        return await handler(request)

    def _create_info_message(self, warnings: List[Warning]) -> str:
        """
        重复调用提示（info）

        提示 LLM 检查之前的调用结果，避免无效重复
        """
        messages = [w.message for w in warnings]
        return f"""
ℹ️ 重复调用提示
================================================================================
{chr(10).join(f"   - {msg}" for msg in messages)}
================================================================================
"""

    def _cleanup_old_warnings(self, request: ModelRequest) -> None:
        """
        清理所有旧的警告消息
        
        策略：
        - 分离警告消息和其他消息
        - 清除所有旧的警告消息（不保留）
        - 重新组合消息列表（只保留非警告消息）
        - 新的警告消息会在后续步骤中注入
        
        Args:
            request: LLM 请求对象（会被修改）
        """
        other_messages = []
        removed_count = 0
        
        for msg in request.messages:
            if isinstance(msg, SystemMessage):
                content = str(msg.content)
                if (
                    '⚠️' in content
                    or '🚨' in content
                    or '上下文感知' in content
                    or '重复调用提示' in content
                ):
                    # 跳过所有提示/警告消息（清除）
                    removed_count += 1
                else:
                    # 保留非提示/警告的 SystemMessage
                    other_messages.append(msg)
            else:
                # 保留所有非 SystemMessage 的消息
                other_messages.append(msg)
        
        if removed_count > 0:
            logger.info(f"Cleaned up {removed_count} old warning messages before injecting new ones")
        
        # 重新组合：只保留非警告消息（新的警告消息会在后续步骤中注入）
        request.messages = other_messages


if __name__ == "__main__":
    # 测试代码
    print("=" * 60)
    print("测试：ContextAwareMiddleware")
    print("=" * 60)

    # 创建模拟组件
    from langchain_agentx import LlmCallback
    import time

    # 测试1：创建 Middleware
    print("\n1. 测试创建 ContextAwareMiddleware")
    try:
        callback = LlmCallback()
        monitor = ContextMonitor(callback)
        middleware = ContextAwareMiddleware(monitor)
        print("  ✅ ContextAwareMiddleware 创建成功")
    except Exception as e:
        print(f"  ❌ 创建失败: {e}")
        import traceback
        traceback.print_exc()

    # 测试2：格式化警告
    print("\n2. 测试 _create_info_message")
    dup_warning = Warning(
        warning_id="test_dup",
        warning_type="duplicate_calls",
        severity="info",
        message="检测到重复调用工具 'read_file'（相同参数，3次）。",
        timestamp=time.time()
    )
    info_msg = middleware._create_info_message([dup_warning])
    print(f"  信息消息长度: {len(info_msg)}")

    # 测试3：模拟 wrap_model_call（无警告）
    print("\n3. 测试 wrap_model_call（无警告）")

    class MockModelRequest:
        def __init__(self):
            self.messages = []

    class MockModelResponse:
        pass

    request = MockModelRequest()

    def mock_handler(req):
        return MockModelResponse()

    response = middleware.wrap_model_call(request, mock_handler)
    print(f"  messages 数量: {len(request.messages)}")

    # 测试4：模拟有重复调用警告的场景
    print("\n4. 测试 wrap_model_call（有重复调用警告）")

    from langchain_agentx import PerformanceMonitor, ToolCallRecord

    callback.performance_monitor = PerformanceMonitor()
    current_time = time.time()
    # 添加重复调用记录（同工具+同参数）
    for i in range(2):
        record = ToolCallRecord(
            tool_name="read_file",
            start_time=current_time + i * 0.1,
            end_time=current_time + i * 0.1 + 0.05,
            input_args={"path": "/README.md"},
            output_length=100
        )
        record.finalize(record.end_time, 100)
        callback.performance_monitor.records.append(record)

    monitor.warnings_issued.clear()
    request2 = MockModelRequest()
    response2 = middleware.wrap_model_call(request2, mock_handler)

    has_hint = any("重复调用提示" in str(getattr(m, "content", "")) for m in request2.messages)
    print(f"  是否包含重复调用提示: {has_hint}")

    print("\n" + "=" * 60)
    print("测试完成")
    print("=" * 60)

