"""Milvus 访问层：无状态 Store + 短生命周期 gRPC 连接（避免长期持有失效 channel）。"""
from __future__ import annotations

import os
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Callable, Iterator, TypeVar

from pymilvus import AnnSearchRequest, DataType, MilvusClient, RRFRanker, Function, FunctionType

QUERY_MAX_LIMIT = 16384
T = TypeVar("T")


@dataclass(frozen=True)
class MilvusSettings:
    host: str
    port: str
    collection_name: str
    uri: str
    timeout: float

    @classmethod
    def from_env(cls) -> "MilvusSettings":
        host = os.getenv("MILVUS_HOST", "localhost")
        port = os.getenv("MILVUS_PORT", "19530")
        collection = os.getenv("MILVUS_COLLECTION", "kb_documents")
        timeout = float(os.getenv("MILVUS_TIMEOUT", "30"))
        return cls(
            host=host,
            port=port,
            collection_name=collection,
            uri=f"http://{host}:{port}",
            timeout=timeout,
        )


@contextmanager
def milvus_client_session(settings: MilvusSettings | None = None) -> Iterator[MilvusClient]:
    """一次 RPC 会话：创建连接，用完后关闭，不缓存 gRPC channel。"""
    cfg = settings or MilvusSettings.from_env()
    client = MilvusClient(uri=cfg.uri, timeout=cfg.timeout)
    try:
        yield client
    finally:
        client.close()


def _normalize_filter(filter_expr: str) -> str:
    return filter_expr.strip() if filter_expr.strip() else "id >= 0"


