"""MemoryRetriever - 混合检索器

并行查询三层记忆，按权重融合排序。
"""

from __future__ import annotations

import asyncio
import logging
import math
from dataclasses import replace
from datetime import datetime
from typing import Any

from agentkit.memory.base import Memory, MemoryItem
from agentkit.memory.working import WorkingMemory
from agentkit.memory.episodic import EpisodicMemory
from agentkit.memory.semantic import SemanticMemory
from agentkit.memory.query_transformer import QueryTransformerBase
from agentkit.memory.rag_loop import RAGSelfCorrectionLoop
from agentkit.memory.relevance_scorer import RelevanceScorer
from agentkit.memory.knowledge_base import KnowledgeBase, QueryResult
from agentkit.memory.multi_source_retriever import MultiSourceRetriever
from agentkit.tools.base import Tool

logger = logging.getLogger(__name__)


def _estimate_tokens(text: str) -> int:
    """Estimate token count for mixed Chinese/English text.

    Chinese characters typically use 1-2 tokens each.
    English words typically use 1 token each.
    """
    cjk_count = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
    non_cjk = text
    for c in text:
        if '\u4e00' <= c <= '\u9fff':
            non_cjk = non_cjk.replace(c, ' ')
    word_count = len(non_cjk.split())
    return cjk_count * 2 + word_count


