"""
Default KB Tool Runner for Mid Platform.
[AC-MARH-05] Agent 默认 KB 检索工具调用。
[AC-MARH-06] KB 失败时可观测降级。

当 Agent 模式处理开放咨询时，默认尝试调用 KB 检索工具获取事实依据。
"""

import logging
import time
import uuid
from dataclasses import dataclass, field
from typing import Any

from app.models.mid.schemas import ToolCallStatus, ToolCallTrace, ToolType
from app.services.mid.timeout_governor import TimeoutGovernor

logger = logging.getLogger(__name__)

DEFAULT_KB_TOP_K = 5
DEFAULT_KB_TIMEOUT_MS = 2000


@dataclass
class KbToolResult:
    """KB 检索结果。"""
    success: bool
    hits: list[dict[str, Any]] = field(default_factory=list)
    error: str | None = None
    fallback_reason_code: str | None = None
    duration_ms: int = 0
    tool_trace: ToolCallTrace | None = None


@dataclass
class KbToolConfig:
    """KB 工具配置。"""
    enabled: bool = True
    top_k: int = DEFAULT_KB_TOP_K
    timeout_ms: int = DEFAULT_KB_TIMEOUT_MS
    min_score_threshold: float = 0.5


class DefaultKbToolRunner:
    """
    [AC-MARH-05] Agent 默认 KB 检索工具执行器。

    Features:
    - Agent 模式下默认调用 KB 检索
    - 支持超时控制
    - 失败时返回可观测降级信号
    - 记录 kb_tool_called, kb_hit 状态
    """

    def __init__(
        self,
        timeout_governor: TimeoutGovernor | None = None,
        config: KbToolConfig | None = None,
    ):
        self._timeout_governor = timeout_governor or TimeoutGovernor()
        self._config = config or KbToolConfig()
        self._vector_retriever = None

    async def execute(
        self,
        tenant_id: str,
        query: str,
        metadata_filter: dict[str, Any] | None = None,
    ) -> KbToolResult:
        """
        [AC-MARH-05] 执行 KB 检索。

        Args:
            tenant_id: 租户 ID
            query: 检索查询
            metadata_filter: 元数据过滤条件

        Returns:
            KbToolResult 包含检索结果和追踪信息
        """
        if not self._config.enabled:
            logger.info(f"[AC-MARH-05] KB tool disabled for tenant={tenant_id}")
            return KbToolResult(
                success=False,
                fallback_reason_code="KB_DISABLED",
            )

        start_time = time.time()
        tool_trace_id = str(uuid.uuid4())

        logger.info(
            f"[AC-MARH-05] Starting KB retrieval: tenant={tenant_id}, "
            f"query={query[:50]}..., top_k={self._config.top_k}"
        )

        try:
            hits = await self._retrieve_with_timeout(
                tenant_id=tenant_id,
                query=query,
                metadata_filter=metadata_filter,
            )

            duration_ms = int((time.time() - start_time) * 1000)
            kb_hit = len(hits) > 0

            tool_trace = ToolCallTrace(
                tool_name="kb_retrieval",
                tool_type=ToolType.INTERNAL,
                duration_ms=duration_ms,
                status=ToolCallStatus.OK,
                args_digest=f"query={query[:50]}",
                result_digest=f"hits={len(hits)}",
            )

            logger.info(
                f"[AC-MARH-05] KB retrieval completed: tenant={tenant_id}, "
                f"hits={len(hits)}, duration_ms={duration_ms}, kb_hit={kb_hit}"
            )

            return KbToolResult(
                success=True,
                hits=hits,
                duration_ms=duration_ms,
                tool_trace=tool_trace,
            )

        except TimeoutError:
            duration_ms = int((time.time() - start_time) * 1000)
            logger.warning(
                f"[AC-MARH-06] KB retrieval timeout: tenant={tenant_id}, "
                f"duration_ms={duration_ms}"
            )

            tool_trace = ToolCallTrace(
                tool_name="kb_retrieval",
                tool_type=ToolType.INTERNAL,
                duration_ms=duration_ms,
                status=ToolCallStatus.TIMEOUT,
                error_code="KB_TIMEOUT",
            )

            return KbToolResult(
                success=False,
                error="KB retrieval timeout",
                fallback_reason_code="KB_TIMEOUT",
                duration_ms=duration_ms,
                tool_trace=tool_trace,
            )

        except Exception as e:
            duration_ms = int((time.time() - start_time) * 1000)
            logger.error(
                f"[AC-MARH-06] KB retrieval failed: tenant={tenant_id}, "
                f"error={e}"
            )

            tool_trace = ToolCallTrace(
                tool_name="kb_retrieval",
                tool_type=ToolType.INTERNAL,
                duration_ms=duration_ms,
                status=ToolCallStatus.ERROR,
                error_code="KB_ERROR",
            )

            return KbToolResult(
                success=False,
                error=str(e),
                fallback_reason_code="KB_ERROR",
                duration_ms=duration_ms,
                tool_trace=tool_trace,
            )

    async def _retrieve_with_timeout(
        self,
        tenant_id: str,
        query: str,
        metadata_filter: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """带超时控制的检索。"""
        import asyncio

        timeout_seconds = self._config.timeout_ms / 1000.0

        try:
            return await asyncio.wait_for(
                self._do_retrieve(tenant_id, query, metadata_filter),
                timeout=timeout_seconds,
            )
        except asyncio.TimeoutError:
            raise TimeoutError("KB retrieval timeout")

    async def _do_retrieve(
        self,
        tenant_id: str,
        query: str,
        metadata_filter: dict[str, Any] | None = None,
    ) -> list[dict[str, Any]]:
        """执行实际检索。"""
        if self._vector_retriever is None:
            from app.services.retrieval.vector_retriever import get_vector_retriever
            self._vector_retriever = await get_vector_retriever()

        from app.services.retrieval.base import RetrievalContext

        ctx = RetrievalContext(
            tenant_id=tenant_id,
            query=query,
            metadata_filter=metadata_filter,
        )

        result = await self._vector_retriever.retrieve(ctx)

        hits = []
        for hit in result.hits:
            if hit.score >= self._config.min_score_threshold:
                hits.append({
                    "id": hit.metadata.get("chunk_id", str(uuid.uuid4())),
                    "content": hit.text,
                    "score": hit.score,
                    "metadata": hit.metadata,
                })

        return hits[:self._config.top_k]

    def get_config(self) -> KbToolConfig:
        """获取当前配置。"""
        return self._config


_default_kb_tool_runner: DefaultKbToolRunner | None = None


def get_default_kb_tool_runner(
    timeout_governor: TimeoutGovernor | None = None,
    config: KbToolConfig | None = None,
) -> DefaultKbToolRunner:
    """获取或创建 DefaultKbToolRunner 实例。"""
    global _default_kb_tool_runner
    if _default_kb_tool_runner is None:
        _default_kb_tool_runner = DefaultKbToolRunner(
            timeout_governor=timeout_governor,
            config=config,
        )
    return _default_kb_tool_runner
