"""
RAG Context Builder

Transforms search results into formatted context for LLM generation,
maintaining citation tracking throughout.
"""

import logging
from dataclasses import dataclass, field

from sage.data.storage.citations import ChunkCitation
from sage.data.storage.vector_store import SearchResult

logger = logging.getLogger(__name__)


@dataclass
class ContextChunk:
    """
    A chunk of context with citation information.

    Represents a single piece of retrieved content ready
    for inclusion in the LLM prompt.
    """

    content: str
    doc_id: str
    section: str = ""
    page: int | None = None
    score: float = 0.0
    chunk_id: str = ""

    # Content metadata
    has_table: bool = False
    has_equation: bool = False

    # Citation tracking
    citation: ChunkCitation | None = None

    @property
    def citation_key(self) -> str:
        """Generate a citation key for this chunk."""
        parts = [self.doc_id]
        if self.section:
            parts.append(self.section)
        if self.page:
            parts.append(f"p.{self.page}")
        return ", ".join(parts)

    def to_citation(self) -> ChunkCitation:
        """Convert to a ChunkCitation object."""
        if self.citation:
            return self.citation

        return ChunkCitation(
            doc_id=self.doc_id,
            section=self.section,
            page=self.page,
            has_table=self.has_table,
            has_equation=self.has_equation,
            confidence=self.score,
        )


@dataclass
class RAGContext:
    """
    Complete context for a RAG query.

    Contains all retrieved chunks with citations,
    formatted for LLM consumption.
    """

    query: str
    chunks: list[ContextChunk] = field(default_factory=list)
    formatted_context: str = ""

    # Retrieval metadata
    total_retrieved: int = 0
    avg_score: float = 0.0
    unique_sources: int = 0

    # Quality indicators
    has_relevant_content: bool = True
    confidence_level: str = "high"  # high, medium, low

    @property
    def citations(self) -> list[ChunkCitation]:
        """Get all citations from chunks."""
        return [chunk.to_citation() for chunk in self.chunks]

    @property
    def source_doc_ids(self) -> list[str]:
        """Get unique source document IDs."""
        seen = set()
        result = []
        for chunk in self.chunks:
            if chunk.doc_id not in seen:
                seen.add(chunk.doc_id)
                result.append(chunk.doc_id)
        return result

    def get_chunks_by_doc(self, doc_id: str) -> list[ContextChunk]:
        """Get all chunks from a specific document."""
        return [c for c in self.chunks if c.doc_id == doc_id]


