# -*- coding: utf-8 -*-
"""
记忆检索工具
Agent 通过此工具按关键词检索跨会话长期记忆。

使用 BM25 + LLM rerank 两阶段召回：
1. BM25Engine 检索 top_k * 5 候选
2. LLMReranker 重排，取 top_k
3. rerank 失败则降级返回 BM25 top_k

构造函数接收 store 与可选 reranker，由 main.py 在初始化时注入。
"""
from typing import Any, Dict, List

from .base import Tool, coerce_int


class MemorySearchTool(Tool):
    """记忆检索工具：按关键词查询跨会话记忆。

    依赖 MemoryStore（必须），可选 LLMReranker（无则纯 BM25 检索）。
    """

    def __init__(self, store=None, reranker=None, bm25_factory=None):
        """
        Args:
            store: MemoryStore 实例
            reranker: 可选 LLMReranker 实例，None 表示禁用 rerank
            bm25_factory: 可选，返回新 BM25Engine 的工厂函数（便于测试 mock）；
                         None 表示用默认 BM25Engine
        """
        self._store = store
        self._reranker = reranker
        self._bm25_factory = bm25_factory

    @property
    def name(self) -> str:
        return "memory_search"

    @property
    def description(self) -> str:
        return (
            "按关键词检索跨会话长期记忆。"
            "参数：query（查询文本）、top_k（返回数量，默认 3，最大 10）。"
            "返回匹配的记忆条目（含 title 与 content）。"
        )

    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "查询文本（支持中英文混合）",
                },
                "top_k": {
                    "type": "integer",
                    "description": "返回的最大结果数（默认 3，最大 10）",
                    "default": 3,
                },
            },
            "required": ["query"],
        }

    def execute(self, query: str, top_k: int = 3) -> Dict[str, Any]:
        """执行记忆检索"""
        if self._store is None:
            return {"success": False, "error": "记忆库未加载"}
        if not query or not query.strip():
            return {"success": False, "error": "query 不能为空"}

        # 规整 top_k：1~10
        k = coerce_int(top_k, default=3)
        k = max(1, min(10, k))

        all_entries = self._store.list_all()
        if not all_entries:
            return {
                "success": True,
                "results": [],
                "message": "记忆库为空",
            }

        # BM25 召回 top_k * 5 候选
        candidate_count = min(len(all_entries), k * 5)
        bm25 = self._create_bm25()
        bm25.index([e.to_index_entry() for e in all_entries])
        bm25_hits = bm25.search(query, top_k=candidate_count)

        if not bm25_hits:
            return {
                "success": True,
                "results": [],
                "message": "未找到相关记忆",
            }

        # 候选条目
        candidates = [all_entries[h.index] for h in bm25_hits]

        # LLM rerank（可选）
        if self._reranker is not None and self._reranker.enabled and len(candidates) > 1:
            try:
                reranked_indices = self._reranker.rerank(query, candidates, top_k=k)
                final_entries = [candidates[i] for i in reranked_indices]
            except Exception:
                # 任何异常都降级到 BM25 顺序
                final_entries = candidates[:k]
        else:
            final_entries = candidates[:k]

        # 序列化结果
        results = []
        for e in final_entries:
            results.append({
                "id": e.id,
                "title": e.title,
                "content": e.content,
                "category": e.category,
                "tags": e.tags,
                "created_at": e.created_at,
            })

        return {
            "success": True,
            "results": results,
            "total": len(results),
        }

    def _create_bm25(self):
        """创建 BM25Engine 实例（支持测试 mock）"""
        if self._bm25_factory is not None:
            return self._bm25_factory()
        from ..memory import BM25Engine
        return BM25Engine()
