"""Milvus vector store operations for document embeddings."""

import asyncio
import json
import time
from dataclasses import dataclass
from uuid import UUID

import structlog
from pymilvus import (
    Collection,
    CollectionSchema,
    DataType,
    FieldSchema,
    connections,
    utility,
)

from src.core.config import settings
from src.core.exceptions import IndexingError
from src.extraction.models import ExtractedRecord
from src.knowledge.embeddings import EmbeddingService, record_to_text

logger = structlog.get_logger(__name__)

# Milvus JSON field limit is 65536 bytes; use 60KB threshold for safety margin
MILVUS_JSON_MAX_BYTES = 60000


def truncate_metadata_to_fit(metadata: dict, max_bytes: int = MILVUS_JSON_MAX_BYTES) -> dict:
    """Truncate metadata dict to fit within Milvus JSON field size limit.

    Strategy:
    1. Remove large string values (>1000 chars) first
    2. If still too large, remove list/dict values
    3. Final fallback: keep only essential fields

    Args:
        metadata: Original metadata dict
        max_bytes: Maximum allowed size in bytes (default: 60KB)

    Returns:
        Truncated metadata dict that fits within the limit
    """
    def get_json_size(d: dict) -> int:
        return len(json.dumps(d, ensure_ascii=False, default=str).encode('utf-8'))

    # Check if already within limit
    original_size = get_json_size(metadata)
    if original_size <= max_bytes:
        return metadata

    result = metadata.copy()
    logger.warning(
        "metadata_exceeds_milvus_limit",
        original_size_bytes=original_size,
        max_bytes=max_bytes,
        message="Truncating metadata to fit Milvus JSON field limit"
    )

    # Phase 1: Truncate large string values
    for key, value in list(result.items()):
        if isinstance(value, str) and len(value) > 1000:
            result[key] = value[:997] + "..."
            if get_json_size(result) <= max_bytes:
                return result

    # Phase 2: Remove list/dict values that are large
    for key, value in list(result.items()):
        if isinstance(value, (list, dict)):
            value_size = len(json.dumps(value, ensure_ascii=False, default=str).encode('utf-8'))
            if value_size > 5000:
                del result[key]
                if get_json_size(result) <= max_bytes:
                    return result

    # Phase 3: Remove remaining large values progressively
    while get_json_size(result) > max_bytes and result:
        # Find largest value and remove it
        largest_key = None
        largest_size = 0
        for key, value in result.items():
            value_size = len(json.dumps(value, ensure_ascii=False, default=str).encode('utf-8'))
            if value_size > largest_size:
                largest_size = value_size
                largest_key = key
        if largest_key:
            del result[largest_key]
        else:
            break

    # Mark that metadata was truncated
    if result != metadata:
        result['_metadata_truncated'] = True

    return result


@dataclass
class MilvusEntity:
    """Milvus entity structure for vector storage.

    Unified schema supporting both Excel and Word documents.

    Attributes:
        vector: 3072-dimensional embedding vector
        record_id: UUID linking to PostgreSQL extracted_records table
        document_id: UUID for filtering by document
        tenant_id: Tenant UUID for data isolation (data-isolation-spec Section 3)
        node_id: Node UUID for data isolation (data-isolation-spec Section 3)
        filename: For display in search results
        document_type: Type of document ("excel" or "word")
        text_content: Original text for result display
        metadata: JSON dict with type-specific fields (flexible schema)
            Base fields for Excel: {"sheet_name": str, "row_number": int, "has_symbols": bool, "element_type": str}
            Additional fields from DETAILED pipeline:
                - Images: "image_number", "position_key", "image_format", "width", "height", "cell_reference", "from_cell", "to_cell", "mime_type"
                - Flowcharts: "chunk_type", "shape_count", "connector_count", "connection_count", "bounding_box"
                - Shapes: Shape-specific annotations and positions
                - Note: image_base64 is NOT stored to avoid bloating Milvus
            Base fields for Word: {"chunk_index": int, "contains_table": bool, "contains_image": bool}
        domain_category: Domain category (finance, hr, engineering, etc.)
        domain_terms: List of domain-specific terms for pre-filtering
        tags: Semantic tags for filtering (e.g., req:functional, risk:security)
    """

    vector: list[float]
    record_id: str
    document_id: str
    tenant_id: str
    node_id: str
    filename: str
    document_type: str
    text_content: str
    metadata: dict
    domain_category: str = "general"
    domain_terms: list[str] = None
    tags: list[str] = None

    def __post_init__(self):
        """Initialize domain_terms and tags to empty lists if None."""
        if self.domain_terms is None:
            self.domain_terms = []
        if self.tags is None:
            self.tags = []