class MilvusStore:
    """Milvus 集合读写；本身不持有连接，所有 IO 经 milvus_client_session。"""

    def __init__(self, settings: MilvusSettings | None = None):
        self._settings = settings or MilvusSettings.from_env()

    @property
    def collection_name(self) -> str:
        return self._settings.collection_name

    def _run(self, operation: Callable[[MilvusClient], T]) -> T:
        with milvus_client_session(self._settings) as client:
            return operation(client)

    @staticmethod
    def _ensure_loaded(client: MilvusClient, collection_name: str) -> None:
        """幂等加载 collection。

        Milvus 服务重启后 collection 默认 released，直接 search/query 会抛
        'call load() before search'。已 loaded 的 collection 重复调用会被
        服务端识别为无操作。collection 不存在 / 正在加载等异常忽略，真正
        的状态错误会在后续操作时抛出明确信息。
        """
        try:
            client.load_collection(collection_name)
        except Exception:
            pass

    def _run_read(self, operation: Callable[[MilvusClient], T]) -> T:
        """读取类操作统一入口：先幂等 load collection，再执行。"""
        with milvus_client_session(self._settings) as client:
            MilvusStore._ensure_loaded(client, self.collection_name)
            return operation(client)

    @contextmanager
    def session(self) -> Iterator[MilvusClient]:
        """同一业务流（如整次上传）内复用一条连接，用毕即关。"""
        with milvus_client_session(self._settings) as client:
            yield client

    @staticmethod
    def ensure_collection(client: MilvusClient, collection_name: str, dense_dim: int) -> None:
        # 不论 collection 是否已存在，最终都要做索引对账：
        # - 历史已存在的 collection 可能因旧版异常处理被静默吞掉索引创建失败
        # - 新建过程中若 create_collection 部分成功（collection 已建、索引未建全）也需要补救
        if not client.has_collection(collection_name):
            schema = client.create_schema(auto_id=True, enable_dynamic_field=True)
            schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)
            schema.add_field("dense_embedding", DataType.FLOAT_VECTOR, dim=dense_dim)
            schema.add_field("sparse_embedding", DataType.SPARSE_FLOAT_VECTOR)
            schema.add_field(
                "text",
                DataType.VARCHAR,
                max_length=65535,
                enable_analyzer=True,
                analyzer_params={"type": "standard"},
                enable_match=True,
            )
            schema.add_field("document_id", DataType.INT64)
            schema.add_field("filename", DataType.VARCHAR, max_length=255)
            schema.add_field("file_type", DataType.VARCHAR, max_length=50)
            schema.add_field("file_path", DataType.VARCHAR, max_length=1024)
            schema.add_field("page_number", DataType.INT64)
            schema.add_field("chunk_idx", DataType.INT64)
            schema.add_field("chunk_id", DataType.VARCHAR, max_length=512)
            schema.add_field("parent_chunk_id", DataType.VARCHAR, max_length=512)
            schema.add_field("root_chunk_id", DataType.VARCHAR, max_length=512)
            schema.add_field("chunk_level", DataType.INT64)

            bm25_function = Function(
                name="text_bm25_emb",
                function_type=FunctionType.BM25,
                input_field_names=["text"],
                output_field_names=["sparse_embedding"],
            )
            schema.add_function(bm25_function)

            index_params = client.prepare_index_params()
            index_params.add_index(
                field_name="dense_embedding",
                index_type="HNSW",
                metric_type="IP",
                params={"M": 16, "efConstruction": 256},
            )
            index_params.add_index(
                field_name="sparse_embedding",
                index_type="SPARSE_INVERTED_INDEX",
                metric_type="BM25",
                params={"drop_ratio_build": 0.2},
            )
            try:
                client.create_collection(
                    collection_name=collection_name,
                    schema=schema,
                    index_params=index_params,
                )
            except Exception as e:
                # 仅当 collection 已实际建好时才吞掉异常（典型场景：Milvus Lite on Windows
                # 重命名 manifest.json 时偶发 WinError 183）。此时索引是否齐全会由后续
                # _ensure_indexes 对账补建，避免历史那种"半建状态被永久隐藏"的缺陷。
                # collection 未建成功时必须 raise，不能掩盖真正的创建失败。
                if not client.has_collection(collection_name):
                    raise

        MilvusStore._ensure_indexes(client, collection_name)
        MilvusStore._ensure_loaded(client, collection_name)

    @staticmethod
    def _ensure_indexes(client: MilvusClient, collection_name: str) -> None:
        """索引对账：dense_embedding 与 sparse_embedding 任一缺失则单独补建。

        默认情况下索引名等于字段名（create_index 未显式指定 index_name），
        因此直接用字段名判断是否已存在。
        """
        try:
            existing = set(client.list_indexes(collection_name))
        except Exception:
            existing = set()

        if "dense_embedding" not in existing:
            dense_params = client.prepare_index_params()
            dense_params.add_index(
                field_name="dense_embedding",
                index_type="HNSW",
                metric_type="IP",
                params={"M": 16, "efConstruction": 256},
            )
            client.create_index(collection_name=collection_name, index_params=dense_params)

        if "sparse_embedding" not in existing:
            sparse_params = client.prepare_index_params()
            sparse_params.add_index(
                field_name="sparse_embedding",
                index_type="SPARSE_INVERTED_INDEX",
                metric_type="BM25",
                params={"drop_ratio_build": 0.2},
            )
            client.create_index(collection_name=collection_name, index_params=sparse_params)

    def init_collection(self, dense_dim: int | None = None) -> None:
        if dense_dim is None:
            dense_dim = int(os.getenv("DENSE_EMBEDDING_DIM", "1024"))

        def _init(client: MilvusClient) -> None:
            self.ensure_collection(client, self.collection_name, dense_dim)

        self._run(_init)

    def insert(self, data: list[dict]):
        return self._run(lambda client: client.insert(self.collection_name, data))

    def delete_by_document(self, document_id: int) -> int:
        def _delete(client: MilvusClient) -> int:
            result = client.delete(
                collection_name=self.collection_name,
                filter=f"document_id == {document_id}",
            )
            return getattr(result, "delete_count", len(result)) if result else 0

        return self._run(_delete)

    def delete_by_ids(self, ids: list[int | str]) -> int:
        if not ids:
            return 0

        def _delete(client: MilvusClient) -> int:
            result = client.delete(collection_name=self.collection_name, ids=ids)
            return getattr(result, "delete_count", len(result)) if result else 0

        return self._run(_delete)

    def query(
        self,
        filter_expr: str = "",
        output_fields: list[str] | None = None,
        limit: int = 10000,
        offset: int = 0,
    ):
        expr = _normalize_filter(filter_expr)
        fields = output_fields or ["filename", "file_type"]

        def _query(client: MilvusClient):
            return client.query(
                collection_name=self.collection_name,
                filter=expr,
                output_fields=fields,
                limit=min(limit, QUERY_MAX_LIMIT),
                offset=offset,
            )

        return self._run_read(_query)

    def query_all(self, filter_expr: str = "", output_fields: list[str] | None = None) -> list:
        """分页拉取；单次 session 内完成，避免每页新建连接。"""
        fields = output_fields or ["filename", "file_type"]
        expr = _normalize_filter(filter_expr)

        def _query_all(client: MilvusClient) -> list:
            out: list = []
            offset = 0
            while True:
                batch = client.query(
                    collection_name=self.collection_name,
                    filter=expr,
                    output_fields=fields,
                    limit=QUERY_MAX_LIMIT,
                    offset=offset,
                )
                if not batch:
                    break
                out.extend(batch)
                if len(batch) < QUERY_MAX_LIMIT:
                    break
                offset += len(batch)
            return out

        return self._run_read(_query_all)

    def get_chunks_by_ids(self, chunk_ids: list[str]) -> list[dict]:
        ids = [item for item in chunk_ids if item]
        if not ids:
            return []
        quoted_ids = ", ".join(f'"{item}"' for item in ids)
        return self.query(
            filter_expr=f"chunk_id in [{quoted_ids}]",
            output_fields=[
                "text",
                "filename",
                "file_type",
                "page_number",
                "chunk_id",
                "parent_chunk_id",
                "root_chunk_id",
                "chunk_level",
                "chunk_idx",
            ],
            limit=len(ids),
        )

    def hybrid_retrieve(
        self,
        dense_embedding: list[float],
        query: str,
        top_k: int = 5,
        rrf_k: int = 60,
        filter_expr: str = "",
    ) -> list[dict]:
        output_fields = [
            "text",
            "document_id",
            "category_id",
            "filename",
            "file_type",
            "page_number",
            "chunk_id",
            "parent_chunk_id",
            "root_chunk_id",
            "chunk_level",
            "chunk_idx",
        ]
        dense_search = AnnSearchRequest(
            data=[dense_embedding],
            anns_field="dense_embedding",
            param={"metric_type": "IP", "params": {"ef": 64}},
            limit=top_k * 2,
            expr=filter_expr,
        )
        sparse_search = AnnSearchRequest(
            data=[query],
            anns_field="sparse_embedding",
            param={"metric_type": "BM25", "params": {"drop_ratio_search": 0.2}},
            limit=top_k * 2,
            expr=filter_expr,
        )
        reranker = RRFRanker(k=rrf_k)

        def _search(client: MilvusClient):
            return client.hybrid_search(
                collection_name=self.collection_name,
                reqs=[dense_search, sparse_search],
                ranker=reranker,
                limit=top_k,
                output_fields=output_fields,
            )

        results = self._run_read(_search)
        formatted_results = []
        for hits in results:
            for hit in hits:
                entity = hit.get("entity", hit)
                formatted_results.append({
                    "id": hit.get("id"),
                    "text": entity.get("text", ""),
                    "document_id": entity.get("document_id"),
                    "category_id": entity.get("category_id"),
                    "filename": entity.get("filename", ""),
                    "file_type": entity.get("file_type", ""),
                    "page_number": entity.get("page_number", 0),
                    "chunk_id": entity.get("chunk_id", ""),
                    "parent_chunk_id": entity.get("parent_chunk_id", ""),
                    "root_chunk_id": entity.get("root_chunk_id", ""),
                    "chunk_level": entity.get("chunk_level", 0),
                    "chunk_idx": entity.get("chunk_idx", 0),
                    "score": hit.get("distance", 0.0),
                })
        return formatted_results

    def dense_retrieve(
        self,
        dense_embedding: list[float],
        top_k: int = 5,
        filter_expr: str = "",
    ) -> list[dict]:
        def _search(client: MilvusClient):
            return client.search(
                collection_name=self.collection_name,
                data=[dense_embedding],
                anns_field="dense_embedding",
                search_params={"metric_type": "IP", "params": {"ef": 64}},
                limit=top_k,
                output_fields=[
                    "text",
                    "document_id",
                    "category_id",
                    "filename",
                    "file_type",
                    "page_number",
                    "chunk_id",
                    "parent_chunk_id",
                    "root_chunk_id",
                    "chunk_level",
                    "chunk_idx",
                ],
                filter=filter_expr,
            )

        results = self._run_read(_search)
        formatted_results = []
        for hits in results:
            for hit in hits:
                formatted_results.append({
                    "id": hit.get("id"),
                    "text": hit.get("entity", {}).get("text", ""),
                    "document_id": hit.get("entity", {}).get("document_id"),
                    "category_id": hit.get("entity", {}).get("category_id"),
                    "filename": hit.get("entity", {}).get("filename", ""),
                    "file_type": hit.get("entity", {}).get("file_type", ""),
                    "page_number": hit.get("entity", {}).get("page_number", 0),
                    "chunk_id": hit.get("entity", {}).get("chunk_id", ""),
                    "parent_chunk_id": hit.get("entity", {}).get("parent_chunk_id", ""),
                    "root_chunk_id": hit.get("entity", {}).get("root_chunk_id", ""),
                    "chunk_level": hit.get("entity", {}).get("chunk_level", 0),
                    "chunk_idx": hit.get("entity", {}).get("chunk_idx", 0),
                    "score": hit.get("distance", 0.0),
                })
        return formatted_results

    def drop_collection(self) -> None:
        def _drop(client: MilvusClient) -> None:
            if client.has_collection(self.collection_name):
                client.drop_collection(self.collection_name)

        self._run(_drop)


# 兼容旧名；全项目共用同一无状态 Store 实例即可（不缓存连接）
MilvusManager = MilvusStore

_store: MilvusStore | None = None


def get_milvus_store() -> MilvusStore:
    global _store
    if _store is None:
        _store = MilvusStore()
    return _store
