src/rag/storage.py

import asyncio
import logging
from collections import OrderedDict, defaultdict
from pathlib import Path
from textwrap import dedent
from typing import Any, Dict, List, Optional, Tuple, Type, Union
from uuid import uuid4

import chromadb
import numpy as np
import pandas as pd
from chromadb.api.types import QueryResult
from chromadb.config import Settings
from datasets.utils import metadata
from langchain.chains.base import Chain
from langchain.schema import Document
from langchain_core.callbacks.manager import CallbackManagerForChainRun
from langchain_core.embeddings import Embeddings
from langchain_core.language_models import BaseLanguageModel
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain_openai.embeddings import OpenAIEmbeddings
from pandas.core.algorithms import rank
from pydantic import BaseModel, ConfigDict, Field, conbytes

from src.chains import PresentationAnalysis, SlideAnalysis
from src.chains.prompts import JsonH1AndGDPrompt
from src.config.model_setup import EmbeddingConfig
from src.config.navigator import Navigator
from src.rag import BaseScorer, HyperbolicScorer, ScorerTypes
from src.rag.preprocess import RegexQueryPreprocessor
from src.rag.score import ExponentialScorer, MinScorer

logger = logging.getLogger(__name__)


class SlideChunk(BaseModel):
    """Container for slide chunk data ready for ChromaDB

    Each chunk represents a logical part of a slide with its metadata.
    Chunks from the same slide share same slide_id but have different chunk_types.
    """

    id: str = Field(description="Unique identifier for the chunk")
    text: str = Field(description="Text content to embed")
    metadata: Dict[str, str] = Field(
        description="Associated metadata including slide relationships"
    )

    model_config = ConfigDict(frozen=True)


class ScoredChunk(BaseModel):
    """Container for retrieved chunk with similarity score"""

    document: Document
    score: float

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @property
    def slide_id(self) -> str:
        """Get slide identifier from metadata"""
        return self.document.metadata["slide_id"]

    @property
    def pdf_path(self) -> str:
        return self.document.metadata["pdf_path"]

    @property
    def pdf_name(self) -> str:
        return Path(self.pdf_path).stem

    @property
    def chunk_type(self) -> str:
        """Get chunk type from metadata"""
        return self.document.metadata["chunk_type"]

    @property
    def page_num(self) -> int:
        return int(self.document.metadata["page_num"])


class SearchResult(BaseModel):
    """Container for search results with metadata"""

    chunks: List[ScoredChunk]
    metadata: Dict = Field(default_factory=dict)

    model_config = ConfigDict(arbitrary_types_allowed=True)


class SearchResultPage(BaseModel):
    """Container for search results with full slide context"""

    matched_chunk: ScoredChunk = Field(description="Best matching chunk for this slide")
    slide_chunks: Dict[str, Document] = Field(
        default_factory=dict, description="All chunks from the same slide"
    )
    metadata: Dict = Field(default_factory=dict)
    chunk_distances: Dict[str, Optional[float]] = Field(
        description="Distance scores by chunk type (None if not matched)"
    )

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @property
    def slide_id(self):
        return self.matched_chunk.slide_id

    @property
    def pdf_path(self) -> Path:
        return Path(self.matched_chunk.pdf_path)

    @property
    def pdf_name(self):
        return self.matched_chunk.pdf_name

    @property
    def best_score(self):
        return self.matched_chunk.score

    @property
    def page_num(self):
        return self.matched_chunk.page_num


class SearchResultPresentation(BaseModel):
    """Container for presentation-level search results

    Represents all matching slides from a single presentation
    """

    slides: List[SearchResultPage] = Field(
        description="Matching slides from this presentation"
    )
    scorer: ScorerTypes = MinScorer()
    metadata: Dict = Field(default_factory=dict)

    model_config = ConfigDict(arbitrary_types_allowed=True)

    def __getitem__(self, idx) -> SearchResultPage:
        return self.slides[idx]

    def __len__(self) -> int:
        return len(self.slides)

    def set_scorer(self, scorer: BaseScorer):
        self.scorer = scorer

    @property
    def rank_score(self) -> float:
        if self.scorer is None:
            raise AttributeError("Scorer not set")
        return self.scorer.compute_score(self.slide_scores)

    @property
    def pdf_path(self) -> Path:
        return Path(self.slides[0].pdf_path)

    @property
    def title(self) -> str:
        return self.pdf_path.stem

    @property
    def slide_scores(self):
        return [s.best_score for s in self.slides]

    @property
    def best_distance(self) -> float:
        return min(slide.best_score for slide in self.slides)

    @property
    def best_slide(self) -> SearchResultPage:
        return min(self.slides, key=lambda slide: slide.best_score)

    @property
    def mean_score(self) -> float:
        scores = [s.best_score for s in self.slides]
        return sum(scores) / len(scores) if len(scores) else float("inf")

    def format_as_text(self) -> str:
        """Format search results as text for LLM consumption."""
        text_parts = [f"Presentation: {self.title}\n"]

        for slide in self.slides:
            text_parts.append(f"\nSlide {slide.page_num}:")

            # Add all available chunks in a structured way
            for chunk_type, doc in slide.slide_chunks.items():
                if doc.page_content.strip():
                    text_parts.append(f"\n{chunk_type.replace('_', ' ').title()}:")
                    text_parts.append(doc.page_content.strip())

        return "\n".join(text_parts)


