import json

from pymilvus import DataType, MilvusClient

from kbqa.rag.types import PreparedChunk, VectorHit


class MilvusVectorStore:
    collection_name = "kbqa_chunks"

    def __init__(self, client: MilvusClient, embedding_dim: int) -> None:
        self.client = client
        self.embedding_dim = embedding_dim

    def ensure_collection(self) -> None:
        if self.client.has_collection(self.collection_name):
            self.client.load_collection(self.collection_name)
            return
        schema = self.client.create_schema(auto_id=False, enable_dynamic_field=False)
        schema.add_field(
            field_name="chunk_id",
            datatype=DataType.VARCHAR,
            is_primary=True,
            max_length=36,
        )
        schema.add_field(
            field_name="document_id",
            datatype=DataType.VARCHAR,
            max_length=36,
        )
        schema.add_field(field_name="order_index", datatype=DataType.INT64)
        schema.add_field(
            field_name="vector",
            datatype=DataType.FLOAT_VECTOR,
            dim=self.embedding_dim,
        )
        index_params = self.client.prepare_index_params()
        index_params.add_index(
            field_name="vector",
            index_type="AUTOINDEX",
            metric_type="COSINE",
        )
        self.client.create_collection(
            collection_name=self.collection_name,
            schema=schema,
            index_params=index_params,
        )
        self.client.load_collection(self.collection_name)

    def upsert(self, chunks: list[PreparedChunk], vectors: list[list[float]]) -> None:
        if len(chunks) != len(vectors):
            raise ValueError("chunks and vectors must have equal lengths")
        if not chunks:
            return
        data = [
            {
                "chunk_id": chunk.id,
                "document_id": chunk.document_id,
                "order_index": chunk.order_index,
                "vector": vector,
            }
            for chunk, vector in zip(chunks, vectors, strict=True)
        ]
        self.client.upsert(collection_name=self.collection_name, data=data)

    def search(self, vector: list[float], limit: int) -> list[VectorHit]:
        results = self.client.search(
            collection_name=self.collection_name,
            data=[vector],
            anns_field="vector",
            limit=limit,
            output_fields=["chunk_id", "document_id", "order_index"],
            search_params={"metric_type": "COSINE"},
        )
        hits: list[VectorHit] = []
        for item in results[0] if results else []:
            entity = item.get("entity", {})
            hits.append(
                VectorHit(
                    chunk_id=str(entity["chunk_id"]),
                    document_id=str(entity["document_id"]),
                    order_index=int(entity["order_index"]),
                    score=float(item["distance"]),
                )
            )
        return hits

    def delete_document(self, document_id: str) -> None:
        encoded = json.dumps(document_id)
        self.client.delete(
            collection_name=self.collection_name,
            filter=f"document_id == {encoded}",
        )

    def ping(self) -> None:
        self.client.list_collections()
