"""
RAG Response Formatting

Formats RAG responses with proper citations and structure
for both human reading and downstream processing.
"""

import re
from dataclasses import dataclass, field
from enum import Enum

from sage.data.storage.citations import ChunkCitation
from sage.rag.context import RAGContext
from sage.rag.validation import CitationExtractor


class OutputFormat(str, Enum):
    """Output format options."""

    MARKDOWN = "markdown"
    PLAIN = "plain"
    HTML = "html"
    JSON = "json"


@dataclass
class FormattedResponse:
    """A formatted RAG response with citations."""

    answer: str
    citations: list[ChunkCitation] = field(default_factory=list)
    sources_section: str = ""
    full_response: str = ""

    # Metadata
    confidence: float = 1.0
    grounded: bool = True
    format: OutputFormat = OutputFormat.MARKDOWN


class CitationFormatter:
    """Formats citations in various styles."""

    def format_inline(self, citation: ChunkCitation) -> str:
        """Format citation for inline use."""
        parts = [citation.doc_id]
        if citation.section:
            parts.append(citation.section)
        if citation.page:
            parts.append(f"p.{citation.page}")
        return f"[{', '.join(parts)}]"

    def format_full(self, citation: ChunkCitation) -> str:
        """Format citation for references section."""
        parts = [citation.doc_id]
        if citation.section_path:
            parts.append(citation.section_path)
        elif citation.section:
            parts.append(citation.section)
        if citation.page:
            parts.append(f"page {citation.page}")
        return ", ".join(parts)

    def format_list(
        self,
        citations: list[ChunkCitation],
        format_type: OutputFormat = OutputFormat.MARKDOWN,
    ) -> str:
        """Format a list of citations."""
        if not citations:
            return ""

        # Deduplicate by doc_id + section
        seen = set()
        unique = []
        for c in citations:
            key = (c.doc_id, c.section, c.page)
            if key not in seen:
                seen.add(key)
                unique.append(c)

        if format_type == OutputFormat.MARKDOWN:
            lines = ["**Sources:**"]
            for i, c in enumerate(unique, 1):
                lines.append(f"{i}. {self.format_full(c)}")
            return "\n".join(lines)

        elif format_type == OutputFormat.HTML:
            items = [f"<li>{self.format_full(c)}</li>" for c in unique]
            return f"<h4>Sources:</h4><ol>{''.join(items)}</ol>"

        else:  # PLAIN or JSON
            lines = ["Sources:"]
            for i, c in enumerate(unique, 1):
                lines.append(f"  {i}. {self.format_full(c)}")
            return "\n".join(lines)


