"""HTTP RAG Service - 通过 HTTP 调用业务系统知识库 API

配置驱动，不直接依赖业务系统代码，通过 base_url + api_key 连接。
"""

import logging
from typing import Any

import httpx

logger = logging.getLogger(__name__)


class HttpRAGService:
    """HTTP 客户端，调用业务系统的知识库检索 API

    适配任意提供以下接口的知识库服务：
    - POST {base_url}/search  → 语义检索
    - POST {base_url}/ingest  → 文档写入（可选）

    典型配置（agentkit.yaml）::

        memory:
          semantic:
            enabled: true
            base_url: "http://localhost:8000/api/knowledge"
            api_key: "${GEO_API_KEY}"
            knowledge_base_ids:
              - "industry-kb-id"
              - "enterprise-kb-id"
            timeout: 30
            contextual_chunking: false
    """

    def __init__(
        self,
        base_url: str,
        api_key: str | None = None,
        knowledge_base_ids: list[str] | None = None,
        timeout: int = 30,
        contextual_chunking: bool = False,
        llm_gateway: Any = None,
    ):
        """
        Args:
            base_url: 知识库 API 基础地址，如 http://localhost:8000/api/knowledge
            api_key: 认证 API Key（放在 Authorization: Bearer 头）
            knowledge_base_ids: 默认检索的知识库 ID 列表
            timeout: HTTP 请求超时秒数
        """
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._knowledge_base_ids = knowledge_base_ids or []
        self._timeout = timeout
        self._client: httpx.AsyncClient | None = None
        self._contextual_chunking = contextual_chunking
        self._llm_gateway = llm_gateway

    def _get_client(self) -> httpx.AsyncClient:
        """懒初始化 httpx 客户端"""
        if self._client is None or self._client.is_closed:
            headers: dict[str, str] = {"Content-Type": "application/json"}
            if self._api_key:
                headers["Authorization"] = f"Bearer {self._api_key}"
            self._client = httpx.AsyncClient(
                base_url=self._base_url,
                headers=headers,
                timeout=self._timeout,
            )
        return self._client

    async def search(
        self,
        query: str,
        knowledge_base_ids: list[str] | None = None,
        top_k: int = 5,
    ) -> list[dict[str, Any]]:
        """语义检索知识库

        Args:
            query: 检索查询
            knowledge_base_ids: 知识库 ID 列表（默认使用配置值）
            top_k: 返回结果数量

        Returns:
            检索结果列表，每项包含 content/score/document_id 等字段
        """
        kb_ids = knowledge_base_ids or self._knowledge_base_ids
        payload = {
            "query": query,
            "knowledge_base_ids": kb_ids,
            "top_k": top_k,
        }

        client = self._get_client()
        try:
            resp = await client.post("/search", json=payload)
            resp.raise_for_status()
            data = resp.json()

            # 兼容两种响应格式：
            # 1. {"results": [...]} — GEO 标准 SearchResponse
            # 2. [...] — 直接返回列表
            if isinstance(data, dict) and "results" in data:
                results = data["results"]
            elif isinstance(data, list):
                results = data
            else:
                logger.warning(f"Unexpected search response format: {type(data)}")
                return []

            # 标准化为 SemanticMemory 期望的格式
            normalized = []
            for r in results:
                if isinstance(r, dict):
                    normalized.append({
                        "id": r.get("chunk_id", r.get("id", "")),
                        "content": r.get("content", ""),
                        "score": float(r.get("score", 0.0)),
                        "source": r.get("source", "rag"),
                        "document_id": r.get("document_id", ""),
                        "document_title": r.get("document_title", ""),
                        "metadata": r.get("metadata", {}),
                    })
            return normalized

        except httpx.HTTPStatusError as e:
            logger.error(f"RAG search HTTP error: {e.response.status_code} — {e.response.text[:200]}")
            return []
        except httpx.RequestError as e:
            logger.error(f"RAG search request error: {e}")
            return []
        except Exception as e:
            logger.error(f"RAG search unexpected error: {e}")
            return []

    async def enhanced_search(
        self,
        query: str,
        knowledge_base_ids: list[str] | None = None,
        top_k: int = 5,
        use_rerank: bool = True,
        use_compression: bool = False,
    ) -> list[dict[str, Any]]:
        """增强语义检索知识库（支持 rerank 和 compression）

        对每个知识库分别调用 /bases/{kb_id}/retrieve 接口，
        合并结果后按 score 降序返回 top_k 条。

        Args:
            query: 检索查询
            knowledge_base_ids: 知识库 ID 列表（默认使用配置值）
            top_k: 返回结果数量
            use_rerank: 是否启用 rerank 重排序
            use_compression: 是否启用上下文压缩

        Returns:
            检索结果列表，每项包含 content/score/document_id 等字段
        """
        kb_ids = knowledge_base_ids or self._knowledge_base_ids
        if not kb_ids:
            return []

        payload = {
            "query": query,
            "top_k": top_k,
            "use_rerank": use_rerank,
            "use_compression": use_compression,
        }

        client = self._get_client()
        all_results: list[dict[str, Any]] = []

        for kb_id in kb_ids:
            try:
                resp = await client.post(f"/bases/{kb_id}/retrieve", json=payload)
                resp.raise_for_status()
                data = resp.json()

                # 兼容两种响应格式
                if isinstance(data, dict) and "results" in data:
                    results = data["results"]
                elif isinstance(data, list):
                    results = data
                else:
                    logger.warning(f"Unexpected enhanced_search response format: {type(data)}")
                    continue

                # 标准化
                for r in results:
                    if isinstance(r, dict):
                        all_results.append({
                            "id": r.get("chunk_id", r.get("id", "")),
                            "content": r.get("content", ""),
                            "score": float(r.get("score", 0.0)),
                            "source": r.get("source", "rag"),
                            "document_id": r.get("document_id", ""),
                            "document_title": r.get("document_title", ""),
                            "knowledge_base_id": kb_id,
                            "metadata": r.get("metadata", {}),
                        })

            except httpx.HTTPStatusError as e:
                if e.response.status_code == 404:
                    # This KB doesn't support enhanced search — fall back to
                    # standard search for THIS KB only, not all KBs.
                    logger.info(
                        f"Enhanced search not available for KB {kb_id}, "
                        f"using standard search"
                    )
                    std_result = await self.search(
                        query, knowledge_base_ids=[kb_id], top_k=top_k
                    )
                    all_results.extend(std_result)
                else:
                    logger.error(
                        f"RAG enhanced_search HTTP error for KB {kb_id}: "
                        f"{e.response.status_code} — {e.response.text[:200]}"
                    )
                    raise
            except httpx.RequestError as e:
                logger.error(f"RAG enhanced_search request error for KB {kb_id}: {e}")
                raise
            except Exception as e:
                logger.error(f"RAG enhanced_search unexpected error for KB {kb_id}: {e}")
                raise

        # 按 score 降序排序，返回 top_k
        all_results.sort(key=lambda x: x["score"], reverse=True)
        return all_results[:top_k]

    async def ingest(
        self,
        key: str,
        value: Any,
        metadata: dict[str, Any] | None = None,
    ) -> dict[str, Any] | None:
        """写入文档到知识库（可选操作）

        When contextual_chunking is enabled and llm_gateway is configured,
        the document content is enhanced with contextual prefixes before ingestion.

        Args:
            key: 文档标题或标识
            value: 文档内容
            metadata: 额外元数据

        Returns:
            写入结果，或 None 表示写入不可用
        """
        kb_ids = self._knowledge_base_ids
        if not kb_ids:
            logger.warning("HttpRAGService.ingest: no knowledge_base_ids configured")
            return None

        content = str(value)

        # Apply contextual chunking if enabled
        if self._contextual_chunking and self._llm_gateway:
            from agentkit.memory.contextual_retrieval import ContextualChunker

            chunker = ContextualChunker(llm_gateway=self._llm_gateway)
            # Simple chunking: split by paragraphs
            raw_chunks = [c.strip() for c in content.split("\n\n") if c.strip()]
            if raw_chunks:
                enhanced = await chunker.enhance_chunks(
                    document=content, chunks=raw_chunks, metadata=metadata
                )
                # Rejoin enhanced chunks
                content = "\n\n".join(chunk.enhanced_content for chunk in enhanced)

        payload = {
            "title": key,
            "content": content,
            "source_type": "text",
            "metadata": metadata or {},
        }

        client = self._get_client()
        try:
            # 写入到第一个配置的知识库
            kb_id = kb_ids[0]
            resp = await client.post(f"/bases/{kb_id}/documents", json=payload)
            resp.raise_for_status()
            return resp.json()
        except httpx.HTTPStatusError as e:
            logger.error(f"RAG ingest HTTP error: {e.response.status_code}")
            return None
        except Exception as e:
            logger.error(f"RAG ingest error: {e}")
            return None

    async def health_check(self) -> bool:
        """检查知识库服务是否可用"""
        client = self._get_client()
        try:
            resp = await client.get("/bases")
            return resp.status_code in (200, 401)  # 401 = 服务在但需认证
        except Exception:
            return False

    async def close(self) -> None:
        """关闭 HTTP 客户端"""
        if self._client and not self._client.is_closed:
            await self._client.aclose()
            self._client = None

    async def __aenter__(self) -> "HttpRAGService":
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.close()
