"""Milvus 知识向量索引适配器。

MySQL 保存完整文档和发布状态，Milvus 只保存用于语义检索的知识片段及向量。
领域层通过 ``KnowledgeIndex`` 接口使用本类，因此不会依赖 Milvus SDK 细节。
"""

import json
from threading import Lock
from typing import Any, cast

from pymilvus import DataType, MilvusClient  # type: ignore[import-untyped]

from zbt.core.config import Settings
from zbt.core.errors import AppError
from zbt.domains.knowledge.models import KnowledgeChunk, KnowledgeSearchHit
from zbt.infrastructure.embedding.bge_model import load_bge_m3_model


class MilvusKnowledgeIndex:
    """使用本地 BGE-M3 生成稠密向量，并写入 Milvus。"""

    def __init__(self, settings: Settings) -> None:
        self._settings = settings
        self._collection_name = f"{settings.milvus_collection_prefix}knowledge_chunks"
        self._client: MilvusClient | None = None
        self._embedding_model: Any | None = None
        self._client_lock = Lock()
        self._embedding_model_lock = Lock()

    def warmup(self) -> None:
        """预连接 Milvus，并在后台把本地 BGE-M3 权重装入内存。"""
        client = self._get_client()
        self._ensure_collection(client)
        self._get_embedding_model()

    def replace_document(
        self,
        document_id: str,
        chunks: list[KnowledgeChunk],
    ) -> None:
        """重新生成整篇文档的向量，并用当前版本片段替换旧数据。"""

        if not chunks:
            raise AppError("KNOWLEDGE_CHUNKS_EMPTY", "知识文档没有可索引内容", 422)
        try:
            client = self._get_client()
            self._ensure_collection(client)
            vectors = self._embed([chunk.content for chunk in chunks])
            # 先删除旧片段再 upsert，避免文档缩短后多余片段仍被检索到。
            client.delete(
                collection_name=self._collection_name,
                filter=f'document_id == "{self._escape_filter(document_id)}"',
            )
            client.upsert(
                collection_name=self._collection_name,
                data=[
                    {
                        "id": chunk.id,
                        "document_id": chunk.document_id,
                        "document_version": chunk.document_version,
                        "ordinal": chunk.ordinal,
                        "title": chunk.title,
                        "content": chunk.content,
                        "document_type": chunk.document_type,
                        "source_name": chunk.source_name,
                        "product_code": chunk.product_code or "",
                        "vector": vector,
                    }
                    for chunk, vector in zip(chunks, vectors, strict=True)
                ],
            )
            client.flush(collection_name=self._collection_name)
        except AppError:
            raise
        except Exception as error:
            raise AppError(
                "KNOWLEDGE_INDEX_FAILED",
                "知识索引失败，请检查本地Milvus和BGE-M3配置",
                503,
                retryable=True,
            ) from error

    def search(
        self,
        query: str,
        *,
        document_ids: set[str],
        limit: int,
    ) -> list[KnowledgeSearchHit]:
        """只在服务层允许的文档编号范围内执行余弦相似度检索。"""

        if not document_ids:
            return []
        try:
            client = self._get_client()
            self._ensure_collection(client)
            vector = self._embed([query])[0]
            results = client.search(
                collection_name=self._collection_name,
                data=[vector],
                filter=f"document_id in {json.dumps(sorted(document_ids))}",
                limit=limit,
                output_fields=[
                    "document_id",
                    "document_version",
                    "ordinal",
                    "title",
                    "content",
                    "document_type",
                    "source_name",
                    "product_code",
                ],
                search_params={"metric_type": "COSINE"},
            )
            return [
                self._to_search_hit(hit)
                for hit in (results[0] if results else [])
            ]
        except AppError:
            raise
        except Exception as error:
            raise AppError(
                "KNOWLEDGE_SEARCH_FAILED",
                "知识检索失败，请检查本地Milvus和BGE-M3配置",
                503,
                retryable=True,
            ) from error

    def _get_client(self) -> MilvusClient:
        if self._client is None:
            with self._client_lock:
                if self._client is None:
                    self._client = MilvusClient(
                        uri=self._settings.milvus_uri,
                        token=self._settings.milvus_token,
                    )
        return self._client

    def _get_embedding_model(self) -> Any:
        """首次使用时才加载 BGE-M3，并通过锁避免并发请求重复加载模型。"""

        if self._embedding_model is None:
            with self._embedding_model_lock:
                if self._embedding_model is None:
                    self._embedding_model = load_bge_m3_model(
                        self._settings.bge_m3_model_path,
                        self._settings.bge_m3_device,
                    )
        return self._embedding_model

    def _embed(self, texts: list[str]) -> list[list[float]]:
        result = self._get_embedding_model().encode(
            texts,
            return_dense=True,
            return_sparse=False,
            return_colbert_vecs=False,
        )
        dense_vectors = result.get("dense_vecs")
        if dense_vectors is None:
            raise AppError("EMBEDDING_EMPTY", "BGE-M3未返回稠密向量", 503, retryable=True)
        raw_vectors = (
            dense_vectors.tolist() if hasattr(dense_vectors, "tolist") else list(dense_vectors)
        )
        return cast(list[list[float]], raw_vectors)

    def _ensure_collection(self, client: MilvusClient) -> None:
        """确保知识集合和向量索引存在；已存在时不重复创建。"""

        if client.has_collection(self._collection_name):
            return
        schema = MilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
        schema.add_field(
            field_name="id",
            datatype=DataType.VARCHAR,
            is_primary=True,
            max_length=128,
        )
        schema.add_field(
            field_name="document_id",
            datatype=DataType.VARCHAR,
            max_length=64,
        )
        schema.add_field(field_name="document_version", datatype=DataType.INT64)
        schema.add_field(field_name="ordinal", datatype=DataType.INT64)
        schema.add_field(field_name="title", datatype=DataType.VARCHAR, max_length=512)
        schema.add_field(field_name="content", datatype=DataType.VARCHAR, max_length=65535)
        schema.add_field(
            field_name="document_type",
            datatype=DataType.VARCHAR,
            max_length=64,
        )
        schema.add_field(
            field_name="source_name",
            datatype=DataType.VARCHAR,
            max_length=512,
        )
        schema.add_field(
            field_name="product_code",
            datatype=DataType.VARCHAR,
            max_length=128,
        )
        schema.add_field(
            field_name="vector",
            datatype=DataType.FLOAT_VECTOR,
            dim=self._settings.embedding_dimension,
        )
        index_params = MilvusClient.prepare_index_params()
        index_params.add_index(
            field_name="vector",
            index_type="AUTOINDEX",
            metric_type="COSINE",
        )
        client.create_collection(
            collection_name=self._collection_name,
            schema=schema,
            index_params=index_params,
        )

    @staticmethod
    def _to_search_hit(hit: dict[str, Any]) -> KnowledgeSearchHit:
        entity = cast(dict[str, Any], hit.get("entity", {}))
        product_code = str(entity.get("product_code") or "")
        return KnowledgeSearchHit(
            chunk_id=str(hit["id"]),
            document_id=str(entity["document_id"]),
            document_version=int(entity["document_version"]),
            ordinal=int(entity["ordinal"]),
            title=str(entity["title"]),
            content=str(entity["content"]),
            document_type=str(entity["document_type"]),
            source_name=str(entity["source_name"]),
            product_code=product_code or None,
            score=float(hit.get("distance", 0.0)),
        )

    @staticmethod
    def _escape_filter(value: str) -> str:
        return value.replace("\\", "\\\\").replace('"', '\\"')