class ResponseFormatter:
    """
    Formats RAG responses for output.

    Combines answer text with citations and metadata
    into a well-structured response.

    Example:
        formatter = ResponseFormatter()
        formatted = formatter.format(
            answer="The subcritical mass limit...",
            context=rag_context,
            confidence=0.95
        )
        print(formatted.full_response)
    """

    def __init__(
        self,
        output_format: OutputFormat = OutputFormat.MARKDOWN,
        include_confidence: bool = True,
        include_sources_section: bool = True,
    ):
        """
        Initialize response formatter.

        Args:
            output_format: Output format to use
            include_confidence: Whether to include confidence indicator
            include_sources_section: Whether to add sources section
        """
        self.output_format = output_format
        self.include_confidence = include_confidence
        self.include_sources_section = include_sources_section
        self.citation_formatter = CitationFormatter()
        self.citation_extractor = CitationExtractor()

    def format(
        self,
        answer: str,
        context: RAGContext,
        confidence: float | None = None,
        grounded: bool = True,
    ) -> FormattedResponse:
        """
        Format a RAG response.

        Args:
            answer: Raw answer text from LLM
            context: RAGContext with source citations
            confidence: Optional confidence score
            grounded: Whether response is grounded

        Returns:
            FormattedResponse with formatted text
        """
        # Clean up the answer
        answer = self._clean_answer(answer)

        # Extract and map citations
        citations = self._extract_used_citations(answer, context)

        # Build sources section
        sources_section = ""
        if self.include_sources_section and citations:
            sources_section = self.citation_formatter.format_list(
                citations, self.output_format
            )

        # Build full response
        full_response = self._build_full_response(
            answer=answer,
            sources_section=sources_section,
            confidence=confidence if self.include_confidence else None,
            grounded=grounded,
        )

        return FormattedResponse(
            answer=answer,
            citations=citations,
            sources_section=sources_section,
            full_response=full_response,
            confidence=confidence or context.avg_score,
            grounded=grounded,
            format=self.output_format,
        )

    def _clean_answer(self, answer: str) -> str:
        """Clean up answer text."""
        # Remove excessive whitespace
        answer = re.sub(r"\n{3,}", "\n\n", answer)
        answer = answer.strip()
        return answer

    def _extract_used_citations(
        self, answer: str, context: RAGContext
    ) -> list[ChunkCitation]:
        """Extract citations actually used in the answer."""
        # Get citations from answer text
        cited_text = self.citation_extractor.extract(answer)
        cited_doc_ids = self.citation_extractor.extract_doc_ids(cited_text)

        # Match to context citations
        used_citations = []
        for chunk in context.chunks:
            citation = chunk.to_citation()
            # Check if this doc was cited
            for doc_id in cited_doc_ids:
                if self._match_doc_id(doc_id, citation.doc_id):
                    used_citations.append(citation)
                    break

        # If no explicit citations matched, include all context sources
        if not used_citations:
            used_citations = context.citations

        return used_citations

    def _match_doc_id(self, cited: str, actual: str) -> bool:
        """Check if cited doc ID matches actual."""
        cited_lower = cited.lower().strip()
        actual_lower = actual.lower().strip()

        # Exact match
        if cited_lower == actual_lower:
            return True

        # Partial match
        if cited_lower in actual_lower or actual_lower in cited_lower:
            return True

        return False

    def _build_full_response(
        self,
        answer: str,
        sources_section: str,
        confidence: float | None,
        grounded: bool,
    ) -> str:
        """Build the complete formatted response."""
        parts = [answer]

        # Add confidence indicator if requested
        if confidence is not None:
            confidence_str = self._format_confidence(confidence, grounded)
            if confidence_str:
                parts.append("")
                parts.append(confidence_str)

        # Add sources section
        if sources_section:
            parts.append("")
            parts.append(sources_section)

        if self.output_format == OutputFormat.MARKDOWN:
            return "\n".join(parts)
        elif self.output_format == OutputFormat.HTML:
            return "<div class='rag-response'>" + \
                   "<br><br>".join(parts) + "</div>"
        else:
            return "\n".join(parts)

    def _format_confidence(self, confidence: float, grounded: bool) -> str:
        """Format confidence indicator."""
        if self.output_format == OutputFormat.MARKDOWN:
            if confidence >= 0.8 and grounded:
                return "_Response confidence: High_"
            elif confidence >= 0.5:
                return "_Response confidence: Medium - please verify critical details_"
            else:
                return "_Response confidence: Low - additional sources recommended_"

        elif self.output_format == OutputFormat.HTML:
            level = "high" if confidence >= 0.8 else "medium" if confidence >= 0.5 else "low"
            return f"<span class='confidence confidence-{level}'>Confidence: {level.title()}</span>"

        else:
            if confidence >= 0.8 and grounded:
                return "Confidence: High"
            elif confidence >= 0.5:
                return "Confidence: Medium"
            else:
                return "Confidence: Low"


@dataclass
class RAGResponseOutput:
    """
    Complete RAG response output for API consumers.

    Contains all information needed to display and use the response.
    """

    # Core response
    answer: str
    query: str

    # Citations
    citations: list[dict] = field(default_factory=list)
    sources_used: list[str] = field(default_factory=list)

    # Quality indicators
    confidence: float = 1.0
    grounded: bool = True
    validation_passed: bool = True
    validation_issues: list[str] = field(default_factory=list)

    # Metadata
    model: str = ""
    latency_ms: float = 0.0
    tokens_used: int = 0

    # Safety classification
    classification: str | None = None  # "green", "yellow", or "red"
    escalation_request_id: str | None = None

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        result = {
            "answer": self.answer,
            "query": self.query,
            "citations": self.citations,
            "sources_used": self.sources_used,
            "confidence": self.confidence,
            "grounded": self.grounded,
            "validation_passed": self.validation_passed,
            "validation_issues": self.validation_issues,
            "metadata": {
                "model": self.model,
                "latency_ms": self.latency_ms,
                "tokens_used": self.tokens_used,
            },
        }

        # Add safety classification if present
        if self.classification:
            result["classification"] = self.classification
        if self.escalation_request_id:
            result["escalation_request_id"] = self.escalation_request_id

        return result


def format_rag_response(
    answer: str,
    context: RAGContext,
    confidence: float = 1.0,
    grounded: bool = True,
    format_type: OutputFormat = OutputFormat.MARKDOWN,
) -> FormattedResponse:
    """
    Convenience function to format a RAG response.

    Args:
        answer: Raw answer text
        context: RAGContext with sources
        confidence: Confidence score
        grounded: Whether response is grounded
        format_type: Output format

    Returns:
        FormattedResponse
    """
    formatter = ResponseFormatter(output_format=format_type)
    return formatter.format(answer, context, confidence, grounded)