def _truncate_to_byte_length(s: str, max_bytes: int) -> str:
    """Truncate a string so its UTF-8 byte length does not exceed max_bytes."""
    encoded = s.encode("utf-8")
    if len(encoded) <= max_bytes:
        return s
    # Truncate bytes and decode safely (ignore partial multibyte chars)
    return encoded[:max_bytes].decode("utf-8", errors="ignore")


def normalize_tags(raw_tags: list, max_length: int = 50) -> list[str]:
    """Convert TagAssignment objects or dicts to simple tag name strings.

    The SemanticTagger returns TagAssignment objects with structure:
    {"tag_id": "...", "name": "req:functional", "confidence": 0.95, ...}

    Milvus expects simple string arrays: ["req:functional", "topic:auth"]

    Args:
        raw_tags: List of tags (TagAssignment dicts, strings, or objects)
        max_length: Max byte length per tag string (Milvus VARCHAR limit is in bytes)

    Returns:
        List of tag name strings suitable for Milvus storage
    """
    if not raw_tags:
        return []

    normalized = []
    for tag in raw_tags:
        if isinstance(tag, str):
            normalized.append(_truncate_to_byte_length(tag, max_length))
        elif isinstance(tag, dict):
            tag_name = tag.get("name")
            if tag_name:
                normalized.append(_truncate_to_byte_length(tag_name, max_length))
        elif hasattr(tag, "name"):
            normalized.append(_truncate_to_byte_length(tag.name, max_length))
        else:
            logger.warning("unknown_tag_format", tag_type=type(tag).__name__)

    return normalized


