from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from pymilvus import DataType, MilvusClient

from app.embeddings import EmbeddingProvider
from app.schemas import Evidence


@dataclass(frozen=True)
class IndexedDocument:
    """Milvus 入库对象，承载受控元数据用于按服务 / OS 过滤。"""

    content: str
    source: str
    doc_type: str
    chunk_index: int
    policy_type: str
    category: str  # 课程版语义为"适用类目"；本作业改用为 client_os 过滤值
    service_name: str  # 新增字段，用于按服务过滤避免跨服务召回
    metadata: dict[str, Any]


class MilvusVectorStore:
    def __init__(
        self,
        uri: str,
        token: str,
        collection_name: str,
        embeddings: EmbeddingProvider,
    ) -> None:
        kwargs: dict[str, Any] = {"uri": uri}
        if token:
            kwargs["token"] = token
        self.client = MilvusClient(**kwargs)
        self.collection_name = collection_name
        self.embeddings = embeddings

    def is_ready(self) -> bool:
        self.client.list_collections()
        return True

    def create_collection(self, recreate: bool = False) -> None:
        exists = self.client.has_collection(collection_name=self.collection_name)
        # 课程数据允许显式重建；生产环境应创建新 Collection 后切换 Alias。
        if exists and recreate:
            self.client.drop_collection(collection_name=self.collection_name)
            exists = False
        if exists:
            self.client.load_collection(collection_name=self.collection_name)
            return

        schema = MilvusClient.create_schema(
            auto_id=True,
            enable_dynamic_field=False,
        )
        schema.add_field("id", DataType.INT64, is_primary=True)
        schema.add_field("content", DataType.VARCHAR, max_length=8192)
        schema.add_field("source", DataType.VARCHAR, max_length=1024)
        schema.add_field("doc_type", DataType.VARCHAR, max_length=64)
        schema.add_field("chunk_index", DataType.INT64)
        schema.add_field("policy_type", DataType.VARCHAR, max_length=64)
        schema.add_field("category", DataType.VARCHAR, max_length=64)
        schema.add_field("service_name", DataType.VARCHAR, max_length=64)
        schema.add_field("metadata", DataType.JSON)
        schema.add_field(
            "embedding",
            DataType.FLOAT_VECTOR,
            # Schema 维度必须与当前 Embedding Provider 完全一致。
            dim=self.embeddings.dimension,
        )

        index_params = self.client.prepare_index_params()
        index_params.add_index(
            field_name="embedding",
            index_name="embedding_hnsw",
            index_type="HNSW",
            metric_type="COSINE",
            params={"M": 16, "efConstruction": 200},
        )
        # 三个反向索引覆盖常见过滤维度：政策类型 / 服务名称 / 客户端 OS。
        # 任意写入端只能填入受控字符串，禁止任意表达式进入 Milvus。
        index_params.add_index(
            field_name="policy_type",
            index_name="policy_type_inverted",
            index_type="INVERTED",
        )
        index_params.add_index(
            field_name="service_name",
            index_name="service_name_inverted",
            index_type="INVERTED",
        )
        index_params.add_index(
            field_name="category",
            index_name="category_inverted",
            index_type="INVERTED",
        )

        self.client.create_collection(
            collection_name=self.collection_name,
            schema=schema,
            index_params=index_params,
        )
        self.client.load_collection(collection_name=self.collection_name)

    def insert_documents(self, documents: list[IndexedDocument]) -> int:
        if not documents:
            return 0
        # 先批量生成向量，再用 strict zip 防止文档与向量静默错位。
        vectors = self.embeddings.embed_documents([item.content for item in documents])
        rows = []
        for item, vector in zip(documents, vectors, strict=True):
            rows.append(
                {
                    "content": item.content,
                    "source": item.source,
                    "doc_type": item.doc_type,
                    "chunk_index": item.chunk_index,
                    "policy_type": item.policy_type,
                    "category": item.category,
                    "service_name": item.service_name,
                    "metadata": item.metadata,
                    "embedding": vector,
                }
            )
        self.client.insert(collection_name=self.collection_name, data=rows)
        self.client.flush(collection_name=self.collection_name)
        return len(rows)

    @staticmethod
    def _validate_filter_value(field: str, value: str) -> str:
        """白名单过滤值：只允许字母数字下划线连字符，避免 Milvus 表达式注入。"""
        if not value:
            return ""
        if not value.replace("_", "").replace("-", "").isalnum():
            raise ValueError(f"{field} 非法：{value!r}")
        return value

    def search(
        self,
        query: str,
        top_k: int = 5,
        service_name: str = "",
        client_os: str = "",
        doc_type: str = "",
    ) -> list[Evidence]:
        query_vector = self.embeddings.embed_query(query)

        # Filter 只由已校验的结构化字段生成，不接收任意 Milvus 表达式。
        clauses: list[str] = []
        normalized_service = self._validate_filter_value("service_name", service_name)
        if normalized_service:
            clauses.append(f'service_name == "{normalized_service}"')
        normalized_os = self._validate_filter_value("client_os", client_os)
        if normalized_os:
            clauses.append(f'category == "{normalized_os}"')
        normalized_doc_type = self._validate_filter_value("doc_type", doc_type)
        if normalized_doc_type:
            clauses.append(f'doc_type == "{normalized_doc_type}"')
        filter_expression = " and ".join(clauses)

        results = self.client.search(
            collection_name=self.collection_name,
            data=[query_vector],
            anns_field="embedding",
            filter=filter_expression,
            limit=top_k,
            output_fields=[
                "content",
                "source",
                "doc_type",
                "chunk_index",
                "policy_type",
                "category",
                "service_name",
                "metadata",
            ],
            search_params={"metric_type": "COSINE", "params": {"ef": 64}},
        )
        evidence: list[Evidence] = []
        # Store 层统一转换 Evidence，Graph 不依赖 PyMilvus 的原始返回结构。
        for hit in results[0]:
            entity = hit.get("entity", {})
            evidence.append(
                Evidence(
                    source_type="milvus",
                    source=entity.get("source", "unknown"),
                    content=entity.get("content", ""),
                    score=float(hit.get("distance", 0.0)),
                    metadata={
                        "id": hit.get("id"),
                        "doc_type": entity.get("doc_type"),
                        "chunk_index": entity.get("chunk_index"),
                        "policy_type": entity.get("policy_type"),
                        "category": entity.get("category"),
                        "service_name": entity.get("service_name"),
                        **(entity.get("metadata") or {}),
                    },
                )
            )
        return evidence