"""文档向量化并写入 Milvus - 支持密集+稀疏向量"""
import os

from backend.indexing.embedding import EmbeddingService, embedding_service as _default_embedding_service
from backend.indexing.milvus_client import MilvusStore, get_milvus_store


class MilvusWriter:
    """文档向量化并写入 Milvus 服务 - 支持混合检索"""

    def __init__(self, embedding_service: EmbeddingService = None, milvus_manager: MilvusStore = None):
        self.embedding_service = embedding_service or _default_embedding_service
        self.milvus_manager = milvus_manager or get_milvus_store()

    def write_chunks(self, chunks: list[dict], batch_size: int = 50, progress_callback=None):
        """将 L3 叶子分块写入 Milvus（稠密向量 + 服务端 BM25 稀疏向量）。"""
        if not chunks:
            return []

        leaf_chunks = [c for c in chunks if c.get("chunk_level") == 3]
        if not leaf_chunks:
            return []

        dense_dim = int(os.getenv("DENSE_EMBEDDING_DIM", "1024"))
        total = len(leaf_chunks)
        vector_ids = []

        with self.milvus_manager.session() as client:
            MilvusStore.ensure_collection(client, self.milvus_manager.collection_name, dense_dim)

            for i in range(0, total, batch_size):
                batch = leaf_chunks[i : i + batch_size]
                texts = [doc["text"] for doc in batch]
                dense_embeddings = self.embedding_service.get_embeddings(texts)

                insert_data = [
                    {
                        "dense_embedding": dense_emb,
                        "text": doc["text"],
                        "document_id": doc.get("document_id"),
                        "category_id": doc.get("category_id"),
                        "filename": doc["filename"],
                        "file_type": doc["file_type"],
                        "file_path": doc.get("file_path", ""),
                        "page_number": doc.get("page_number", 0),
                        "chunk_idx": doc.get("chunk_idx", 0),
                        "chunk_id": doc.get("chunk_id", ""),
                        "parent_chunk_id": doc.get("parent_chunk_id", ""),
                        "root_chunk_id": doc.get("root_chunk_id", ""),
                        "chunk_level": doc.get("chunk_level", 0),
                    }
                    for doc, dense_emb in zip(batch, dense_embeddings)
                ]

                result = client.insert(self.milvus_manager.collection_name, insert_data)
                ids = result.get("ids", []) if isinstance(result, dict) else getattr(result, "primary_keys", [])
                vector_ids.extend(ids)

                if progress_callback:
                    processed = min(i + batch_size, total)
                    progress_callback(processed, total)

        # 返回所有 leaf_chunks 对应的 vector_id 列表
        id_map = {leaf_chunks[i]["chunk_id"]: vector_ids[i] for i in range(len(vector_ids))}
        return [id_map.get(c["chunk_id"]) for c in leaf_chunks]

    def write_documents(self, documents: list[dict], batch_size: int = 50, progress_callback=None):
        """兼容旧接口：传入 list[dict] 直接写入（要求 dict 已包含 document_id 等字段）。"""
        return self.write_chunks(documents, batch_size, progress_callback)