class VectorStore:
    """Manages Milvus vector storage operations for document records.

    Handles embedding generation, collection creation, and batch vector upload
    with retry logic for robust indexing.
    """

    def __init__(self) -> None:
        """Initialize VectorStore with Milvus connection and EmbeddingService."""
        self.logger = structlog.get_logger()
        self.embedding_service = EmbeddingService()
        self.collection_name = settings.milvus_collection
        self.vector_dim = settings.embedding_dimensions
        self._connected = False

    def _connect(self) -> None:
        """Connect to Milvus server if not already connected."""
        if not self._connected:
            connections.connect(
                alias="default",
                host=settings.milvus_host,
                port=settings.milvus_port,
            )
            self._connected = True
            self.logger.info(
                "milvus_connected",
                host=settings.milvus_host,
                port=settings.milvus_port,
            )

    def _disconnect(self) -> None:
        """Disconnect from Milvus server."""
        if self._connected:
            connections.disconnect("default")
            self._connected = False
            self.logger.info("milvus_disconnected")

    async def create_collection(self) -> None:
        """Create Milvus collection with schema if it doesn't exist.

        Creates the collection with:
        - 3072-dimensional vector field
        - Scalar fields for metadata (record_id, document_id, filename, etc.)
        - HNSW index for fast similarity search
        - Cosine distance metric

        Idempotent: Does nothing if collection already exists.

        Raises:
            IndexingError: If collection creation fails
        """
        self._connect()

        try:
            # Check if collection already exists
            if utility.has_collection(self.collection_name):
                self.logger.info(
                    "collection_already_exists",
                    collection_name=self.collection_name,
                )
                return

            # Define unified schema fields for both Excel and Word documents
            fields = [
                FieldSchema(
                    name="id",
                    dtype=DataType.INT64,
                    is_primary=True,
                    auto_id=True,
                ),
                FieldSchema(
                    name="vector",
                    dtype=DataType.FLOAT_VECTOR,
                    dim=self.vector_dim,
                ),
                FieldSchema(
                    name="record_id",
                    dtype=DataType.VARCHAR,
                    max_length=36,
                ),
                FieldSchema(
                    name="document_id",
                    dtype=DataType.VARCHAR,
                    max_length=36,
                ),
                FieldSchema(
                    name="filename",
                    dtype=DataType.VARCHAR,
                    max_length=255,
                ),
                FieldSchema(
                    name="document_type",
                    dtype=DataType.VARCHAR,
                    max_length=20,
                ),
                FieldSchema(
                    name="text_content",
                    dtype=DataType.VARCHAR,
                    max_length=65535,
                ),
                FieldSchema(
                    name="metadata",
                    dtype=DataType.JSON,
                ),
                # Domain-aware RAG fields
                FieldSchema(
                    name="domain_category",
                    dtype=DataType.VARCHAR,
                    max_length=64,
                    default_value="general",
                ),
                FieldSchema(
                    name="domain_terms",
                    dtype=DataType.ARRAY,
                    element_type=DataType.VARCHAR,
                    max_capacity=50,
                    max_length=100,
                ),
                # Semantic tagging fields
                FieldSchema(
                    name="tags",
                    dtype=DataType.ARRAY,
                    element_type=DataType.VARCHAR,
                    max_capacity=20,
                    max_length=100,  # Increased to support longer multilingual tags (e.g., Japanese)
                ),
                # Tenant + node isolation fields (data-isolation-spec Section 3.2)
                FieldSchema(
                    name="tenant_id",
                    dtype=DataType.VARCHAR,
                    max_length=36,
                ),
                FieldSchema(
                    name="node_id",
                    dtype=DataType.VARCHAR,
                    max_length=36,
                ),
            ]

            # Create schema
            schema = CollectionSchema(
                fields,
                description="DSOL document records with embeddings",
            )

            # Create collection
            collection = Collection(
                name=self.collection_name,
                schema=schema,
            )

            # Create HNSW index for vector similarity search
            index_params = {
                "metric_type": "COSINE",
                "index_type": "HNSW",
                "params": {
                    "M": 16,  # Number of bi-directional links
                    "efConstruction": 200,  # Size of dynamic candidate list
                },
            }

            collection.create_index(
                field_name="vector",
                index_params=index_params,
            )

            # Create INVERTED indexes for domain fields (for efficient filtering)
            collection.create_index(
                field_name="domain_category",
                index_params={"index_type": "INVERTED"},
            )

            collection.create_index(
                field_name="domain_terms",
                index_params={"index_type": "INVERTED"},
            )

            collection.create_index(
                field_name="tags",
                index_params={"index_type": "INVERTED"},
            )

            # Tenant isolation indexes (data-isolation-spec Section 3.2)
            collection.create_index(
                field_name="tenant_id",
                index_params={"index_type": "INVERTED"},
            )

            collection.create_index(
                field_name="node_id",
                index_params={"index_type": "INVERTED"},
            )

            # Load collection into memory for searching
            collection.load()

            self.logger.info(
                "collection_created",
                collection_name=self.collection_name,
                vector_dim=self.vector_dim,
                index_type="HNSW",
            )

        except Exception as e:
            self.logger.error(
                "collection_creation_failed",
                collection_name=self.collection_name,
                error=str(e),
            )
            raise IndexingError(
                f"Failed to create Milvus collection: {e}",
                details={"collection_name": self.collection_name, "error": str(e)},
            ) from e

    async def index_records(
        self,
        records: list[ExtractedRecord] | list[dict],
        document_id: UUID,
        record_ids: list[UUID],
        document_type: str = "excel",
        filename: str = "",
        tenant_id: str = "",
        node_id: str = "",
    ) -> int:
        """Embed and index records in Milvus.

        Supports both Excel ExtractedRecords and Word chunk dictionaries.

        Processes records in batches of 100:
        1. Convert records to text representation
        2. Generate embeddings via Azure OpenAI
        3. Upload vectors to Milvus with metadata

        Args:
            records: List of ExtractedRecord objects (Excel) or chunk dicts (Word)
            document_id: Document UUID for metadata
            record_ids: PostgreSQL record IDs from StructuredStore
            document_type: "excel" or "word" (default: "excel")
            filename: Document filename (required for Word documents)
            tenant_id: Tenant UUID string (required for data isolation)
            node_id: Node UUID string (required for data isolation)

        Returns:
            Total number of vectors indexed

        Raises:
            IndexingError: If indexing fails
            ValueError: If records and record_ids lengths don't match
        """
        if len(records) != len(record_ids):
            raise ValueError(
                f"Records count ({len(records)}) must match "
                f"record_ids count ({len(record_ids)})"
            )

        if not records:
            self.logger.info("no_records_to_index")
            return 0

        # Ensure collection exists before indexing
        await self.create_collection()

        self._connect()

        start_time = time.perf_counter()
        total_indexed = 0
        batch_size = 100

        try:
            # Process in batches of 100
            for i in range(0, len(records), batch_size):
                batch_start = time.perf_counter()

                batch_records = records[i : i + batch_size]
                batch_ids = record_ids[i : i + batch_size]

                # Convert records to text based on document type
                if document_type == "excel":
                    texts = [record_to_text(record) for record in batch_records]
                else:  # Word chunks
                    texts = [record.get("text", "") for record in batch_records]

                # Generate embeddings
                embeddings = await self.embedding_service.embed_batch(texts)

                # Prepare Milvus entities
                entities = self._prepare_milvus_entities(
                    batch_records,
                    embeddings,
                    batch_ids,
                    document_id,
                    document_type,
                    filename,
                    tenant_id=tenant_id,
                    node_id=node_id,
                )

                # Upload to Milvus with retry
                await self._upsert_vectors(entities)

                total_indexed += len(batch_records)
                batch_duration_ms = (time.perf_counter() - batch_start) * 1000

                self.logger.info(
                    "batch_indexed",
                    batch_num=i // batch_size + 1,
                    batch_size=len(batch_records),
                    total_indexed=total_indexed,
                    total_records=len(records),
                    batch_duration_ms=round(batch_duration_ms, 2),
                )

            # Calculate overall metrics
            total_duration_s = time.perf_counter() - start_time

            self.logger.info(
                "indexing_complete",
                total_records=total_indexed,
                duration_seconds=round(total_duration_s, 2),
                records_per_second=round(total_indexed / total_duration_s, 2),
            )

            return total_indexed

        except Exception as e:
            self.logger.error(
                "indexing_failed",
                total_indexed=total_indexed,
                total_records=len(records),
                error=str(e),
            )
            raise IndexingError(
                f"Indexing failed after {total_indexed}/{len(records)} records: {e}",
                details={
                    "total_indexed": total_indexed,
                    "total_records": len(records),
                    "error": str(e),
                },
            ) from e

    async def index_with_vectors(
        self,
        records: list[ExtractedRecord] | list[dict],
        vectors: list[list[float]],
        document_id: UUID,
        record_ids: list[UUID],
        document_type: str = "excel",
        filename: str = "",
        domain_data: list[dict] | None = None,
        tenant_id: str = "",
        node_id: str = "",
    ) -> int:
        """Index records with pre-computed embedding vectors in Milvus.

        Use this method when embeddings have already been generated (e.g., by
        GenerateEmbeddingsOperator in Airflow). Skips the embedding generation
        step for efficiency.

        Args:
            records: List of ExtractedRecord objects (Excel) or chunk dicts (Word)
            vectors: Pre-computed embedding vectors (3072 dims each)
            document_id: Document UUID for metadata
            record_ids: PostgreSQL record IDs from StructuredStore
            document_type: "excel" or "word" (default: "excel")
            filename: Document filename (required for Word documents)
            domain_data: Optional list of dicts with domain_category and domain_terms
            tenant_id: Tenant UUID string (required for data isolation)
            node_id: Node UUID string (required for data isolation)

        Returns:
            Total number of vectors indexed

        Raises:
            IndexingError: If indexing fails
            ValueError: If records, vectors, and record_ids lengths don't match
        """
        if len(records) != len(vectors) or len(records) != len(record_ids):
            raise ValueError(
                f"Records ({len(records)}), vectors ({len(vectors)}), and "
                f"record_ids ({len(record_ids)}) counts must match"
            )

        if not records:
            self.logger.info("no_records_to_index")
            return 0

        # Ensure collection exists before indexing
        await self.create_collection()

        self._connect()

        start_time = time.perf_counter()

        try:
            # Prepare Milvus entities with pre-computed vectors
            entities = self._prepare_milvus_entities(
                records,
                vectors,
                record_ids,
                document_id,
                document_type,
                filename,
                domain_data,
                tenant_id=tenant_id,
                node_id=node_id,
            )

            # Upload to Milvus with retry
            await self._upsert_vectors(entities)

            total_indexed = len(entities)
            total_duration_s = time.perf_counter() - start_time

            self.logger.info(
                "indexing_with_vectors_complete",
                total_records=total_indexed,
                duration_seconds=round(total_duration_s, 2),
                records_per_second=round(total_indexed / total_duration_s, 2)
                if total_duration_s > 0
                else 0,
            )

            return total_indexed

        except Exception as e:
            self.logger.error(
                "indexing_with_vectors_failed",
                total_records=len(records),
                error=str(e),
            )
            raise IndexingError(
                f"Indexing with pre-computed vectors failed: {e}",
                details={
                    "total_records": len(records),
                    "error": str(e),
                },
            ) from e

    def _prepare_milvus_entities(
        self,
        records: list[ExtractedRecord] | list[dict],
        embeddings: list[list[float]],
        record_ids: list[UUID],
        document_id: UUID,
        document_type: str = "excel",
        filename: str = "",
        domain_data: list[dict] | None = None,
        tenant_id: str = "",
        node_id: str = "",
    ) -> list[MilvusEntity]:
        """Prepare Milvus entities from records and embeddings.

        Supports both Excel ExtractedRecords and Word chunk dictionaries.

        Args:
            records: List of ExtractedRecord objects (Excel) or chunk dicts (Word)
            embeddings: List of embedding vectors (3072 dims each)
            record_ids: PostgreSQL record UUIDs
            document_id: Document UUID
            document_type: "excel" or "word"
            filename: Document filename (required for Word docs)
            domain_data: Optional list of dicts with domain_category and domain_terms
            tenant_id: Tenant UUID string for data isolation
            node_id: Node UUID string for data isolation

        Returns:
            List of MilvusEntity objects ready for upload
        """
        entities = []

        for idx, (record, embedding, record_id) in enumerate(zip(records, embeddings, record_ids)):
            # Extract domain data for this record (if provided)
            domain_category = "general"
            domain_terms = []
            tags = []
            if domain_data and idx < len(domain_data):
                domain_category = domain_data[idx].get("domain_category", "general")
                domain_terms = [
                    _truncate_to_byte_length(t, 100) for t in domain_data[idx].get("domain_terms", [])
                ]
            if document_type == "excel":
                # Excel record - ExtractedRecord object
                has_symbols = hasattr(record, "resolved_content") and bool(
                    record.resolved_content
                )
                text_content = record_to_text(record)
                # Milvus VARCHAR limit is 65535 bytes
                # For Excel cells, use 8000 char limit (8000 × 3 = 24000 bytes worst case)
                if len(text_content) > 8000:
                    text_content = text_content[:7997] + "..."

                # Start with base metadata
                metadata = {
                    "sheet_name": record._source["sheet"],
                    "row_number": record._source["row"],
                    "has_symbols": has_symbols,
                }

                # Add cell colors if available
                if "cell_colors" in record._source:
                    metadata["cell_colors"] = record._source["cell_colors"]

                # Add record_json for standard table records (LARGE TABLE pipeline only)
                # Only add for records WITHOUT _legacy_metadata (excludes DETAILED pipeline: images, flowcharts, shapes)
                if not (hasattr(record, '_legacy_metadata') and record._legacy_metadata):
                    # Validate JSON serialization and size
                    try:
                        record_json_str = json.dumps(record.content, ensure_ascii=False, default=str)
                        json_size_kb = len(record_json_str.encode('utf-8')) / 1024

                        # Milvus JSON field limit: 65KB (leave margin for safety)
                        if json_size_kb > 50:
                            logger.warning(
                                f"record_json_too_large",
                                sheet=record._source["sheet"],
                                row=record._source["row"],
                                size_kb=round(json_size_kb, 1),
                                message="Truncating to first 20 fields"
                            )
                            # Store first 20 fields as fallback
                            truncated = dict(list(record.content.items())[:20])
                            metadata['record_json'] = truncated
                            metadata['record_json_truncated'] = True
                        else:
                            metadata['record_json'] = record.content

                    except (TypeError, ValueError) as e:
                        logger.warning(
                            f"record_json_serialization_error",
                            sheet=record._source["sheet"],
                            row=record._source["row"],
                            error=str(e),
                            message="Storing as string"
                        )
                        metadata['record_json'] = str(record.content)

                # Preserve ALL detailed metadata from DETAILED pipeline if available
                if hasattr(record, '_legacy_metadata') and record._legacy_metadata:
                    legacy_meta = record._legacy_metadata
                    element_type = legacy_meta.get('element_type', 'Text')
                    metadata['element_type'] = element_type

                    # Remove row_number for ALL DETAILED pipeline records
                    # (DETAILED chunks are not table rows, they're images/flowcharts/shapes/text-by-sheet)
                    metadata.pop('row_number', None)
                    metadata.pop('has_symbols', None)

                    # Get full metadata from DETAILED pipeline
                    full_metadata = legacy_meta.get('full_metadata', {})
                    if full_metadata:
                        for key, value in full_metadata.items():
                            # Skip fields already added, internal fields, or image_base64 (too large for Milvus)
                            if key not in ['sheet_name', 'sheet_number', 'image_base64'] and not key.startswith('_'):
                                metadata[key] = value

                file_name = record._source["filename"]

                # Extract tags from record if available (normalize to simple strings)
                raw_tags = getattr(record, "tags", []) or []
                tags = normalize_tags(raw_tags)

            else:  # Word record - dict from SemanticChunker
                text_content = record.get("text", "")
                # Milvus VARCHAR limit is 65535 bytes (not characters)
                # Use 1.5x factor for mixed ASCII/multibyte text
                # 65535 / 1.5 = 43690, use 43000 for safety margin
                if len(text_content) > 43000:
                    text_content = text_content[:42997] + "..."

                # Extract Word-specific metadata
                word_metadata = record.get("metadata", {})
                metadata = {
                    "page_number": word_metadata.get("page_number"),
                    "contains_table": word_metadata.get("contains_table", False),
                }
                file_name = filename

                # Extract tags from record if available (normalize to simple strings)
                raw_tags = record.get("tags", []) or []
                tags = normalize_tags(raw_tags)

            # Ensure metadata fits within Milvus JSON field limit (65536 bytes)
            metadata = truncate_metadata_to_fit(metadata)

            entity = MilvusEntity(
                vector=embedding,
                record_id=str(record_id),
                document_id=str(document_id),
                tenant_id=tenant_id,
                node_id=node_id,
                filename=file_name,
                document_type=document_type,
                text_content=text_content,
                metadata=metadata,
                domain_category=domain_category,
                domain_terms=domain_terms,
                tags=tags,
            )

            entities.append(entity)

        return entities

    async def _upsert_vectors(
        self,
        entities: list[MilvusEntity],
        max_retries: int = 3,
        batch_size: int = 500,
    ) -> None:
        """Upload vectors to Milvus with batching and retry logic.

        Inserts in batches to avoid exceeding the gRPC 64MB message limit.

        Args:
            entities: List of MilvusEntity objects to upload
            max_retries: Number of retries for transient errors
            batch_size: Number of entities per insert batch

        Raises:
            IndexingError: If upload fails after all retries
        """
        collection = Collection(self.collection_name)

        total_batches = (len(entities) - 1) // batch_size + 1
        self.logger.info(
            "inserting_vectors",
            total=len(entities),
            batch_size=batch_size,
            batches=total_batches,
        )

        for batch_num in range(total_batches):
            start = batch_num * batch_size
            end = min(start + batch_size, len(entities))
            batch = entities[start:end]

            # Prepare data in columnar format for Milvus with unified schema
            # Note: Exclude 'id' field as it's auto_id=True
            # Order must match collection schema field order
            data = [
                [e.vector for e in batch],
                [e.record_id for e in batch],
                [e.document_id for e in batch],
                [e.filename for e in batch],
                [e.document_type for e in batch],
                [e.text_content for e in batch],
                [e.metadata for e in batch],
                [e.domain_category for e in batch],
                [e.domain_terms for e in batch],
                [e.tags for e in batch],
                [e.tenant_id for e in batch],
                [e.node_id for e in batch],
            ]

            for attempt in range(max_retries):
                try:
                    collection.insert(data)

                    self.logger.debug(
                        "batch_inserted",
                        batch=batch_num + 1,
                        total_batches=total_batches,
                        count=len(batch),
                        attempt=attempt + 1,
                    )

                    break

                except Exception as e:
                    is_last_attempt = attempt == max_retries - 1

                    if is_last_attempt:
                        self.logger.error(
                            "insert_failed_permanently",
                            batch=batch_num + 1,
                            count=len(batch),
                            error=str(e),
                            attempts=max_retries,
                        )
                        raise IndexingError(
                            f"Failed to insert batch {batch_num + 1} after {max_retries} attempts: {e}",
                            details={
                                "batch": batch_num + 1,
                                "count": len(batch),
                                "attempts": max_retries,
                                "error": str(e),
                            },
                        ) from e

                    backoff_seconds = 2**attempt

                    self.logger.warning(
                        "insert_attempt_failed_retrying",
                        batch=batch_num + 1,
                        count=len(batch),
                        error=str(e),
                        attempt=attempt + 1,
                        max_retries=max_retries,
                        backoff_seconds=backoff_seconds,
                    )

                    await asyncio.sleep(backoff_seconds)

        # Flush once after all batches to make data available for queries
        collection.flush()

        self.logger.info(
            "all_vectors_inserted",
            total=len(entities),
            batches=total_batches,
        )

    def __enter__(self):
        """Context manager entry."""
        self._connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self._disconnect()