class ScoredPresentations(BaseModel):
    """Container for search results with scoring mechanism

    presentations are sorted
    """

    presentations: List[SearchResultPresentation]
    scorer: ScorerTypes = ExponentialScorer()

    model_config = ConfigDict(arbitrary_types_allowed=True)

    def model_post_init(self, __context: Any) -> None:
        self.sort_presentations()

        for p in self.presentations:
            p.set_scorer(self.scorer)

        return super().model_post_init(__context)

    def __getitem__(self, idx) -> SearchResultPresentation:
        return self.presentations[idx]

    def __len__(self):
        return len(self.presentations)

    def sort_presentations(self):
        self.presentations.sort(key=lambda p: self.scorer.compute_score(p.slide_scores))

    def set_scorer(self, scorer: BaseScorer):
        self.scorer = scorer
        self.sort_presentations()

    def get_scores(self) -> List[float]:
        return [self.scorer.compute_score(p.slide_scores) for p in self.presentations]


class SlideIndexer:
    """Process slides into chunks suitable for ChromaDB storage"""

    def __init__(self, collection_name: str):
        """Initialize indexer with semantic section types"""
        # Main content sections from SlideDescription
        self._content_sections = ["text_content", "visual_content"]

        # General description sections
        self._general_sections = [
            "topic_overview",
            "conclusions_and_insights",
            "layout_and_composition",
        ]

        self._chunk_types = self._content_sections + self._general_sections
        self.collection_name = collection_name

    def _create_chunk_id(self, slide: SlideAnalysis, chunk_type: str) -> str:
        """Create unique identifier for a chunk

        Format: collection_name__presentation_name__page_num__chunk_type
        """
        # Get presentation name from path
        pres_name = slide.pdf_path.stem
        clean_name = "".join(c for c in pres_name if c.isalnum())
        return f"{self.collection_name}__{clean_name}__{slide.page_num}__{chunk_type}"

    def _get_base_id(self, chunk_id: str) -> str:
        """Extract base identifier without short ID

        Args:
            chunk_id: Full chunk identifier

        Returns:
            Base identifier (presentation_name__page_num__chunk_type)
        """
        # Split by double underscore and take all parts except short ID
        return "__".join(chunk_id.split("__")[1:])

    def _prepare_chunk_metadata(
        self, slide: SlideAnalysis, chunk_type: str
    ) -> Dict[str, str]:
        """Prepare metadata for a chunk"""
        metadata = dict(
            # Basic slide info
            pdf_path=str(slide.pdf_path),
            page_num=str(slide.page_num),  # BUG: why str?
            # Chunk specific
            chunk_type=chunk_type,
            slide_id=f"{slide.pdf_path.stem}__{slide.page_num}",
            section_category=(
                "content"
                if chunk_type in self._content_sections
                else "general_description"
            ),
            # Analysis metadata
            prompt=slide.vision_prompt if slide.vision_prompt else "",
        )

        # # Add any response metadata if present
        # if slide.response_metadata:
        #     for key, value in slide.response_metadata.items():
        #         metadata[f"response_{key}"] = str(value)

        return metadata

    def process_slide(self, slide: SlideAnalysis) -> List[SlideChunk]:
        """Process single slide into chunks

        Args:
            slide: Slide analysis results

        Returns:
            List of chunks ready for embedding
        """
        try:
            chunks = []

            # Get parsed content
            content = slide.parsed_output

            # Process main content sections
            content_dict = content.model_dump()
            for section in self._content_sections:
                if text := content_dict.get(section, "").strip():
                    chunk_id = self._create_chunk_id(slide, section)
                    metadata = self._prepare_chunk_metadata(slide, section)

                    chunks.append(SlideChunk(id=chunk_id, text=text, metadata=metadata))

            # Process general description sections
            general_dict = content.general_description.model_dump()
            for section in self._general_sections:
                if text := general_dict.get(section, "").strip():
                    chunk_id = self._create_chunk_id(slide, section)
                    metadata = self._prepare_chunk_metadata(slide, section)

                    chunks.append(SlideChunk(id=chunk_id, text=text, metadata=metadata))

            # # Add raw LLM output as separate chunk if needed
            # if slide.llm_output.strip():
            #     chunk_id = self._create_chunk_id(slide, "raw_llm_output")
            #     metadata = self._prepare_chunk_metadata(slide, "raw_llm_output")

            #     chunks.append(SlideChunk(
            #         id=chunk_id,
            #         text=slide.llm_output,
            #         metadata=metadata
            #     ))

            if len(chunks):
                logger.info(
                    f"Created {len(chunks)} chunks for slide {slide.page_num} "
                    f"of '{slide.pdf_path.stem}'"
                )
            else:
                logger.warning(
                    f"Created {len(chunks)} chunks for slide {slide.page_num} "
                    f"of '{slide.pdf_path.stem}'"
                )
            return chunks

        except Exception as e:
            logger.error(f"Failed to process slide {slide.page_num}: {str(e)}")
            return []


