"""
SAGE NCS Knowledge Base Search

High-level search interface that integrates the vector store, embeddings,
and document processing pipeline for Nuclear Criticality Safety documents.
"""

import asyncio
import logging
from dataclasses import dataclass, field
from pathlib import Path

from sage.data.ingestion.pdf import PDFExtractor
from sage.data.processing.chunking import (
    ChunkingConfig,
    TechnicalChunker,
)
from sage.data.storage.embeddings import (
    EmbeddingConfig,
    OpenAIEmbeddings,
    SparseEncoder,
)
from sage.data.storage.vector_store import (
    HybridRetriever,
    QdrantMode,
    QdrantVectorStore,
    SearchResult,
    VectorStoreConfig,
)

logger = logging.getLogger(__name__)


@dataclass
class SearchConfig:
    """Configuration for the SAGE search system."""

    # Qdrant settings
    qdrant_mode: QdrantMode = QdrantMode.MEMORY
    qdrant_host: str = "localhost"
    qdrant_port: int = 6333
    qdrant_url: str | None = None
    qdrant_api_key: str | None = None
    collection_name: str = "sage_ncs"

    # Embedding settings
    openai_api_key: str | None = None
    embedding_model: str = "text-embedding-3-small"
    embedding_dimensions: int = 1536

    # Chunking settings
    chunk_size: int = 512
    max_chunk_size: int = 1024
    chunk_overlap: float = 0.1

    # Search settings
    default_top_k: int = 10
    use_hybrid_search: bool = True


@dataclass
class IndexedDocument:
    """Represents a document that has been indexed."""

    doc_id: str
    title: str
    source: str
    file_path: Path
    total_pages: int
    total_chunks: int
    total_chars: int


@dataclass
class SearchResponse:
    """Response from a search query."""

    query: str
    results: list[SearchResult] = field(default_factory=list)
    total_found: int = 0
    search_type: str = "hybrid"

    def format_results(self, max_content_length: int = 300) -> str:
        """Format search results for display."""
        lines = [f"Query: {self.query}", f"Found: {self.total_found} results", ""]

        for i, result in enumerate(self.results, 1):
            content_preview = result.content[:max_content_length]
            if len(result.content) > max_content_length:
                content_preview += "..."

            lines.append(f"[{i}] Score: {result.score:.3f}")
            lines.append(f"    Doc: {result.doc_id}, Page: {result.page_number}")
            lines.append(f"    {content_preview}")
            lines.append("")

        return "\n".join(lines)