class MemoryRetriever:
    """混合检索器 - 并行查询三层记忆，按权重融合排序

    检索策略：
    1. 并行查询 Working/Episodic/Semantic 三层
    2. 按权重融合排序（默认 Working 0.2, Episodic 0.4, Semantic 0.4）
    3. 时间衰减：越久远的记忆权重越低
    4. 上下文窗口管理：总 token 不超过预算
    """

    def __init__(
        self,
        working_memory: WorkingMemory | None = None,
        episodic_memory: EpisodicMemory | None = None,
        semantic_memory: SemanticMemory | None = None,
        weights: dict[str, float] | None = None,
        query_transformer: QueryTransformerBase | None = None,
        context_template: str = "structured",
        enable_self_correction: bool = False,
        max_correction_retries: int = 3,
        knowledge_sources: dict[str, KnowledgeBase] | None = None,
    ):
        self._working = working_memory
        self._episodic = episodic_memory
        self._semantic = semantic_memory
        self._weights = weights or {
            "working": 0.2,
            "episodic": 0.4,
            "semantic": 0.4,
        }
        self._query_transformer = query_transformer
        self._context_template = context_template
        self._enable_self_correction = enable_self_correction
        self._correction_loop: RAGSelfCorrectionLoop | None = None
        if enable_self_correction:
            self._correction_loop = RAGSelfCorrectionLoop(
                retriever=self,
                scorer=RelevanceScorer(),
                query_transformer=query_transformer,
                max_retries=max_correction_retries,
            )
        self._multi_source_retriever = MultiSourceRetriever(sources=knowledge_sources)

    async def retrieve(
        self,
        query: str,
        top_k: int = 5,
        token_budget: int = 3000,
        filters: dict[str, Any] | None = None,
        _skip_correction: bool = False,
        sources: list[str] | None = None,
        source_weights: dict[str, float] | None = None,
    ) -> list[MemoryItem]:
        """混合检索三层记忆

        Args:
            query: 检索查询
            top_k: 返回最大结果数
            token_budget: token 预算
            filters: 过滤条件
            _skip_correction: 内部参数，CRAG 循环内部调用时跳过自纠正
            sources: 指定信息源列表，如 ["feishu", "local:合规文档"]。
                     传入时仅从指定源检索，不查三层记忆。
            source_weights: 信息源权重映射，用于多源检索时调整分数
        """
        # 多源检索路径：指定了 sources 时委托给 MultiSourceRetriever
        if sources is not None:
            return await self._retrieve_from_sources(
                query, top_k, token_budget, sources, source_weights
            )
        # Self-correction loop (CRAG)
        if (
            self._enable_self_correction
            and self._correction_loop is not None
            and not _skip_correction
        ):
            result = await self._correction_loop.retrieve_with_correction(
                query, top_k=top_k, token_budget=token_budget, filters=filters
            )
            if result.degraded:
                logger.warning(
                    f"RAG self-correction degraded after {result.total_retries} retries"
                )
            return result.items
        # Query transformation
        if self._query_transformer is not None:
            transformed = await self._query_transformer.transform(query)
            search_query = transformed.main_query
            sub_queries = transformed.sub_queries
        else:
            search_query = query
            sub_queries = []

        # Primary search with main query
        all_items = await self._search_layers(search_query, top_k, filters)

        # Sub-query search in parallel
        if sub_queries:
            sub_tasks = [
                self._search_layers(sq, top_k, filters) for sq in sub_queries
            ]
            sub_results = await asyncio.gather(*sub_tasks, return_exceptions=True)
            for result in sub_results:
                if isinstance(result, Exception):
                    logger.warning(f"Sub-query search failed: {result}")
                    continue
                all_items.extend(result)

        # Deduplicate by key (keep highest score)
        seen: dict[str, MemoryItem] = {}
        for item in all_items:
            if item.key not in seen or item.score > seen[item.key].score:
                seen[item.key] = item
        all_items = list(seen.values())

        # 按分数排序
        all_items.sort(key=lambda x: x.score, reverse=True)

        # Token 预算管理
        selected = []
        total_tokens = 0
        for item in all_items:
            text = str(item.value)
            estimated_tokens = _estimate_tokens(text)
            if total_tokens + estimated_tokens > token_budget:
                continue
            selected.append(item)
            total_tokens += estimated_tokens
            if len(selected) >= top_k:
                break

        return selected

    async def _search_layers(
        self,
        query: str,
        top_k: int = 5,
        filters: dict[str, Any] | None = None,
    ) -> list[MemoryItem]:
        """Search all configured memory layers with a single query"""
        tasks = []
        layer_names = []

        if self._working:
            tasks.append(self._working.search(query, top_k=top_k, filters=filters))
            layer_names.append("working")
        if self._episodic:
            tasks.append(self._episodic.search(query, top_k=top_k, filters=filters))
            layer_names.append("episodic")
        if self._semantic:
            tasks.append(self._semantic.search(query, top_k=top_k, filters=filters))
            layer_names.append("semantic")

        if not tasks:
            return []

        # 并行查询
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # 融合排序
        all_items = []
        for layer_name, result in zip(layer_names, results):
            if isinstance(result, Exception):
                logger.error(f"Memory search failed for {layer_name}: {result}")
                continue
            weight = self._weights.get(layer_name, 0.3)
            for item in result:
                weighted = replace(item, score=item.score * weight)
                all_items.append(weighted)

        return all_items

    async def _retrieve_from_sources(
        self,
        query: str,
        top_k: int = 5,
        token_budget: int = 3000,
        sources: list[str] | None = None,
        source_weights: dict[str, float] | None = None,
    ) -> list[MemoryItem]:
        """从指定信息源检索，将 QueryResult 转换为 MemoryItem

        Args:
            query: 检索查询
            top_k: 返回最大结果数
            token_budget: token 预算
            sources: 信息源名称列表
            source_weights: 信息源权重映射
        """
        kb_results = await self._multi_source_retriever.search(
            query, top_k=top_k, sources=sources, weights=source_weights
        )

        # QueryResult → MemoryItem
        items = []
        for r in kb_results:
            items.append(MemoryItem(
                key=r.source_id,
                value=r.content,
                metadata={
                    **r.metadata,
                    "source": "rag",
                    "source_name": r.source_name,
                    "doc_id": r.doc_id,
                    "document_title": r.title,
                },
                score=r.score,
            ))

        # Token 预算管理
        selected = []
        total_tokens = 0
        for item in items:
            text = str(item.value)
            estimated_tokens = _estimate_tokens(text)
            if total_tokens + estimated_tokens > token_budget:
                continue
            selected.append(item)
            total_tokens += estimated_tokens
            if len(selected) >= top_k:
                break

        return selected

    @property
    def multi_source_retriever(self) -> MultiSourceRetriever:
        """获取多源检索器，用于直接注册/注销信息源"""
        return self._multi_source_retriever

    async def get_context_string(
        self,
        query: str,
        top_k: int = 5,
        token_budget: int = 3000,
    ) -> str:
        """获取格式化的上下文字符串

        根据 context_template 选择输出格式：
        - "structured": 带来源标注的结构化格式
        - "flat": 纯文本拼接（向后兼容）
        """
        items = await self.retrieve(query, top_k, token_budget)

        if not items:
            return ""

        if self._context_template == "flat":
            parts = [str(item.value) for item in items]
            return "\n\n".join(parts)

        # Structured format
        parts: list[str] = []
        for item in items:
            header = self._format_structured_header(item)
            parts.append(f"{header}\n{item.value}")

        result = "\n\n".join(parts)

        # Respect token budget — truncate if formatted output exceeds it
        estimated_tokens = _estimate_tokens(result)
        # Safety limit: also check character count as a ceiling.
        # This handles edge cases like very long unbroken strings.
        max_chars = token_budget * 4
        if estimated_tokens > token_budget or len(result) > max_chars:
            result = result[:max_chars]

        return result

    @staticmethod
    def _format_structured_header(item: MemoryItem) -> str:
        """根据 MemoryItem 的 metadata 生成结构化标题行"""
        source = item.metadata.get("source", "")
        score = item.score

        if source == "rag":
            kb_type = item.metadata.get("kb_type", "知识库")
            document_title = item.metadata.get("document_title", "未知文档")
            return f"### 知识库参考 [来源: {kb_type} | 相关度: {score:.2f} | 文档: {document_title}]"
        elif source == "graph":
            return f"### 知识图谱 [实体: {item.key} | 相关度: {score:.2f}]"
        elif source == "episodic":
            task_type = item.metadata.get("task_type", "未知")
            return f"### 过往经验 [来源: 情景记忆 | 任务类型: {task_type}]"
        elif source == "working":
            return f"### 工作记忆 [键: {item.key}]"
        else:
            return f"### 参考 [来源: {source} | 相关度: {score:.2f}]"

    async def store_episode(
        self, key: str, value: Any, metadata: dict[str, Any] | None = None
    ) -> None:
        """Store an episode into episodic memory if available.

        Public API that delegates to the underlying EpisodicMemory, avoiding
        the need for callers to access the private ``_episodic`` attribute.
        """
        if self._episodic is not None:
            await self._episodic.store(key, value, metadata)

    def create_retrieve_tool(self, max_calls: int = 3) -> Tool | None:
        """Create a retrieve_knowledge tool if semantic memory is configured.

        Returns None if no semantic memory is available (tool not applicable).
        """
        if self._semantic is None:
            return None
        return RetrieveKnowledgeTool(retriever=self, max_calls=max_calls)


