"""命中处理 — 模型优化模式（LLM 生成回答）vs 直接回答模式（返回匹配段落）。

- model_opt: 调用 LLM Gateway，基于检索结果生成自然语言回答
- direct:    直接拼接匹配段落返回（零 LLM 调用，零额外延迟）

模式选择由 KB 设置 ``default_hit_processing`` 决定，Agent 运行时可通过
``RetrievalRequest.hit_processing_mode`` 覆盖。
"""

from __future__ import annotations

import hashlib
import logging
from typing import Any

from pydantic import BaseModel, ConfigDict

from agentkit.rag_platform.models import QueryResult

logger = logging.getLogger(__name__)

HIT_PROCESSING_MODEL_OPT = "model_opt"
HIT_PROCESSING_DIRECT = "direct"

# ponytail: 硬编码默认模型 — 升级路径为从 KB 设置 / agentkit.yaml 读取
_DEFAULT_LLM_MODEL = "gpt-4o-mini"

# RAG prompt 模板 — 指示 LLM 基于参考资料回答
_RAG_PROMPT_TEMPLATE = (
    "你是一个知识库问答助手。请根据以下检索到的参考资料回答用户问题。\n\n"
    "参考资料：\n{context}\n\n"
    "用户问题：{query}\n\n"
    "请基于参考资料回答，如果资料中没有相关信息请说明。"
)


class HitProcessingResult(BaseModel):
    """命中处理结果。"""

    model_config = ConfigDict()

    mode: str  # "model_opt" | "direct"
    answer: str  # LLM 生成的回答或直接拼接的段落
    sources: list[QueryResult]  # 检索到的来源
    cached: bool = False  # 是否命中缓存


class HitProcessor:
    """命中处理器 — 根据 mode 处理检索结果。

    Args:
        llm_gateway: LLM 网关实例（model_opt 模式使用），需支持 ``async chat(messages, model)`` 接口
        cache_enabled: 是否启用命中处理缓存（默认 True）
        model: LLM 模型名（默认 gpt-4o-mini）
    """

    def __init__(
        self,
        llm_gateway: Any = None,
        cache_enabled: bool = True,
        model: str = _DEFAULT_LLM_MODEL,
    ) -> None:
        self._llm = llm_gateway
        self._cache_enabled = cache_enabled
        self._model = model
        # ponytail: 进程内 dict 缓存 — 无 TTL，无容量上限
        # 升级路径：迁移到 Redis，带 TTL + LRU 淘汰
        self._cache: dict[str, HitProcessingResult] = {}

    async def process(
        self,
        query: str,
        results: list[QueryResult],
        mode: str = HIT_PROCESSING_MODEL_OPT,
    ) -> HitProcessingResult:
        """处理检索结果。

        Args:
            query: 用户查询文本
            results: 检索结果列表
            mode: 命中处理模式 — "model_opt"（LLM 生成）或 "direct"（直接返回段落）

        Returns:
            HitProcessingResult 包含回答、来源和缓存标记
        """
        if not results:
            return HitProcessingResult(
                mode=mode, answer="未找到相关内容。", sources=[], cached=False
            )

        # 缓存检查
        cache_key: str | None = None
        if self._cache_enabled:
            cache_key = self._cache_key(query, results, mode)
            cached = self._cache.get(cache_key)
            if cached is not None:
                return cached.model_copy(update={"cached": True})

        if mode == HIT_PROCESSING_DIRECT:
            result = self._process_direct(results)
        else:
            result = await self._process_model_opt(query, results)

        # 写入缓存
        if cache_key is not None:
            self._cache[cache_key] = result

        return result

    def _process_direct(self, results: list[QueryResult]) -> HitProcessingResult:
        """直接回答模式 — 拼接匹配段落，按分数降序排列。"""
        sorted_results = sorted(results, key=lambda r: r.score, reverse=True)
        parts: list[str] = []
        for i, r in enumerate(sorted_results, start=1):
            parts.append(f"[{i}] {r.content}")
        answer = "\n\n".join(parts)
        return HitProcessingResult(
            mode=HIT_PROCESSING_DIRECT,
            answer=answer,
            sources=sorted_results,
            cached=False,
        )

    async def _process_model_opt(
        self,
        query: str,
        results: list[QueryResult],
    ) -> HitProcessingResult:
        """模型优化模式 — LLM 基于检索结果生成回答。

        无 LLM 网关或调用失败时降级为 direct 模式（不丢失检索结果）。
        """
        if self._llm is None:
            logger.warning("No LLM gateway configured, falling back to direct mode")
            return self._process_direct(results)

        # 构建 context — 按分数降序
        sorted_results = sorted(results, key=lambda r: r.score, reverse=True)
        context_parts: list[str] = []
        for i, r in enumerate(sorted_results, start=1):
            context_parts.append(f"[{i}] {r.content}")
        context = "\n\n".join(context_parts)

        prompt = _RAG_PROMPT_TEMPLATE.format(context=context, query=query)
        messages = [
            {"role": "system", "content": "你是一个知识库问答助手，请基于参考资料回答问题。"},
            {"role": "user", "content": prompt},
        ]

        try:
            response = await self._llm.chat(messages=messages, model=self._model)
            answer = response.content
        except Exception as e:
            logger.warning("LLM call failed, falling back to direct mode: %s", e)
            return self._process_direct(results)

        return HitProcessingResult(
            mode=HIT_PROCESSING_MODEL_OPT,
            answer=answer,
            sources=sorted_results,
            cached=False,
        )

    @staticmethod
    def _cache_key(
        query: str,
        results: list[QueryResult],
        mode: str,
    ) -> str:
        """生成缓存键 — 基于 mode + query + chunk_ids 的 SHA256 哈希。"""
        chunk_ids = ",".join(sorted(r.chunk_id for r in results))
        raw = f"{mode}:{query}:{chunk_ids}"
        return hashlib.sha256(raw.encode()).hexdigest()


__all__ = [
    "HIT_PROCESSING_DIRECT",
    "HIT_PROCESSING_MODEL_OPT",
    "HitProcessingResult",
    "HitProcessor",
]