class ContextBuilder:
    """
    Builds RAG context from search results.

    Transforms SearchResult objects into formatted context
    with proper citation tracking and quality assessment.

    Example:
        builder = ContextBuilder()
        context = builder.build(
            query="What is the subcritical limit for U-235?",
            search_results=[...],
            max_chunks=10
        )
        print(context.formatted_context)
    """

    def __init__(
        self,
        max_context_length: int = 8000,
        chunk_separator: str = "\n\n---\n\n",
        include_scores: bool = False,
        score_threshold: float = 0.3,
    ):
        """
        Initialize context builder.

        Args:
            max_context_length: Maximum context length in characters
            chunk_separator: Separator between chunks
            include_scores: Whether to include scores in formatted output
            score_threshold: Minimum score to include a chunk
        """
        self.max_context_length = max_context_length
        self.chunk_separator = chunk_separator
        self.include_scores = include_scores
        self.score_threshold = score_threshold

    def build(
        self,
        query: str,
        search_results: list[SearchResult],
        max_chunks: int | None = None,
    ) -> RAGContext:
        """
        Build RAG context from search results.

        Args:
            query: Original user query
            search_results: Results from vector search
            max_chunks: Maximum number of chunks to include

        Returns:
            RAGContext with formatted content and citations
        """
        # Filter by score threshold
        filtered_results = [
            r for r in search_results if r.score >= self.score_threshold
        ]

        if not filtered_results:
            return self._build_empty_context(query, search_results)

        # Convert to context chunks
        chunks = self._convert_to_chunks(filtered_results)

        # Apply max chunks limit
        if max_chunks:
            chunks = chunks[:max_chunks]

        # Truncate to fit context length
        chunks = self._truncate_to_length(chunks)

        # Format the context
        formatted = self._format_chunks(chunks)

        # Calculate metrics
        scores = [c.score for c in chunks]
        avg_score = sum(scores) / len(scores) if scores else 0.0
        unique_docs = len({c.doc_id for c in chunks})

        # Assess confidence
        confidence = self._assess_confidence(chunks, avg_score)

        return RAGContext(
            query=query,
            chunks=chunks,
            formatted_context=formatted,
            total_retrieved=len(search_results),
            avg_score=avg_score,
            unique_sources=unique_docs,
            has_relevant_content=True,
            confidence_level=confidence,
        )

    def _convert_to_chunks(self, results: list[SearchResult]) -> list[ContextChunk]:
        """Convert SearchResult objects to ContextChunks."""
        chunks = []
        for result in results:
            chunk = ContextChunk(
                content=result.content,
                doc_id=result.doc_id,
                section=result.section_path,
                page=result.page_number,
                score=result.score,
                chunk_id=result.chunk_id,
                has_table=result.has_table,
                has_equation=result.has_equation,
            )

            # Build citation from result metadata if available
            if result.metadata:
                chunk.citation = ChunkCitation(
                    doc_id=result.doc_id,
                    section=result.section_path,
                    page=result.page_number,
                    has_table=result.has_table,
                    has_equation=result.has_equation,
                    confidence=result.score,
                )

            chunks.append(chunk)

        return chunks

    def _truncate_to_length(self, chunks: list[ContextChunk]) -> list[ContextChunk]:
        """Truncate chunks to fit within max context length."""
        result = []
        current_length = 0
        separator_len = len(self.chunk_separator)

        for chunk in chunks:
            chunk_len = len(self._format_single_chunk(chunk))

            if current_length + chunk_len + separator_len > self.max_context_length:
                # Try to include a truncated version
                remaining = self.max_context_length - current_length - separator_len - 100
                if remaining > 200:  # Only include if meaningful content remains
                    truncated = ContextChunk(
                        content=chunk.content[:remaining] + "... [truncated]",
                        doc_id=chunk.doc_id,
                        section=chunk.section,
                        page=chunk.page,
                        score=chunk.score,
                        chunk_id=chunk.chunk_id,
                        has_table=chunk.has_table,
                        has_equation=chunk.has_equation,
                        citation=chunk.citation,
                    )
                    result.append(truncated)
                break

            result.append(chunk)
            current_length += chunk_len + separator_len

        return result

    def _format_chunks(self, chunks: list[ContextChunk]) -> str:
        """Format chunks into a context string."""
        formatted_chunks = [self._format_single_chunk(c) for c in chunks]
        return self.chunk_separator.join(formatted_chunks)

    def _format_single_chunk(self, chunk: ContextChunk) -> str:
        """Format a single chunk with citation header."""
        # Build header
        header_parts = [f"[{chunk.doc_id}"]
        if chunk.section:
            header_parts.append(f", {chunk.section}")
        if chunk.page:
            header_parts.append(f", p.{chunk.page}")
        header_parts.append("]")

        if self.include_scores:
            header_parts.append(f" (score: {chunk.score:.3f})")

        # Add content type indicators
        indicators = []
        if chunk.has_table:
            indicators.append("[TABLE]")
        if chunk.has_equation:
            indicators.append("[EQUATION]")

        header = "".join(header_parts)
        if indicators:
            header = f"{header} {' '.join(indicators)}"

        return f"{header}\n{chunk.content}"

    def _assess_confidence(
        self, chunks: list[ContextChunk], avg_score: float
    ) -> str:
        """Assess overall confidence level."""
        if not chunks:
            return "low"

        # High confidence: multiple relevant chunks with good scores
        if len(chunks) >= 3 and avg_score >= 0.7:
            return "high"

        # Medium confidence: some relevant content
        if len(chunks) >= 1 and avg_score >= 0.5:
            return "medium"

        return "low"

    def _build_empty_context(
        self, query: str, original_results: list[SearchResult]
    ) -> RAGContext:
        """Build context when no results meet threshold."""
        return RAGContext(
            query=query,
            chunks=[],
            formatted_context="No relevant content found in the knowledge base for this query.",
            total_retrieved=len(original_results),
            avg_score=0.0,
            unique_sources=0,
            has_relevant_content=False,
            confidence_level="low",
        )

    def build_multi_query_context(
        self,
        queries: list[str],
        results_per_query: list[list[SearchResult]],
        max_chunks_per_query: int = 5,
    ) -> RAGContext:
        """
        Build context from multiple query results.

        Useful for query expansion or multi-faceted questions.

        Args:
            queries: List of queries (original + expansions)
            results_per_query: Search results for each query
            max_chunks_per_query: Max chunks from each query

        Returns:
            Combined RAGContext
        """
        # Collect all chunks, deduplicating by chunk_id
        seen_ids = set()
        all_chunks = []

        for _query_idx, results in enumerate(results_per_query):
            filtered = [r for r in results if r.score >= self.score_threshold]
            converted = self._convert_to_chunks(filtered[:max_chunks_per_query])

            for chunk in converted:
                if chunk.chunk_id not in seen_ids:
                    seen_ids.add(chunk.chunk_id)
                    all_chunks.append(chunk)

        # Sort by score
        all_chunks.sort(key=lambda c: c.score, reverse=True)

        # Truncate and format
        all_chunks = self._truncate_to_length(all_chunks)
        formatted = self._format_chunks(all_chunks)

        scores = [c.score for c in all_chunks]
        avg_score = sum(scores) / len(scores) if scores else 0.0

        return RAGContext(
            query=queries[0] if queries else "",
            chunks=all_chunks,
            formatted_context=formatted,
            total_retrieved=sum(len(r) for r in results_per_query),
            avg_score=avg_score,
            unique_sources=len({c.doc_id for c in all_chunks}),
            has_relevant_content=len(all_chunks) > 0,
            confidence_level=self._assess_confidence(all_chunks, avg_score),
        )


def build_context_from_results(
    query: str,
    results: list[SearchResult],
    max_chunks: int = 10,
    max_length: int = 8000,
) -> RAGContext:
    """
    Convenience function to build context from search results.

    Args:
        query: User query
        results: Search results
        max_chunks: Maximum chunks to include
        max_length: Maximum context length

    Returns:
        RAGContext ready for generation
    """
    builder = ContextBuilder(max_context_length=max_length)
    return builder.build(query, results, max_chunks=max_chunks)