class RetrieveKnowledgeTool(Tool):
    """Built-in tool for knowledge base retrieval during ReAct reasoning."""

    def __init__(self, retriever: MemoryRetriever, max_calls: int = 3):
        super().__init__(
            name="retrieve_knowledge",
            description="Search the knowledge base for additional information. Use this tool when you need more context, facts, or details to answer a question accurately.",
            input_schema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The search query to find relevant information in the knowledge base",
                    }
                },
                "required": ["query"],
            },
        )
        self._retriever = retriever
        self._max_calls = max_calls
        self._call_count = 0

    async def execute(self, **kwargs) -> dict:
        query = kwargs.get("query", "")
        if not query:
            return {"error": "query is required", "results": []}

        if self._call_count >= self._max_calls:
            return {"error": f"Maximum retrieval calls ({self._max_calls}) reached", "results": []}

        self._call_count += 1

        try:
            items = await self._retriever.retrieve(query, top_k=5)
            results = []
            for item in items:
                results.append({
                    "content": item.value,
                    "score": item.score,
                    "source": item.metadata.get("source", "unknown"),
                    "document_title": item.metadata.get("document_title", ""),
                })
            return {"query": query, "results": results, "call_count": self._call_count}
        except Exception as e:
            return {"error": str(e), "results": []}