class SAGESearch:
    """
    High-level search interface for SAGE NCS knowledge base.

    This class provides a unified interface for:
    - Indexing documents (PDF extraction, chunking, embedding, storage)
    - Searching the knowledge base (hybrid vector + keyword search)
    - Managing the document collection

    Example:
        config = SearchConfig(
            qdrant_mode=QdrantMode.MEMORY,
            openai_api_key="sk-..."
        )
        search = SAGESearch(config)
        await search.initialize()

        # Index a document
        await search.index_document(Path("nureg-1520.pdf"), doc_id="NUREG-1520")

        # Search
        results = await search.search("criticality safety limits for U-235")
        print(results.format_results())
    """

    def __init__(self, config: SearchConfig | None = None):
        """Initialize SAGE search system."""
        self.config = config or SearchConfig()
        self._initialized = False

        # Components (initialized lazily)
        self._vector_store: QdrantVectorStore | None = None
        self._embeddings: OpenAIEmbeddings | None = None
        self._sparse_encoder: SparseEncoder | None = None
        self._retriever: HybridRetriever | None = None
        self._pdf_extractor: PDFExtractor | None = None
        self._chunker: TechnicalChunker | None = None

        # Track indexed documents
        self._indexed_docs: dict[str, IndexedDocument] = {}

    async def initialize(self) -> None:
        """Initialize all components."""
        if self._initialized:
            return

        logger.info("Initializing SAGE Search system...")

        # Initialize vector store
        vs_config = VectorStoreConfig(
            mode=self.config.qdrant_mode,
            host=self.config.qdrant_host,
            port=self.config.qdrant_port,
            cloud_url=self.config.qdrant_url,
            api_key=self.config.qdrant_api_key,
            collection_name=self.config.collection_name,
            dense_vector_size=self.config.embedding_dimensions,
        )
        self._vector_store = QdrantVectorStore(vs_config)
        await self._vector_store.initialize()
        logger.info("Vector store initialized")

        # Initialize embeddings
        embed_config = EmbeddingConfig(
            model_name=self.config.embedding_model,
            api_key=self.config.openai_api_key,
            dimensions=self.config.embedding_dimensions,
        )
        self._embeddings = OpenAIEmbeddings(embed_config)
        logger.info("Embeddings initialized")

        # Initialize sparse encoder for hybrid search
        if self.config.use_hybrid_search:
            self._sparse_encoder = SparseEncoder()
            logger.info("Sparse encoder initialized")

        # Initialize retriever
        self._retriever = HybridRetriever(
            vector_store=self._vector_store,
            embedding_model=self._embeddings,
            sparse_encoder=self._sparse_encoder,
        )
        logger.info("Hybrid retriever initialized")

        # Initialize document processing
        self._pdf_extractor = PDFExtractor()
        chunking_config = ChunkingConfig(
            target_chunk_size=self.config.chunk_size,
            max_chunk_size=self.config.max_chunk_size,
            overlap_ratio=self.config.chunk_overlap,
        )
        self._chunker = TechnicalChunker(chunking_config)
        logger.info("Document processors initialized")

        self._initialized = True
        logger.info("SAGE Search system ready")

    async def close(self) -> None:
        """Clean up resources."""
        if self._vector_store:
            await self._vector_store.close()
        self._initialized = False

    async def __aenter__(self) -> "SAGESearch":
        """Async context manager entry."""
        await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Async context manager exit."""
        await self.close()

    async def index_document(
        self,
        file_path: Path,
        doc_id: str,
        source: str = "unknown",
        title: str | None = None,
    ) -> IndexedDocument:
        """
        Index a document into the knowledge base.

        Args:
            file_path: Path to the PDF document
            doc_id: Document identifier (e.g., "NUREG-1520")
            source: Document source (e.g., "nrc", "osti", "ansi")
            title: Optional document title

        Returns:
            IndexedDocument with indexing statistics
        """
        if not self._initialized:
            await self.initialize()

        logger.info(f"Indexing document: {doc_id} from {file_path}")

        # Extract text from PDF
        extraction_result = self._pdf_extractor.extract(file_path)
        logger.info(
            f"Extracted {len(extraction_result.text_content):,} chars "
            f"from {extraction_result.total_pages} pages"
        )

        # Chunk the content
        chunks = self._chunker.chunk(
            extraction_result.text_content,
            doc_id=doc_id,
        )
        logger.info(f"Created {len(chunks)} chunks")

        # Generate embeddings
        chunk_texts = [chunk.content for chunk in chunks]
        embeddings = await self._embeddings.embed_batch(chunk_texts)
        logger.info(f"Generated {len(embeddings)} embeddings")

        # Generate sparse vectors if using hybrid search
        sparse_vectors = None
        if self._sparse_encoder:
            sparse_vectors = [
                self._sparse_encoder.encode(text) for text in chunk_texts
            ]

        # Add source to chunks before storing
        for chunk in chunks:
            chunk.source = source

        # Add to vector store
        await self._vector_store.add_chunks(
            chunks=chunks,
            dense_embeddings=embeddings,
            sparse_embeddings=sparse_vectors,
        )
        logger.info(f"Added {len(chunks)} chunks to vector store")

        # Track indexed document
        indexed_doc = IndexedDocument(
            doc_id=doc_id,
            title=title or doc_id,
            source=source,
            file_path=file_path,
            total_pages=extraction_result.total_pages,
            total_chunks=len(chunks),
            total_chars=len(extraction_result.text_content),
        )
        self._indexed_docs[doc_id] = indexed_doc

        return indexed_doc

    async def search(
        self,
        query: str,
        top_k: int | None = None,
        filters: dict | None = None,
        use_hybrid: bool | None = None,
    ) -> SearchResponse:
        """
        Search the knowledge base.

        Args:
            query: Search query
            top_k: Number of results to return
            filters: Optional filters (e.g., {"source": "nrc"})
            use_hybrid: Whether to use hybrid search (default: config setting)

        Returns:
            SearchResponse with results
        """
        if not self._initialized:
            await self.initialize()

        top_k = top_k or self.config.default_top_k
        use_hybrid = use_hybrid if use_hybrid is not None else self.config.use_hybrid_search

        if use_hybrid and self._sparse_encoder:
            results = await self._retriever.search(
                query=query,
                top_k=top_k,
                filters=filters,
            )
            search_type = "hybrid"
        else:
            # Dense-only search
            query_embedding = await self._embeddings.embed(query)
            results = await self._vector_store.search(
                query_embedding=query_embedding,
                top_k=top_k,
                filters=filters,
            )
            search_type = "dense"

        return SearchResponse(
            query=query,
            results=results,
            total_found=len(results),
            search_type=search_type,
        )

    async def search_by_document(
        self,
        query: str,
        doc_id: str,
        top_k: int = 5,
    ) -> SearchResponse:
        """Search within a specific document."""
        return await self.search(
            query=query,
            top_k=top_k,
            filters={"doc_id": doc_id},
        )

    async def search_standards(
        self,
        query: str,
        top_k: int = 10,
    ) -> SearchResponse:
        """Search ANSI/ANS standards."""
        return await self.search(
            query=query,
            top_k=top_k,
            filters={"source": "ansi"},
        )

    async def search_nuregs(
        self,
        query: str,
        top_k: int = 10,
    ) -> SearchResponse:
        """Search NUREG documents."""
        return await self.search(
            query=query,
            top_k=top_k,
            filters={"source": "nrc"},
        )

    def get_indexed_documents(self) -> list[IndexedDocument]:
        """Get list of indexed documents."""
        return list(self._indexed_docs.values())

    async def get_stats(self) -> dict:
        """Get knowledge base statistics."""
        if not self._initialized:
            await self.initialize()

        collection_info = await self._vector_store.get_collection_info()

        return {
            "total_documents": len(self._indexed_docs),
            "total_chunks": collection_info.get("points_count", 0),
            "collection_name": self.config.collection_name,
            "embedding_model": self.config.embedding_model,
            "use_hybrid_search": self.config.use_hybrid_search,
            "documents": [
                {
                    "doc_id": doc.doc_id,
                    "source": doc.source,
                    "chunks": doc.total_chunks,
                }
                for doc in self._indexed_docs.values()
            ],
        }


async def demo():
    """Demo the search system."""
    import os

    # Check for API key
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        print("Set OPENAI_API_KEY environment variable to run demo")
        return

    config = SearchConfig(
        qdrant_mode=QdrantMode.MEMORY,
        openai_api_key=api_key,
    )

    async with SAGESearch(config) as search:
        # Index a document
        pdf_path = Path("/Users/bryan/sage/data/raw/nrc/ML15176A258.pdf")
        if pdf_path.exists():
            doc = await search.index_document(
                file_path=pdf_path,
                doc_id="NUREG-1520-Rev2",
                source="nrc",
                title="Standard Review Plan for Fuel Cycle Facilities",
            )
            print(f"Indexed: {doc.doc_id} ({doc.total_chunks} chunks)")

            # Search
            response = await search.search("criticality safety limits for enriched uranium")
            print(response.format_results())
        else:
            print(f"PDF not found: {pdf_path}")


if __name__ == "__main__":
    asyncio.run(demo())
