"""循环检测中间件 — 替代 DeerFlow LoopDetectionMiddleware"""

import hashlib
import logging
import threading
from collections import OrderedDict
from typing import Optional

from app.agent.messages import AIMessage, SystemMessage

logger = logging.getLogger(__name__)


class LoopDetector:
    """检测并阻止重复工具调用循环"""

    def __init__(self, warn_threshold: int = 3, hard_limit: int = 5,
                 window_size: int = 20):
        self.warn_threshold = warn_threshold
        self.hard_limit = hard_limit
        self.window_size = window_size

        # Per-thread tracking: {thread_id: OrderedDict[hash, count]}
        self._thread_history: dict[str, OrderedDict] = {}
        self._lock = threading.Lock()

    def _hash_tool_call(self, name: str, args: dict) -> str:
        """Generate MD5 hash for a tool call signature"""
        raw = f"{name}:{sorted(args.items())}"
        return hashlib.md5(raw.encode()).hexdigest()

    def _get_thread_history(self, thread_id: str) -> OrderedDict:
        """Get or create history for a thread (LRU bounded)"""
        if thread_id not in self._thread_history:
            self._thread_history[thread_id] = OrderedDict()
        return self._thread_history[thread_id]

    def _record_and_check(self, thread_id: str, call_hash: str) -> int:
        """Record a tool call hash and return its repeat count"""
        with self._lock:
            history = self._get_thread_history(thread_id)

            if call_hash in history:
                history.move_to_end(call_hash)
                history[call_hash] += 1
            else:
                history[call_hash] = 1

            # LRU eviction
            while len(history) > self.window_size:
                history.popitem(last=False)

            return history[call_hash]

    def _cleanup_thread(self, thread_id: str):
        """Remove thread history to prevent memory leak"""
        with self._lock:
            self._thread_history.pop(thread_id, None)

    def check(self, state: dict, thread_id: str = "default") -> dict | None:
        """检测循环，返回修改后的 state 或 None"""
        messages = state.get("messages", [])
        if not messages:
            return None

        last_msg = messages[-1]
        if not isinstance(last_msg, AIMessage) or not last_msg.tool_calls:
            return None

        # Check each tool call for loops
        max_repeat = 0
        for tc in last_msg.tool_calls:
            call_hash = self._hash_tool_call(tc["name"], tc.get("args", {}))
            repeat_count = self._record_and_check(thread_id, call_hash)
            max_repeat = max(max_repeat, repeat_count)

        # Hard limit: strip tool calls, force text-only response
        if max_repeat >= self.hard_limit:
            logger.warning(
                f"Loop hard limit reached (count={max_repeat}), "
                f"stripping tool calls for thread={thread_id}"
            )
            # Replace AIMessage with text-only version + warning
            warning_text = (
                "检测到重复操作，已自动终止工具调用循环。"
                "请基于已有信息直接回答用户问题。"
            )
            new_ai = AIMessage(content=warning_text)
            new_messages = messages[:-1] + [new_ai]
            self._cleanup_thread(thread_id)
            return {**state, "messages": new_messages}

        # Warn threshold: inject warning SystemMessage
        if max_repeat >= self.warn_threshold:
            logger.info(
                f"Loop warning (count={max_repeat}) for thread={thread_id}"
            )
            warning = SystemMessage(
                content="⚠️ 检测到可能重复的操作，请考虑换一种方式完成任务。"
            )
            new_messages = messages + [warning]
            return {**state, "messages": new_messages}

        return None
