"""
InfiniRetri Processor for large document processing.

Orchestrates the complete pipeline: chunking, embedding, compression,
and retrieval for processing documents with 10M+ tokens.

Requirements: MEM-06 (large document processing)
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, List, Optional

from gsd_rlm.memory.retrieval.chunker import Chunk, SemanticChunker
from gsd_rlm.memory.retrieval.compressor import SemanticCompressor

if TYPE_CHECKING:
    from gsd_rlm.agents.definition import LLMProvider


@dataclass
class InfiniRetriResult:
    """Result of processing a document through InfiniRetri.

    Contains the original document metadata, processed chunks,
    and retrieval index for semantic search.

    Attributes:
        document_id: Unique identifier for the source document
        original_tokens: Token count of original document
        chunks: List of processed chunks with summaries
        index: Retrieval index (embedding -> chunk mapping)
        compression_achieved: Actual compression ratio achieved
    """

    document_id: str
    original_tokens: int
    chunks: List[Chunk]
    index: Dict[str, Any]
    compression_achieved: float = 0.0

    @property
    def compressed_tokens(self) -> int:
        """Total tokens in chunk summaries.

        Returns:
            Sum of summary token counts
        """
        total = 0
        for chunk in self.chunks:
            if chunk.summary:
                # Approximate token count
                total += max(1, len(chunk.summary) // 4)
        return total

    def to_dict(self) -> Dict[str, Any]:
        """Convert result to dictionary for serialization.

        Returns:
            Dictionary representation
        """
        return {
            "document_id": self.document_id,
            "original_tokens": self.original_tokens,
            "chunks": [c.to_dict() for c in self.chunks],
            "index": self.index,
            "compression_achieved": self.compression_achieved,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "InfiniRetriResult":
        """Create result from dictionary.

        Args:
            data: Dictionary with result data

        Returns:
            InfiniRetriResult instance
        """
        return cls(
            document_id=data.get("document_id", ""),
            original_tokens=data.get("original_tokens", 0),
            chunks=[Chunk.from_dict(c) for c in data.get("chunks", [])],
            index=data.get("index", {}),
            compression_achieved=data.get("compression_achieved", 0.0),
        )


class InfiniRetriProcessor:
    """Process large documents with semantic chunking and compression.

    Implements the complete InfiniRetri pipeline:
    1. Semantic chunking at sentence boundaries
    2. Optional embedding generation for each chunk
    3. Compression to summaries (56x target)
    4. Importance scoring based on position and content
    5. Retrieval index construction

    Example:
        ```python
        processor = InfiniRetriProcessor(
            chunk_size=512,
            chunk_overlap=50,
            compression_ratio=0.018
        )

        result = await processor.process_document(
            content=large_document,
            document_id="doc-001",
            llm_provider=ollama,
            embedding_model=sentence_transformer
        )

        # Retrieve relevant chunks
        relevant = await processor.retrieve_relevant_chunks(
            result,
            query="What are the main features?",
            k=5
        )
        ```
    """

    def __init__(
        self,
        chunk_size: int = 512,
        chunk_overlap: int = 50,
        compression_ratio: float = 0.018,
    ):
        """Initialize the processor.

        Args:
            chunk_size: Target tokens per chunk (default: 512)
            chunk_overlap: Overlap tokens between chunks (default: 50)
            compression_ratio: Target compression ratio (default: 0.018 = 56x)
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.compression_ratio = compression_ratio

        # Initialize components
        self._chunker = SemanticChunker(
            chunk_size=chunk_size, chunk_overlap=chunk_overlap
        )
        self._compressor = SemanticCompressor(
            target_compression=1.0 / compression_ratio
        )

    async def process_document(
        self,
        content: str,
        document_id: str,
        llm_provider: Optional["LLMProvider"] = None,
        embedding_model: Optional[Any] = None,
    ) -> InfiniRetriResult:
        """Process a large document through the complete pipeline.

        Args:
            content: Document text to process
            document_id: Unique identifier for the document
            llm_provider: Optional LLM for intelligent compression
            embedding_model: Optional model for generating embeddings

        Returns:
            InfiniRetriResult with chunks, index, and compression metrics
        """
        # Step 1: Semantic chunking
        chunks = self._semantic_chunk(content, document_id)

        # Step 2: Generate embeddings (if model provided)
        if embedding_model is not None:
            chunks = await self._generate_embeddings(chunks, embedding_model)

        # Step 3: Compress chunks
        chunks = await self._compress_chunks(chunks, llm_provider)

        # Step 4: Score importance
        chunks = self._score_importance(chunks)

        # Step 5: Build retrieval index
        index = self._build_index(chunks)

        # Calculate metrics
        original_tokens = self._estimate_tokens(content)
        compression_achieved = self._calculate_compression(chunks, original_tokens)

        return InfiniRetriResult(
            document_id=document_id,
            original_tokens=original_tokens,
            chunks=chunks,
            index=index,
            compression_achieved=compression_achieved,
        )

    def _semantic_chunk(self, content: str, document_id: str) -> List[Chunk]:
        """Split content into semantic chunks.

        Args:
            content: Document content
            document_id: Document identifier

        Returns:
            List of Chunk objects
        """
        return self._chunker.chunk(content, document_id)

    async def _generate_embeddings(
        self, chunks: List[Chunk], embedding_model: Any
    ) -> List[Chunk]:
        """Generate embeddings for chunks.

        Args:
            chunks: Chunks to embed
            embedding_model: Model for generating embeddings

        Returns:
            Chunks with embeddings populated
        """
        for chunk in chunks:
            try:
                # Use summary if available, otherwise content
                text = chunk.summary if chunk.summary else chunk.content
                if hasattr(embedding_model, "encode"):
                    # Sentence transformer style
                    embedding = embedding_model.encode(text)
                    chunk.embedding = embedding.tolist()
                elif hasattr(embedding_model, "embed"):
                    # OpenAI style
                    embedding = await embedding_model.embed(text)
                    chunk.embedding = embedding
                else:
                    # Unknown model type - skip embedding
                    chunk.embedding = []
            except Exception:
                # Embedding failed - continue without
                chunk.embedding = []

        return chunks

    async def _compress_chunks(
        self, chunks: List[Chunk], llm_provider: Optional["LLMProvider"]
    ) -> List[Chunk]:
        """Compress chunks using SemanticCompressor.

        Args:
            chunks: Chunks to compress
            llm_provider: Optional LLM provider

        Returns:
            Compressed chunks
        """
        return await self._compressor.compress_chunks(chunks, llm_provider)

    def _score_importance(self, chunks: List[Chunk]) -> List[Chunk]:
        """Score chunk importance based on position and content.

        Importance factors:
        - First and last chunks get bonus (introduction/conclusion)
        - Chunks with questions get bonus
        - Chunks with numbers/data get bonus

        Args:
            chunks: Chunks to score

        Returns:
            Chunks with importance_score populated
        """
        if not chunks:
            return chunks

        n_chunks = len(chunks)

        for i, chunk in enumerate(chunks):
            score = 0.5  # Base score

            # Position bonus
            if i == 0:
                # First chunk - introduction
                score += 0.3
            elif i == n_chunks - 1:
                # Last chunk - conclusion
                score += 0.2
            elif i < n_chunks * 0.1:
                # Early chunks
                score += 0.1

            # Content bonuses
            content = chunk.content.lower()
            summary = chunk.summary.lower() if chunk.summary else ""
            combined = content + " " + summary

            # Questions indicate key information
            if "?" in combined:
                score += 0.1

            # Numbers indicate data/facts
            import re

            if re.search(r"\d+", combined):
                score += 0.05

            # Key phrases
            key_phrases = [
                "important",
                "key",
                "critical",
                "essential",
                "must",
                "should",
            ]
            for phrase in key_phrases:
                if phrase in combined:
                    score += 0.05
                    break

            # Normalize to 0-1
            chunk.importance_score = min(1.0, score)

        return chunks

    def _build_index(self, chunks: List[Chunk]) -> Dict[str, Any]:
        """Build retrieval index from chunks.

        Creates a simple in-memory index mapping:
        - chunk_id -> chunk data
        - embedding vectors for similarity search

        Args:
            chunks: Chunks to index

        Returns:
            Index dictionary
        """
        index = {
            "chunk_ids": [],
            "embeddings": [],
            "summaries": [],
            "importance_scores": [],
        }

        for chunk in chunks:
            index["chunk_ids"].append(chunk.chunk_id)
            index["embeddings"].append(chunk.embedding)
            index["summaries"].append(chunk.summary)
            index["importance_scores"].append(chunk.importance_score)

        return index

    async def retrieve_relevant_chunks(
        self,
        result: InfiniRetriResult,
        query: str,
        k: int = 10,
        embedding_model: Optional[Any] = None,
    ) -> List[Chunk]:
        """Retrieve most relevant chunks for a query.

        Uses cosine similarity between query embedding and chunk
        embeddings to find relevant chunks.

        Args:
            result: Processed document result
            query: Search query
            k: Number of chunks to retrieve
            embedding_model: Model for query embedding (required if chunks have embeddings)

        Returns:
            List of k most relevant chunks
        """
        chunks = result.chunks
        if not chunks:
            return []

        # Check if we have embeddings
        has_embeddings = any(c.embedding for c in chunks)

        if has_embeddings and embedding_model is not None:
            # Semantic search with embeddings
            query_embedding = await self._get_query_embedding(query, embedding_model)
            return self._semantic_search(chunks, query_embedding, k)
        else:
            # Fallback: keyword-based search
            return self._keyword_search(chunks, query, k)

    async def _get_query_embedding(
        self, query: str, embedding_model: Any
    ) -> List[float]:
        """Get embedding for query.

        Args:
            query: Query text
            embedding_model: Model for embedding

        Returns:
            Query embedding vector
        """
        try:
            if hasattr(embedding_model, "encode"):
                embedding = embedding_model.encode(query)
                return embedding.tolist()
            elif hasattr(embedding_model, "embed"):
                return await embedding_model.embed(query)
        except Exception:
            pass

        return []

    def _semantic_search(
        self, chunks: List[Chunk], query_embedding: List[float], k: int
    ) -> List[Chunk]:
        """Search using semantic similarity.

        Args:
            chunks: Chunks to search
            query_embedding: Query vector
            k: Number of results

        Returns:
            Top k chunks by similarity
        """
        if not query_embedding:
            return chunks[:k]

        # Calculate similarities
        similarities = []
        for chunk in chunks:
            if chunk.embedding:
                sim = self._cosine_similarity(query_embedding, chunk.embedding)
                # Boost by importance score
                sim += chunk.importance_score * 0.1
            else:
                sim = 0.0
            similarities.append((chunk, sim))

        # Sort by similarity (descending)
        similarities.sort(key=lambda x: x[1], reverse=True)

        return [chunk for chunk, _ in similarities[:k]]

    def _keyword_search(self, chunks: List[Chunk], query: str, k: int) -> List[Chunk]:
        """Fallback keyword-based search.

        Args:
            chunks: Chunks to search
            query: Query string
            k: Number of results

        Returns:
            Top k chunks by keyword overlap
        """
        query_terms = set(query.lower().split())
        scores = []

        for chunk in chunks:
            # Combine content and summary
            text = (chunk.content + " " + chunk.summary).lower()
            chunk_terms = set(text.split())

            # Calculate overlap
            overlap = len(query_terms & chunk_terms)
            score = overlap / max(1, len(query_terms))

            # Boost by importance
            score += chunk.importance_score * 0.1

            scores.append((chunk, score))

        # Sort by score (descending)
        scores.sort(key=lambda x: x[1], reverse=True)

        return [chunk for chunk, _ in scores[:k]]

    def _cosine_similarity(self, a: List[float], b: List[float]) -> float:
        """Calculate cosine similarity between two vectors.

        Args:
            a: First vector
            b: Second vector

        Returns:
            Cosine similarity (-1 to 1)
        """
        if not a or not b or len(a) != len(b):
            return 0.0

        dot_product = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(x * x for x in b))

        if norm_a == 0 or norm_b == 0:
            return 0.0

        return dot_product / (norm_a * norm_b)

    def _estimate_tokens(self, text: str) -> int:
        """Estimate token count for text.

        Args:
            text: Text to estimate

        Returns:
            Estimated token count
        """
        if not text:
            return 0
        return max(1, len(text) // 4)

    def _calculate_compression(
        self, chunks: List[Chunk], original_tokens: int
    ) -> float:
        """Calculate compression ratio achieved.

        Args:
            chunks: Compressed chunks
            original_tokens: Original document tokens

        Returns:
            Compression ratio (e.g., 56.0 = 56x compression)
        """
        if original_tokens == 0:
            return 0.0

        compressed_tokens = sum(
            max(1, len(c.summary) // 4) if c.summary else 0 for c in chunks
        )

        if compressed_tokens == 0:
            return 0.0

        return original_tokens / compressed_tokens