class ChromaSlideStore:
    """Storage and retrieval of slide chunks using ChromaDB

    IMPORTANT: ChromaDB uses cosine distance, not similarity
    Distance = 1 - similarity

    LOWER distance means MORE similar
    - Distance of 0 means exactly the same
    - Distance of 2 means exactly opposite
    - Distance of 1 means perpendicular (unrelated)
    """

    navigator: Navigator = Navigator()

    def __init__(
        self,
        collection_name: str = "pres1",
        embedding_model: Embeddings = EmbeddingConfig().load_openai(),
        query_preprocessor: Optional[RegexQueryPreprocessor] = RegexQueryPreprocessor(),
    ):
        """Initialize ChromaDB storage"""
        self.navigator = Navigator()
        self._db_path = self.navigator.processed / "chroma"
        self._db_path.mkdir(parents=True, exist_ok=True)

        # Initialize ChromaDB
        self._client = chromadb.PersistentClient(
            path=str(self._db_path),
            settings=Settings(anonymized_telemetry=False, allow_reset=True),
        )

        # Get or create collection
        self._collection = self._client.get_or_create_collection(
            name=collection_name, metadata={"hnsw:space": "cosine"}
        )

        # Initialize OpenAI embeddings
        # self._api_key = os.getenv("OPENAI_API_KEY")
        self._embeddings = embedding_model

        # Initialize query preprocessor
        self.query_preprocessor = query_preprocessor

        # Initialize indexer
        self._indexer = SlideIndexer(collection_name=collection_name)

    def add_slide(self, slide: SlideAnalysis) -> None:
        """Add single slide to storage"""
        # Process slide into chunks
        chunks = self._indexer.process_slide(slide)

        # Skip if no chunks
        if not chunks:
            return

        # Prepare data for ChromaDB
        ids = [chunk.id for chunk in chunks]
        texts = [chunk.text for chunk in chunks]
        metadatas = [chunk.metadata for chunk in chunks]

        # Get embeddings
        embeddings = self._embeddings.embed_documents(texts)

        # Add to ChromaDB
        self._collection.add(
            ids=ids, documents=texts, metadatas=metadatas, embeddings=embeddings
        )

    def _chunk_to_langchain(self, chunk: SlideChunk) -> Document:
        """Convert chunk to LangChain document"""
        return Document(page_content=chunk.text, metadata=chunk.metadata)

    def _get_full_slide(self, slide_id: str) -> Dict[str, Document]:
        """Get all chunks for a specific slide

        Args:
            slide_id: Identifier of the slide in format "presentation_name__page_num"

        Returns:
            Dictionary with chunk types as keys and Document objects as values
        """
        # Get all chunks for this slide
        chunks = self.get_by_metadata({"slide_id": slide_id})

        # Group by chunk type
        return {chunk.metadata["chunk_type"]: chunk for chunk in chunks}

    def _get_embeddings(self, texts: List[str]) -> List[float]:
        """Get embeddings for texts"""
        return self._embeddings.embed_documents(texts)

    def _prepare_where_filter(
        self,
        chunk_types: Optional[List[str]] = None,
        metadata_filter: Optional[Dict] = None,
    ) -> Optional[Dict]:
        """Prepare where filter for ChromaDB query"""
        if not (chunk_types or metadata_filter):
            return None

        where_filter = {}
        if chunk_types:
            where_filter["chunk_type"] = {"$in": chunk_types}
        if metadata_filter:
            where_filter.update(metadata_filter)
        return where_filter

    def _process_chroma_results(self, results: QueryResult) -> List[ScoredChunk]:
        """Convert ChromaDB results to list of scored chunks"""
        scored_chunks = []
        for i in range(len(results["ids"][0])):
            doc = Document(
                page_content=results["documents"][0][i],
                metadata=results["metadatas"][0][i],
            )
            score = results["distances"][0][i]
            scored_chunks.append(ScoredChunk(document=doc, score=score))

        return sorted(scored_chunks, key=lambda chunk: chunk.score)

    def _process_search_results(
        self,
        chunks: List[ScoredChunk],
        query: str,
        n_results: int = 10,
        max_score: float = 1.0,
        metadata: Optional[Dict] = None,
    ) -> SearchResult:
        """Process scored chunks into search results"""
        filtered_chunks = [c for c in chunks if c.score <= max_score][:n_results]

        result_metadata = dict(
            query=query,
            total_chunks=len(chunks),
            filtered_chunks=len(filtered_chunks),
        )
        if metadata:
            result_metadata.update(metadata)

        return SearchResult(
            chunks=filtered_chunks,
            metadata=result_metadata,
        )

    def _process_page_results(
        self,
        search_results: SearchResult,
        n_results: int = 4,
    ) -> List[SearchResultPage]:
        """Process search results into page results"""
        # Group chunks by slide_id while preserving order
        slides_map = OrderedDict()  # type: OrderedDict[str, List[ScoredChunk]]

        for chunk in search_results.chunks:
            if chunk.slide_id not in slides_map:
                slides_map[chunk.slide_id] = []
            slides_map[chunk.slide_id].append(chunk)

        # Process each slide's chunks
        page_results = []
        for slide_id, chunks in slides_map.items():
            chunks.sort(key=lambda x: x.score)
            best_chunk = chunks[0]

            slide_chunks = self._get_full_slide(slide_id)
            chunk_distances = {chunk_type: None for chunk_type in slide_chunks.keys()}

            for chunk in chunks:
                chunk_distances[chunk.chunk_type] = chunk.score

            result = SearchResultPage(
                matched_chunk=best_chunk,
                slide_chunks=slide_chunks,
                chunk_distances=chunk_distances,
                metadata=dict(
                    slide_id=slide_id,
                    best_distance=best_chunk.score,
                    total_chunks=len(slide_chunks),
                    matched_chunks=len(chunks),
                ),
            )
            page_results.append(result)

        return page_results[:n_results] if n_results > 0 else page_results

    def _process_presentation_results(
        self,
        page_results: List[SearchResultPage],
        query: str,
        scorer: BaseScorer = HyperbolicScorer(),
        metadata: Optional[Dict] = None,
    ) -> ScoredPresentations:
        """Process page results into presentation results"""
        presentations_map = (
            OrderedDict()
        )  # type: OrderedDict[str, List[SearchResultPage]]

        for result in page_results:
            pres_name = result.pdf_name
            if pres_name not in presentations_map:
                presentations_map[pres_name] = []
            presentations_map[pres_name].append(result)

        presentation_results = []
        for pres_name, slides in presentations_map.items():
            pres_metadata = dict(
                presentation_name=pres_name,
                total_slides=len(slides),
                query=query,
            )
            if metadata:
                pres_metadata.update(metadata)

            pres_result = SearchResultPresentation(
                slides=slides,
                metadata=pres_metadata,
            )
            presentation_results.append(pres_result)

        return ScoredPresentations(presentations=presentation_results, scorer=scorer)

    async def aquery_storage(
        self,
        query: str,
        chunk_types: Optional[List[str]] = None,
        n_results: int = 10,
        metadata_filter: Optional[Dict] = None,
    ) -> QueryResult:
        """Raw async storage query"""
        # Prepare query
        q_storage = self.query_preprocessor(query) if self.query_preprocessor else query
        query_embedding = await self._embeddings.aembed_query(q_storage)

        # Prepare where filter
        where = self._prepare_where_filter(chunk_types, metadata_filter)

        return self._collection.query(
            query_embeddings=[query_embedding],
            n_results=n_results,
            where=where,
        )

    def query_storage(
        self,
        query: str,
        chunk_types: Optional[List[str]] = None,
        n_results: int = 10,
        metadata_filter: Optional[Dict] = None,
    ) -> QueryResult:
        """Raw sync storage query"""
        # Prepare query
        q_storage = self.query_preprocessor(query) if self.query_preprocessor else query
        query_embedding = self._embeddings.embed_query(q_storage)

        # Prepare where filter
        where = self._prepare_where_filter(chunk_types, metadata_filter)

        return self._collection.query(
            query_embeddings=[query_embedding],
            n_results=n_results,
            where=where,
        )

    async def asearch_query_pages(
        self,
        query: str,
        chunk_types: Optional[List[str]] = None,
        n_results: int = 4,
        max_distance: float = 2.0,
        metadata_filter: Optional[Dict] = None,
    ) -> List[SearchResultPage]:
        """Async search for slides with full context"""
        # Execute query with filters
        raw_results = await self.aquery_storage(
            query, chunk_types, n_results, metadata_filter
        )

        # Process results through pipeline
        chunks = self._process_chroma_results(raw_results)
        search_results = self._process_search_results(
            chunks, query, n_results, max_distance
        )
        return self._process_page_results(search_results, n_results)

    def search_query_pages(
        self,
        query: str,
        chunk_types: Optional[List[str]] = None,
        n_results: int = 4,
        max_distance: float = 2.0,
        metadata_filter: Optional[Dict] = None,
    ) -> List[SearchResultPage]:
        """Sync search for slides with full context"""
        # Execute query with filters
        raw_results = self.query_storage(query, chunk_types, n_results, metadata_filter)

        # Process results through pipeline
        chunks = self._process_chroma_results(raw_results)
        search_results = self._process_search_results(
            chunks, query, n_results, max_distance
        )
        return self._process_page_results(search_results, n_results)

    async def asearch_query_presentations(
        self,
        query: str,
        chunk_types: Optional[List[str]] = None,
        n_results: int = 30,
        scorer: BaseScorer = MinScorer(),
        max_distance: float = 2.0,
        metadata_filter: Optional[Dict] = None,
    ) -> ScoredPresentations:
        """Async search for presentations"""
        # Execute query with filters
        raw_results = await self.aquery_storage(
            query, chunk_types, n_results, metadata_filter
        )

        # Process results through pipeline
        chunks = self._process_chroma_results(raw_results)
        search_results = self._process_search_results(
            chunks, query, n_results, max_distance
        )
        page_results = self._process_page_results(search_results, n_results)
        return self._process_presentation_results(page_results, query, scorer)

    def search_query_presentations(
        self,
        query: str,
        chunk_types: Optional[List[str]] = None,
        n_results: int = 30,
        scorer: BaseScorer = MinScorer(),
        max_distance: float = 2.0,
        metadata_filter: Optional[Dict] = None,
    ) -> ScoredPresentations:
        """Sync search for presentations"""
        # Execute query with filters
        raw_results = self.query_storage(query, chunk_types, n_results, metadata_filter)

        # Process results through pipeline
        chunks = self._process_chroma_results(raw_results)
        search_results = self._process_search_results(
            chunks, query, n_results, max_distance
        )
        page_results = self._process_page_results(search_results, n_results)
        return self._process_presentation_results(page_results, query, scorer)

    def get_by_metadata(
        self, metadata_filter: Dict, n_results: Optional[int] = None
    ) -> List[Document]:
        """Get chunks by metadata filter

        Args:
            metadata_filter: Filter conditions
            n_results: Optional limit on results

        Returns:
            List of LangChain documents
        """
        results = self._collection.get(where=metadata_filter, limit=n_results)

        documents = []
        for i in range(len(results["ids"])):
            doc = Document(
                page_content=results["documents"][i], metadata=results["metadatas"][i]  # type: ignore
            )
            documents.append(doc)

        return documents

    async def add_slide_async(self, slide: SlideAnalysis) -> None:
        """Add single slide to storage asynchronously"""
        # Process slide into chunks
        chunks = self._indexer.process_slide(slide)

        # Skip if no chunks
        if not chunks:
            logger.warning(
                f"Slide {slide.page_num} from '{slide.pdf_path}' had no chunks"
            )
            return

        # Prepare data for ChromaDB
        ids = [chunk.id for chunk in chunks]
        texts = [chunk.text for chunk in chunks]
        metadatas = [chunk.metadata for chunk in chunks]

        # Get embeddings asynchronously
        embeddings = await self._embeddings.aembed_documents(texts)

        # Add to ChromaDB
        self._collection.add(
            ids=ids, documents=texts, metadatas=metadatas, embeddings=embeddings  # type: ignore
        )

    async def process_presentation_async(
        self, presentation: PresentationAnalysis, max_concurrent: int = 5
    ) -> None:
        """Process a single presentation asynchronously with concurrency limit

        Args:
            presentation: Presentation to process
            max_concurrent: Maximum number of slides to process concurrently
        """
        from asyncio import Semaphore, create_task, gather

        # Create semaphore for concurrency control
        semaphore = Semaphore(max_concurrent)

        logger.info(f"Start processing presentation '{presentation.name}'")

        async def process_slide_with_semaphore(slide: SlideAnalysis):
            async with semaphore:
                try:
                    await self.add_slide_async(slide)
                    logger.info(
                        f"Processed slide {slide.page_num} of '{presentation.name}'"
                    )
                except Exception as e:
                    logger.error(
                        f"Failed to process slide {slide.page_num} of "
                        f"'{presentation.name}': {str(e)}"
                    )

        # Create tasks for all slides
        tasks = [
            create_task(process_slide_with_semaphore(slide))
            for slide in presentation.slides
        ]

        # Wait for all tasks to complete
        await gather(*tasks)
        logger.info(f"Completed processing presentation: '{presentation.name}'")

    def validate_presentations(self) -> Tuple[pd.DataFrame, List[str]]:
        """Validate that all presentation slides were properly stored.

        Uses metadata from stored chunks to compare number of pages in presentations.
        Result shows how many pages are in ChromaDB vs expected total pages.

        Returns:
            Tuple containing:
            - DataFrame with presentations statistics:
                Columns:
                - presentation: Presentation name
                - stored_pages: Number of pages found in ChromaDB
                - chunks_per_page: Average chunks per page
                - total_chunks: Total chunks for this presentation
                - chunk_types: Set of unique chunk types
                - min_page: First page number
                - max_page: Last page number
            - List of validation warnings if any inconsistencies found
        """
        # Get all stored chunks
        all_chunks = self._collection.get()

        # Group chunks by presentation
        pres_pages: Dict[str, Set[int]] = defaultdict(set)  # Unique pages
        pres_chunks: Dict[str, int] = defaultdict(int)  # Total chunks
        pres_types: Dict[str, Set[str]] = defaultdict(set)  # Chunk types

        # Process each chunk's metadata
        for metadata in all_chunks["metadatas"]:
            if not metadata:
                continue

            pdf_path = metadata.get("pdf_path", "")
            if not pdf_path:
                continue

            # Extract presentation name from path
            pres_name = Path(pdf_path).stem

            # Track pages, chunks and types
            page_num = int(metadata.get("page_num", -1))
            if page_num >= 0:
                pres_pages[pres_name].add(page_num)

            chunk_type = metadata.get("chunk_type", "unknown")
            pres_types[pres_name].add(chunk_type)

            pres_chunks[pres_name] += 1

        # Compile statistics and warnings
        stats_data = []
        warnings = []

        for pres_name in pres_pages:
            stored_pages = len(pres_pages[pres_name])
            total_chunks = pres_chunks[pres_name]
            chunks_per_page = total_chunks / stored_pages if stored_pages > 0 else 0
            chunk_types = pres_types[pres_name]
            pages = sorted(pres_pages[pres_name])

            stats_data.append(
                {
                    "presentation": pres_name,
                    "stored_pages": stored_pages,
                    "chunks_per_page": round(chunks_per_page, 2),
                    "total_chunks": total_chunks,
                    "chunk_types": chunk_types,
                    "min_page": min(pages) if pages else None,
                    "max_page": max(pages) if pages else None,
                }
            )

            # Check for potential issues
            if (
                chunks_per_page < 3
            ):  # Assuming we should have at least 3 chunks per page
                warnings.append(
                    f"Low chunks per page ({chunks_per_page:.1f}) " f"for '{pres_name}'"
                )

            # Check for page number gaps
            if pages:
                expected_pages = set(range(min(pages), max(pages) + 1))
                missing_pages = expected_pages - pres_pages[pres_name]
                if missing_pages:
                    warnings.append(
                        f"Missing pages {sorted(missing_pages)} in '{pres_name}'"
                    )

            # Check for missing chunk types
            expected_types = {
                "text_content",
                "visual_content",
                "topic_overview",
                "conclusions_and_insights",
                "layout_and_composition",
            }
            missing_types = expected_types - chunk_types
            if missing_types:
                warnings.append(f"Missing chunk types {missing_types} in '{pres_name}'")

        # Create DataFrame from stats
        stats_df = pd.DataFrame(stats_data).sort_values("presentation")

        return stats_df, warnings

    def validate_storage(self) -> Tuple[pd.DataFrame, List[str]]:
        """Helper function to run validation and display results.

        Args:
            store: ChromaSlideStore instance to validate

        Returns:
            Tuple of (statistics DataFrame, list of warnings)
        """
        from IPython.display import display

        stats_df, warnings = self.validate_presentations()

        # Display statistics
        print("\nPresentation Statistics:")
        display(stats_df)

        # Display warnings if any
        if warnings:
            print("\nWarnings:")
            for warning in warnings:
                print(f"- {warning}")
        else:
            print("\nNo validation warnings found.")

        return stats_df, warnings

    def copy_collection(
        self, target_collection_name: str, batch_size: int = 100
    ) -> "ChromaSlideStore":
        """Copy contents of current collection to a new collection.
        Updates document IDs to reflect new collection name.

        Args:
            target_collection_name: Name for the target collection
            batch_size: Number of documents to process in each batch

        Returns:
            New ChromaSlideStore instance with copied data
        """
        # Create new store with same embeddings model
        target_store = ChromaSlideStore(
            collection_name=target_collection_name, embedding_model=self._embeddings
        )

        # Get all documents from source collection
        source_data = self._collection.get(
            include=["embeddings", "metadatas", "documents"]
        )
        total_docs = len(source_data["ids"])

        # Process in batches
        for i in range(0, total_docs, batch_size):
            batch_end = min(i + batch_size, total_docs)

            # Extract batch data
            batch_ids = source_data["ids"][i:batch_end]
            batch_embeddings = source_data["embeddings"][i:batch_end]
            batch_documents = source_data["documents"][i:batch_end]
            batch_metadatas = source_data["metadatas"][i:batch_end]

            # Update IDs to use new collection name
            new_batch_ids = []
            for old_id in batch_ids:
                # Replace old collection name with new one
                # Format: collection_name__presentation_name__page_num__chunk_type
                new_id = old_id.replace(
                    f"{self._collection.name}__",
                    f"{target_collection_name}__",
                    1,  # Replace only first occurrence
                )
                new_batch_ids.append(new_id)

            # Add batch to target collection
            target_store._collection.add(
                ids=new_batch_ids,
                embeddings=batch_embeddings,
                documents=batch_documents,
                metadatas=batch_metadatas,
            )

            logger.info(
                f"Copied batch {i//batch_size + 1}: "
                f"documents {i} to {batch_end} of {total_docs}"
            )

        logger.info(
            f"Successfully copied {total_docs} documents "
            f"from '{self._collection.name}' to '{target_collection_name}'"
        )
        return target_store