async def delete_vectors_by_document_id(
    document_id: UUID,
    tenant_id: str = "",
    node_id: str = "",
) -> int:
    """Delete all vectors for a document from Milvus.

    Args:
        document_id: Document UUID to filter vectors.
        tenant_id: Tenant UUID for scoped deletion (defense-in-depth).
        node_id: Node UUID for scoped deletion (defense-in-depth).

    Returns:
        Number of vectors deleted.

    Raises:
        Exception: If Milvus deletion fails.
    """
    try:
        # Connect to Milvus
        connections.connect(
            alias="default",
            host=settings.milvus_host,
            port=settings.milvus_port,
        )

        collection = Collection(settings.milvus_collection)

        # Delete by document_id, scoped by tenant + node if provided
        expr = f'document_id == "{str(document_id)}"'
        if tenant_id:
            expr += f' and tenant_id == "{tenant_id}"'
        if node_id:
            expr += f' and node_id == "{node_id}"'
        result = collection.delete(expr)

        logger.info(
            "milvus_vectors_deleted",
            document_id=str(document_id),
            tenant_id=tenant_id,
            node_id=node_id,
            delete_count=result.delete_count,
        )

        # Disconnect
        connections.disconnect("default")

        return result.delete_count

    except Exception as e:
        logger.error(
            "milvus_deletion_failed",
            document_id=str(document_id),
            error=str(e),
        )
        raise
