"""Comparative retrieval skeleton for tenant-aware comparative runtime."""

from __future__ import annotations

import asyncio
from dataclasses import dataclass
from typing import Any


from app.services.chatbot.vector_retriever import VectorRetriever


@dataclass(frozen=True)
class ComparativeRetrieverRequest:
    """Tenant-aware retriever input contract for comparative runtime.

    Attributes:
        user_id: Tenant user identifier derived from JWT.
        project_id: Tenant project identifier from canonical request.
        file_hashes: Canonical comparative file hashes after normalization.
        aspect_query: Semantic query for one comparative aspect.
        top_k: Normalized retrieval candidate cap.
        similarity_threshold: Normalized deterministic similarity threshold.
    """

    user_id: str
    project_id: str
    file_hashes: tuple[str, ...]
    aspect_query: str
    top_k: int
    similarity_threshold: float


class ComparativeRetriever:
    """Tenant-aware retrieval boundary for comparative runtime."""

    def __init__(self, vector_retriever: VectorRetriever | None = None) -> None:
        """Initialize the comparative retriever.
        
        Args:
            vector_retriever: Optional injected VectorRetriever instance.
        """
        self._vector_retriever = vector_retriever or VectorRetriever()

    async def retrieve_balanced_chunks(
        self,
        query: str,
        file_hashes: tuple[str, ...],
        project_id: str,
        user_id: str,
        top_k: int,
        similarity_threshold: float,
    ) -> list[dict[str, Any]]:
        """Retrieve balanced chunks for comparative synthesis across multiple files.
        
        Args:
            query: The aspect query to search for.
            file_hashes: Tuple of document hashes to retrieve chunks for.
            project_id: Tenant project ID for isolation.
            user_id: Tenant user ID for isolation.
            top_k: Maximum chunks per document.
            similarity_threshold: Minimum similarity threshold.
            
        Returns:
            list[dict[str, Any]]: List of grouped evidence chunks per document.
        """
        async def fetch_for_hash(file_hash: str) -> dict[str, Any]:
            # Tenant isolation (project_id, user_id) is enforced at the graph level 
            # prior to this step, ensuring file_hashes are already tenant-verified.
            chunks = await asyncio.to_thread(
                self._vector_retriever.similarity_search_by_hash,
                query,  # query is the first arg, then active_hash, then top_k
                file_hash,
                top_k
            )
            
            valid_chunks = [
                chunk for chunk in chunks
                if chunk.get("similarity", 0.0) >= similarity_threshold
            ]
            
            return {
                "file_hash": file_hash,
                "title": valid_chunks[0].get("title", "") if valid_chunks else "",
                "chunks": valid_chunks,
                "total_chunks": len(valid_chunks),
            }

        tasks = [fetch_for_hash(h) for h in file_hashes]
        evidence_groups = await asyncio.gather(*tasks)
        
        return list(evidence_groups)

    async def retrieve(
        self,
        request: ComparativeRetrieverRequest,
        aspect_embeddings: dict[str, list[float]],
    ) -> dict[str, Any]:
        """Return grouped evidence metadata for the comparative skeleton.

        Args:
            request: Tenant-aware comparative retrieval request.
            aspect_embeddings: Embeddings keyed by the semantic aspect query.

        Returns:
            dict[str, Any]: Skeleton grouped evidence structure.
        """
        evidence_groups = await self.retrieve_balanced_chunks(
            query=request.aspect_query,
            file_hashes=request.file_hashes,
            project_id=request.project_id,
            user_id=request.user_id,
            top_k=request.top_k,
            similarity_threshold=request.similarity_threshold,
        )
        return {
            "tenant_context": {
                "user_id": request.user_id,
                "project_id": request.project_id,
            },
            "file_hashes": list(request.file_hashes),
            "aspect_query": request.aspect_query,
            "top_k": request.top_k,
            "similarity_threshold": request.similarity_threshold,
            "aspect_embeddings_requested": list(aspect_embeddings.keys()),
            "evidence_groups": evidence_groups,
            "total_chunks_used": sum(group["total_chunks"] for group in evidence_groups),
        }