class PresentationRetriever(BaseModel):
    """Retriever for slide search that provides formatted context"""

    storage: ChromaSlideStore
    scorer: BaseScorer = ExponentialScorer()
    n_contexts: int = -1
    n_pages: int = -1
    n_query_results: int = 70
    retrieve_page_contexts: bool = True

    model_config = ConfigDict(arbitrary_types_allowed=True)

    @property
    def id(self) -> str:
        return self.__class__.__name__.lower()

    def set_n_query_results(self, n_query_results: int):
        self.n_query_results = n_query_results

    def format_slide(
        self, slide: SearchResultPage, metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        text_parts = (
            []
            if metadata is None
            else [f"{k}: {v}" for k, v in metadata.items()] + ["---"]
        )
        text_parts.append(f"Slide {slide.page_num}:")

        ## Sort chunks by type to ensure consistent ordering
        # sorted_chunks = sorted(
        #     slide.slide_chunks.items(), key=lambda x: x[0]  # Sort by chunk type
        # )
        sorted_chunks = slide.slide_chunks.items()

        # NOTE What if we dont add chunks which did not match

        # Add each chunk's content
        for chunk_type, doc in sorted_chunks:
            if doc.page_content.strip():
                text_parts.append(f"\n{chunk_type.replace('_', ' ').title()}:")
                text_parts.append(doc.page_content.strip())
        return "\n\n".join(text_parts)

    def format_contexts(
        self, pres: SearchResultPresentation, n_contexts: int = -1
    ) -> List[str]:
        """Format presentation results as context for LLM"""
        slide_texts = []

        if n_contexts < 0:
            n_contexts = len(pres.slides)

        # Add content from each matching slide
        for i, slide in enumerate(pres.slides):
            # if i == 0:
            #     slide_text = self.format_slide(slide, metadata=dict(pres_name=pres.title))

            if i >= n_contexts:
                break

            slide_text = self.format_slide(slide)
            slide_texts.append(slide_text)

        return slide_texts

    async def aretrieve(
        self,
        query: str,
        chunk_types: Optional[List[str]] = None,
        n_results: int = 30,
        max_distance: float = 2.0,
        metadata_filter: Optional[Dict] = None,
    ) -> Dict[str, Any]:
        """Retrieve presentations and format context

        Args:
            query: Search query
            chunk_types: Optional list of chunk types to search
            n_results: Number of presentations to return
            max_distance: Maximum distance threshold
            metadata_filter: Optional metadata filters

        Returns:
            Dictionary with presentation results and formatted context
        """

        results = await self.storage.asearch_query_presentations(
            query=query,
            chunk_types=chunk_types,
            n_results=n_results,
            scorer=self.scorer,
            max_distance=max_distance,
            metadata_filter=metadata_filter,
        )

        out = self.results2contexts(results)
        out["pres_results"] = results.model_dump()
        return out

    def retrieve(self, *args, **kwargs) -> Dict[str, Any]:
        """Synchronous wrapper for retrieve"""
        return asyncio.run(self.aretrieve(*args, **kwargs))

    def results2contexts(self, results: ScoredPresentations):
        contexts = []
        n_pres = self.n_contexts if self.n_contexts > 0 else len(results)
        for i, pres in enumerate(results.presentations[:n_pres]):

            # Gather relevant info from presentation
            best_chunk = pres.best_slide.matched_chunk
            pres_info = dict(
                pres_name=pres.title,
                pages=[slide.page_num + 1 for slide in pres.slides],
                best_chunk=dict(
                    chunk_type=best_chunk.chunk_type, distance=best_chunk.score
                ),
            )

            if self.retrieve_page_contexts:
                page_contexts = self.format_contexts(pres, self.n_pages)
                pres_info["contexts"] = page_contexts

            contexts.append(pres_info)

        return dict(
            contexts=contexts,
            # answer=self.format_slide(pres[0], metadata=dict(pres_name=best_pres.title)),
            # contexts=contexts,
        )

    def __call__(self, inputs: Dict[str, Any]):
        return self.retrieve(inputs["question"])

    def set_scorer(self, scorer: ScorerTypes):
        self.scorer = scorer

    def get_log_params(self) -> Dict[str, Any]:
        """Get parameters for MLflow logging"""
        return {
            "type": self.__class__.__name__,
            "n_contexts": self.n_contexts,
            "n_pages": self.n_pages,
            "retrieve_page_contexts": self.retrieve_page_contexts,
        }


class LLMPresentationRetriever(PresentationRetriever):
    """LLM-enhanced retriever that reranks results using structured relevance scoring"""

    class RelevanceRanking(BaseModel):
        class RelevanceEval(BaseModel):
            document_id: int = Field(description="The id of the document")
            relevance: int = Field(description="Relevance score from 1-10")
            explanation: str = Field(
                description="Short passage to clarify relevance score"
            )

        results: list[RelevanceEval]

    llm: ChatOpenAI
    top_k: int = 10

    _parser: JsonOutputParser = JsonOutputParser(pydantic_object=RelevanceRanking)

    rerank_prompt: PromptTemplate = PromptTemplate(
        template=dedent(
            """\
            You are evaluating search results for presentation slides.
            Rate how relevant each document is to the given query.
            The relevance score should be from 1-10 where:
            - 1-3: Low relevance, mostly unrelated content
            - 4-6: Moderate relevance, some related points
            - 7-8: High relevance, clearly addresses the query
            - 9-10: Perfect match, directly answers the query

            Evaluate ALL documents and provide brief explanations.

            Presentations to evaluate:

            {context_str}

            Question: {query_str}

            Output Formatting:
            {format_instructions}
            """
        ),
        input_variables=["context_str", "query_str", "format_instructions"],
    )

    def _format_presentations(self, presentations: List[Dict[str, Any]]) -> str:
        """Format presentations for LLM evaluation"""
        formatted = []
        for i, pres in enumerate(presentations):
            content = [f"Document {i+1}:"]
            content.append(f"Title: {pres['pres_name']}")

            if "contexts" in pres:
                content.append("Content:")
                content.extend(pres["contexts"])

            formatted.append("\n".join(content))

        return "\n\n".join(formatted)

    def _rerank_results(
        self,
        results: List[Dict[str, Any]],
        query: str,
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> List[Dict[str, Any]]:
        """Rerank results using LLM relevance scoring"""
        # Format input for LLM
        context_str = self._format_presentations(results)

        # Get LLM evaluation
        chain = self.rerank_prompt | self.llm.with_structured_output(
            self.RelevanceRanking
        )

        ranking = chain.invoke(
            {
                "context_str": context_str,
                "query_str": query,
                "format_instructions": self._parser.get_format_instructions(),
            },
        )

        if len(ranking.results) != len(results):
            logger.warning(
                f"Reranker returned {len(ranking.results)} results when should {len(results)}"
            )

        # Sort results by relevance score
        sorted_evals = sorted(
            ranking.results,  # pyright: ignore[reportAttributeAccessIssue]
            key=lambda x: x.relevance,
            reverse=True,
        )

        # Reorder original results
        reranked = [
            results[eval.document_id - 1].copy()
            for eval in sorted_evals[: self.top_k]
            if eval.document_id - 1 < len(results)
        ]

        # Add LLM scoring info
        for i in range(min(len(reranked), self.top_k)):
            reranked[i]["llm_score"] = sorted_evals[i].relevance
            reranked[i]["llm_explanation"] = sorted_evals[i].explanation

        return reranked

    def __call__(
        self,
        inputs: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Run the chain"""
        # Get base retrieval results
        base_results = super().retrieve(query=inputs["question"])

        # Rerank using LLM
        if len(base_results["contexts"]) > 1:
            reranked = self._rerank_results(
                base_results["contexts"],
                inputs["question"],
            )
        else:
            reranked = base_results["contexts"]

        # Combine contexts from reranked results
        all_contexts = []
        for result in reranked:
            all_contexts.extend(result["contexts"])

        return dict(
            contexts=reranked,
        )

    async def _arerank_results(
        self,
        results: List[Dict[str, Any]],
        query: str,
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> List[Dict[str, Any]]:
        """Rerank results using LLM relevance scoring asynchronously"""
        # Format input for LLM
        context_str = self._format_presentations(results)

        # Get LLM evaluation asynchronously
        chain = self.rerank_prompt | self.llm.with_structured_output(
            self.RelevanceRanking
        )
        ranking = await chain.ainvoke(
            {
                "context_str": context_str,
                "query_str": query,
                "format_instructions": self._parser.get_format_instructions(),
            },
        )

        if len(ranking.results) != len(results):
            logger.warning(
                f"Reranker returned {len(ranking.results)} results when should {len(results)}"
            )

        # Sort results by relevance score
        sorted_evals = sorted(
            ranking.results,
            key=lambda x: x.relevance,
            reverse=True,
        )

        # Reorder original results
        reranked = [
            results[eval.document_id - 1].copy()
            for eval in sorted_evals[: self.top_k]
            if eval.document_id - 1 < len(results)
        ]

        # Add LLM scoring info
        for i in range(min(len(reranked), self.top_k)):
            reranked[i]["llm_score"] = sorted_evals[i].relevance
            reranked[i]["llm_explanation"] = sorted_evals[i].explanation

        return reranked

    async def aretrieve(
        self,
        query: str,
        chunk_types: Optional[List[str]] = None,
        n_results: int = 30,
        max_distance: float = 2.0,
        metadata_filter: Optional[Dict] = None,
    ) -> Dict[str, Any]:
        """Retrieve presentations and format context asynchronously"""

        base_results = await super().aretrieve(
            query,
            chunk_types,
            n_results,
            max_distance,
            metadata_filter,
        )

        # Rerank using LLM asynchronously
        if len(base_results["contexts"]) > 1:
            reranked = await self._arerank_results(
                base_results["contexts"],
                query,
            )
        else:
            reranked = base_results["contexts"]

        return dict(contexts=reranked)

    def get_log_params(self) -> Dict[str, Any]:
        """Get parameters for MLflow logging including LLM specifics"""
        params = super().get_log_params()
        params.update(
            {
                "llm_model": self.llm.model_name,
                "llm_temperature": self.llm.temperature,
                "top_k": self.top_k,
            }
        )
        return params


RetrieverTypes = Union[PresentationRetriever, LLMPresentationRetriever]

# def create_slides_database(
#     presentations: List[PresentationAnalysis], collection_name: str = "slides"
# ) -> ChromaSlideStore:
#     """Create ChromaDB database from slides

#     Args:
#         presentations: List of analyzed presentations
#         collection_name: Name for ChromaDB collection

#     Returns:
#         Configured ChromaSlideStore instance
#     """
#     from dotenv import load_dotenv

#     load_dotenv()
#     # Initialize store
#     store = ChromaSlideStore(collection_name=collection_name)

#     # Add slides from all presentations
#     for presentation in presentations:
#         print(f"Processing '{presentation.name}'...")
#         for slide in presentation.slides:
#             store.add_slide(slide)

#     return store


async def create_slides_database_async(
    presentations: List[PresentationAnalysis],
    collection_name: str = "slides",
    embedding_model: Optional[Embeddings] = None,
    max_concurrent_slides: int = 5,
) -> ChromaSlideStore:
    """Create ChromaDB database from slides asynchronously

    Args:
        presentations: List of analyzed presentations
        collection_name: Name for ChromaDB collection
        embedding_model: Optional embedding model to use
        max_concurrent_slides: Maximum number of slides to process concurrently

    Returns:
        Configured ChromaSlideStore instance
    """
    from asyncio import create_task, gather

    # Initialize store
    store = ChromaSlideStore(
        collection_name=collection_name,
        embedding_model=embedding_model or EmbeddingConfig().load_openai(),
    )

    for pres in presentations:
        await store.process_presentation_async(
            pres, max_concurrent=max_concurrent_slides
        )

    # # Process presentations concurrently
    # tasks = [
    #     create_task(
    #         store.process_presentation_async(
    #             presentation, max_concurrent=max_concurrent_slides
    #         )
    #     )
    #     for presentation in presentations
    # ]
    #
    # # Wait for all presentations to be processed
    # await gather(*tasks)

    return store


def create_slides_database(
    presentations: List[PresentationAnalysis],
    collection_name: str = "slides",
    embedding_model: Optional[Embeddings] = None,
    max_concurrent_slides: int = 3,
) -> ChromaSlideStore:
    """Synchronous wrapper for create_slides_database_async

    Args:
        presentations: List of analyzed presentations
        collection_name: Name for ChromaDB collection
        embedding_model: Optional embedding model to use
        max_concurrent_slides: Maximum number of slides to process concurrently

    Returns:
        Configured ChromaSlideStore instance
    """
    import asyncio

    from dotenv import load_dotenv

    load_dotenv()

    return asyncio.run(
        create_slides_database_async(
            presentations=presentations,
            collection_name=collection_name,
            embedding_model=embedding_model,
            max_concurrent_slides=max_concurrent_slides,
        )
    )