"""Multi-KB query manager built on ChromaDB."""

from __future__ import annotations

import logging
import os
import re
from typing import Any, Callable, Dict, Iterable, List, Optional

from .embedding import EmbeddingRuntime, resolve_embedding_runtime

logger = logging.getLogger(__name__)


def _keyword_score(query: str, text: str) -> float:
    query_terms = {item.lower() for item in re.findall(r"\w+", query)}
    if not query_terms:
        return 0.0
    text_terms = {item.lower() for item in re.findall(r"\w+", text)}
    overlap = len(query_terms.intersection(text_terms))
    return overlap / max(1, len(query_terms))


class KBManager:
    """Manage multiple ChromaDB stores and query them via one API."""

    def __init__(self, default_embedding_model: Optional[str] = "all-MiniLM-L6-v2") -> None:
        self._kbs: Dict[str, str] = {}
        self._kb_models: Dict[str, Optional[str]] = {}
        self._clients: Dict[str, Any] = {}
        self._default_model = default_embedding_model
        self._embedding_runtimes: Dict[str, EmbeddingRuntime] = {}

    def add_kb(self, name: str, path: str, embedding_model: Optional[str] = None) -> None:
        if not name:
            raise ValueError("知识库名称不能为空。")
        abs_path = os.path.abspath(path)
        if not os.path.isdir(abs_path):
            raise FileNotFoundError(f"知识库路径不存在: {abs_path}")
        self._kbs[name] = abs_path
        self._kb_models[name] = embedding_model
        self._clients.pop(name, None)
        logger.info("添加知识库: %s -> %s", name, abs_path)

    def remove_kb(self, name: str) -> None:
        self._kbs.pop(name, None)
        self._kb_models.pop(name, None)
        self._clients.pop(name, None)

    def list_kbs(self) -> Dict[str, str]:
        return dict(self._kbs)

    def has_kb(self) -> bool:
        return bool(self._kbs)

    def _validate_kb_name(self, kb_name: str) -> None:
        if kb_name not in self._kbs:
            registered = ", ".join(sorted(self._kbs)) or "无"
            raise ValueError(f"未注册的 kb_name: {kb_name}。当前已注册: {registered}")

    def _get_client(self, kb_name: str):
        if kb_name in self._clients:
            return self._clients[kb_name]
        self._validate_kb_name(kb_name)
        import chromadb

        client = chromadb.PersistentClient(path=self._kbs[kb_name])
        self._clients[kb_name] = client
        return client

    def _get_embedding_runtime(self, kb_name: str) -> EmbeddingRuntime:
        runtime = self._embedding_runtimes.get(kb_name)
        if runtime is not None:
            return runtime
        model = self._kb_models.get(kb_name) or self._default_model
        runtime = resolve_embedding_runtime(model)
        self._embedding_runtimes[kb_name] = runtime
        return runtime

    def query_json(
        self,
        query: str,
        top_k: int = 5,
        kb_name: Optional[str] = None,
        collection_name: Optional[str] = None,
        *,
        hybrid_alpha: float = 0.85,
        reranker: Optional[Callable[[str, List[Dict[str, Any]]], List[Dict[str, Any]]]] = None,
    ) -> Dict[str, Any]:
        if not self._kbs:
            return {
                "query": query,
                "total": 0,
                "results": [],
                "errors": ["未配置任何知识库。请先调用 add_kb()."],
            }
        if kb_name:
            self._validate_kb_name(kb_name)

        target_names = [kb_name] if kb_name else sorted(self._kbs.keys())
        n_results = min(max(int(top_k), 1), 20)
        all_results: list[Dict[str, Any]] = []
        errors: list[str] = []

        for name in target_names:
            try:
                client = self._get_client(name)
                runtime = self._get_embedding_runtime(name)
                collections = client.list_collections()
                if not collections:
                    continue

                for entry in collections:
                    current_collection = entry if isinstance(entry, str) else entry.name
                    if collection_name and current_collection != collection_name:
                        continue

                    collection = client.get_collection(
                        name=current_collection,
                        embedding_function=runtime.embedding_function,
                    )
                    queried = collection.query(query_texts=[query], n_results=n_results)
                    docs = queried.get("documents", [[]])[0]
                    metadatas = queried.get("metadatas", [[]])[0]
                    distances = queried.get("distances", [[]])[0]

                    for doc, metadata, distance in zip(docs, metadatas, distances):
                        vector_score = max(0.0, 1.0 - float(distance))
                        keyword_score = _keyword_score(query, doc or "")
                        score = hybrid_alpha * vector_score + (1.0 - hybrid_alpha) * keyword_score
                        all_results.append(
                            {
                                "score": round(score, 6),
                                "vector_score": round(vector_score, 6),
                                "keyword_score": round(keyword_score, 6),
                                "kb": name,
                                "collection": current_collection,
                                "document": doc or "",
                                "metadata": metadata or {},
                            }
                        )
            except Exception as exc:
                errors.append(f"{name}: {exc}")

        all_results.sort(key=lambda item: item["score"], reverse=True)
        ranked = all_results[:n_results]

        if reranker is not None:
            try:
                ranked = reranker(query, ranked)[:n_results]
            except Exception as exc:
                errors.append(f"reranker 执行失败: {exc}")

        for idx, item in enumerate(ranked, start=1):
            item["rank"] = idx

        return {
            "query": query,
            "total": len(ranked),
            "kbs": target_names,
            "results": ranked,
            "errors": errors,
        }

    def query_many(
        self,
        queries: Iterable[str],
        *,
        top_k: int = 5,
        kb_name: Optional[str] = None,
        collection_name: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        return [
            self.query_json(
                query=item,
                top_k=top_k,
                kb_name=kb_name,
                collection_name=collection_name,
            )
            for item in queries
        ]

    def query(
        self,
        query: str,
        top_k: int = 5,
        kb_name: Optional[str] = None,
        collection_name: Optional[str] = None,
    ) -> str:
        payload = self.query_json(query, top_k=top_k, kb_name=kb_name, collection_name=collection_name)
        if payload.get("errors") and not payload.get("results"):
            escaped_error = self._escape_xml("; ".join(payload["errors"]))
            return f"<error>{escaped_error}</error>"

        xml = [
            f'<knowledge_results query="{self._escape_xml(query)}" total="{payload["total"]}" '
            f'kbs="{",".join(payload.get("kbs", []))}">'
        ]
        for item in payload.get("results", []):
            metadata = item.get("metadata", {}) or {}
            xml.append(
                f'  <result rank="{item.get("rank", 0)}" similarity="{item.get("score", 0.0):.3f}" '
                f'kb="{self._escape_xml(str(item.get("kb", "")))}" '
                f'collection="{self._escape_xml(str(item.get("collection", "")))}">'
            )
            xml.append(f'    <title>{self._escape_xml(str(metadata.get("title", "N/A")))}</title>')
            xml.append(f'    <section>{self._escape_xml(str(metadata.get("section", "N/A")))}</section>')
            xml.append(f'    <source>{self._escape_xml(str(metadata.get("source", "")))}</source>')
            xml.append(f'    <content>{self._escape_xml(str(item.get("document", "")))}</content>')
            xml.append("  </result>")
        xml.append("</knowledge_results>")
        return "\n".join(xml)

    def list_collections(self, kb_name: Optional[str] = None) -> str:
        if not self._kbs:
            return "<knowledge_collections></knowledge_collections>"
        names = [kb_name] if kb_name else sorted(self._kbs.keys())
        if kb_name:
            self._validate_kb_name(kb_name)

        lines = ["<knowledge_collections>"]
        for name in names:
            try:
                client = self._get_client(name)
                collections = client.list_collections()
                for entry in collections:
                    current_collection = entry if isinstance(entry, str) else entry.name
                    count = client.get_collection(name=current_collection).count()
                    lines.append(f'  <collection kb="{name}" name="{current_collection}" count="{count}" />')
            except Exception as exc:
                lines.append(f'  <error kb="{name}">{self._escape_xml(str(exc))}</error>')
        lines.append("</knowledge_collections>")
        return "\n".join(lines)

    def close(self) -> None:
        self._clients.clear()
        for runtime in self._embedding_runtimes.values():
            runtime.close()
        self._embedding_runtimes.clear()

    @staticmethod
    def _escape_xml(text: str) -> str:
        return (
            text.replace("&", "&amp;")
            .replace("<", "&lt;")
            .replace(">", "&gt;")
            .replace('"', "&quot;")
        )


__all__ = ["KBManager"]
