"""Semantic search retrieval using Milvus vector similarity."""

from uuid import UUID

import structlog
from pymilvus import Collection, connections

from src.api.schemas.query import SourceReference
from src.core.config import settings
from src.core.exceptions import RetrievalError
from src.knowledge.embeddings import EmbeddingService

logger = structlog.get_logger(__name__)


class SemanticRetriever:
    """Retrieves relevant records using semantic search."""

    def __init__(self, top_k: int = 5):
        """Initialize semantic retriever.

        Args:
            top_k: Number of top results to return (default: 5)
        """
        self.top_k = top_k
        self.embedding_service = EmbeddingService()
        self.logger = logger.bind(component="semantic_retriever")

    async def retrieve(
        self,
        query: str,
        tenant_id: str,
        node_id: str,
        document_ids: list[UUID] | None = None,
        tag_filters: dict[str, list[str]] | None = None,
    ) -> tuple[list[SourceReference], float]:
        """Retrieve relevant records for a query.

        Args:
            query: Natural language query
            tenant_id: Tenant UUID string (REQUIRED for data isolation)
            node_id: Node UUID string (REQUIRED for data isolation)
            document_ids: Optional list of document IDs to filter by
            tag_filters: Optional tag filtering with operators:
                - "contains_any": Match records with ANY of these tags
                - "contains_all": Match records with ALL of these tags

        Returns:
            Tuple of (list of source references, top similarity score)

        Raises:
            RetrievalError: If retrieval fails
            ValueError: If tenant_id or node_id is empty
        """
        if not tenant_id or not node_id:
            raise ValueError("tenant_id and node_id are required for data isolation")

        try:
            self.logger.info("retrieval_started", query=query, document_ids=document_ids)

            # Generate query embedding
            query_embedding = await self.embedding_service.embed_batch([query])

            # Connect to Milvus
            connections.connect(host=settings.milvus_host, port=settings.milvus_port)
            collection = Collection(settings.milvus_collection)
            collection.load()

            # Build filter expression
            filters = []

            # MANDATORY — tenant + node isolation (data-isolation-spec Section 3.3)
            filters.append(f'tenant_id == "{tenant_id}"')
            filters.append(f'node_id == "{node_id}"')

            if document_ids:
                uuid_strs = [f'"{str(doc_id)}"' for doc_id in document_ids]
                filters.append(f'document_id in [{", ".join(uuid_strs)}]')

            # Add tag filters
            if tag_filters:
                if "contains_any" in tag_filters and tag_filters["contains_any"]:
                    tags = tag_filters["contains_any"]
                    escaped_tags = [tag.replace("'", "\\'") for tag in tags]
                    tags_list = ", ".join(f'"{tag}"' for tag in escaped_tags)
                    filters.append(f"array_contains_any(tags, [{tags_list}])")

                if "contains_all" in tag_filters and tag_filters["contains_all"]:
                    tags = tag_filters["contains_all"]
                    for tag in tags:
                        escaped_tag = tag.replace("'", "\\'")
                        filters.append(f'array_contains(tags, "{escaped_tag}")')

            expr = " and ".join(filters) if filters else None
            if expr:
                self.logger.debug("filter_applied", filter_expression=expr)

            # Search Milvus with unified schema
            results = collection.search(
                data=query_embedding,
                anns_field="vector",
                param={"metric_type": "COSINE", "params": {"ef": 64}},
                limit=self.top_k,
                expr=expr,
                output_fields=[
                    "text_content",
                    "filename",
                    "document_type",
                    "metadata",
                    "tags",
                ],
            )

            # Extract top similarity score
            top_similarity = 0.0
            if results and results[0]:
                top_similarity = results[0][0].distance

            # Format results as source references
            sources = []
            if results and results[0]:
                for hit in results[0]:
                    doc_type = hit.entity.get("document_type", "excel")
                    metadata = hit.entity.get("metadata", {})

                    # Format location and sheet based on document type
                    if doc_type == "excel":
                        sheet = metadata.get("sheet_name", "")
                        row_num = metadata.get("row_number", 0)
                        location = f"Row {row_num}" if row_num else ""
                    else:  # Word document
                        sheet = ""
                        page_num = metadata.get("page_number")
                        location = f"Page {page_num}" if page_num else ""

                    sources.append(
                        SourceReference(
                            file=hit.entity.get("filename", ""),
                            sheet=sheet,
                            location=location,
                            context=hit.entity.get("text_content", ""),
                        )
                    )

            self.logger.info(
                "retrieval_completed",
                query=query,
                results_count=len(sources),
                top_similarity=top_similarity,
            )

            connections.disconnect("default")
            return sources, top_similarity

        except Exception as e:
            self.logger.error(
                "retrieval_failed",
                query=query,
                error=str(e),
            )
            # Clean up connection
            try:
                connections.disconnect("default")
            except Exception:
                pass
            raise RetrievalError(f"Failed to retrieve results: {e}") from e