"""
SAGE RAG Pipeline

Main orchestration class that connects retrieval, generation,
and validation into a complete RAG system for NCS.
"""

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

from sage.core.abstention import (
    AbstentionConfig,
    AbstentionHandler,
    AbstentionResult,
    format_abstention_response,
)
from sage.core.classification import (
    ClassificationResult,
    OutputClassification,
    OutputClassifier,
)
from sage.core.escalation import (
    EscalationConfig,
    EscalationEvaluation,
    EscalationWorkflow,
    format_escalation_notice,
)
from sage.data.search import SAGESearch, SearchConfig, SearchResponse
from sage.data.storage.citations import ChunkCitation
from sage.data.storage.vector_store import QdrantMode
from sage.llm import LLMConfig, LLMProvider
from sage.rag.context import ContextBuilder, RAGContext
from sage.rag.formatting import (
    FormattedResponse,
    OutputFormat,
    RAGResponseOutput,
    ResponseFormatter,
)
from sage.rag.generation import GenerationEngine
from sage.rag.prompts import SystemPrompt
from sage.rag.validation import (
    ResponseValidator,
    ValidationLevel,
    ValidationResult,
)

logger = logging.getLogger(__name__)


@dataclass
class RAGConfig:
    """
    Configuration for the SAGE RAG pipeline.

    Combines search and LLM configuration for complete RAG setup.
    """

    # Qdrant settings (inherited from SearchConfig)
    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

    # LLM settings
    anthropic_api_key: str | None = None
    primary_llm_provider: LLMProvider = LLMProvider.ANTHROPIC
    primary_llm_model: str | None = None
    fallback_llm_provider: LLMProvider | None = LLMProvider.OPENAI
    llm_temperature: float = 0.3
    llm_max_tokens: int = 4096

    # Search settings
    default_top_k: int = 10
    use_hybrid_search: bool = True
    chunk_size: int = 512
    max_chunk_size: int = 1024

    # RAG settings
    max_context_length: int = 8000
    context_score_threshold: float = 0.3
    validation_level: ValidationLevel = ValidationLevel.MODERATE

    # Output settings
    output_format: OutputFormat = OutputFormat.MARKDOWN
    include_confidence: bool = True
    include_sources: bool = True

    # Safety settings
    enable_abstention: bool = True
    enable_escalation: bool = True
    abstention_config: AbstentionConfig | None = None
    escalation_config: EscalationConfig | None = None

    def to_search_config(self) -> SearchConfig:
        """Convert to SearchConfig."""
        return SearchConfig(
            qdrant_mode=self.qdrant_mode,
            qdrant_host=self.qdrant_host,
            qdrant_port=self.qdrant_port,
            qdrant_url=self.qdrant_url,
            qdrant_api_key=self.qdrant_api_key,
            collection_name=self.collection_name,
            openai_api_key=self.openai_api_key,
            embedding_model=self.embedding_model,
            embedding_dimensions=self.embedding_dimensions,
            chunk_size=self.chunk_size,
            max_chunk_size=self.max_chunk_size,
            default_top_k=self.default_top_k,
            use_hybrid_search=self.use_hybrid_search,
        )

    def to_llm_config(self) -> LLMConfig:
        """Convert to LLMConfig."""
        return LLMConfig(
            anthropic_api_key=self.anthropic_api_key,
            openai_api_key=self.openai_api_key,
            primary_provider=self.primary_llm_provider,
            primary_model=self.primary_llm_model,
            fallback_provider=self.fallback_llm_provider,
            temperature=self.llm_temperature,
            max_tokens=self.llm_max_tokens,
        )


@dataclass
class RAGResponse:
    """
    Complete RAG response with all metadata.

    Contains the answer, citations, validation results,
    and generation metadata.
    """

    # Core response
    answer: str
    query: str

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

    # Quality indicators
    confidence: float = 1.0
    grounded: bool = True

    # Validation
    validation: ValidationResult | None = None

    # Safety checks
    abstention: AbstentionResult | None = None
    escalation: EscalationEvaluation | None = None
    classification: ClassificationResult | None = None
    escalation_request_id: str | None = None

    # Context used
    context: RAGContext | None = None

    # Generation metadata
    model: str = ""
    provider: LLMProvider = LLMProvider.ANTHROPIC
    latency_ms: float = 0.0
    prompt_tokens: int = 0
    completion_tokens: int = 0

    @property
    def is_valid(self) -> bool:
        """Check if response passed validation."""
        return self.validation.is_valid if self.validation else True

    @property
    def total_tokens(self) -> int:
        """Total tokens used."""
        return self.prompt_tokens + self.completion_tokens

    @property
    def output_classification(self) -> OutputClassification:
        """Get the output classification level."""
        if self.classification:
            return self.classification.classification
        return OutputClassification.GREEN  # Default if not classified

    @property
    def should_abstain(self) -> bool:
        """Check if response was abstained."""
        return self.abstention.should_abstain if self.abstention else False

    @property
    def requires_escalation(self) -> bool:
        """Check if response requires escalation."""
        return self.escalation.should_escalate if self.escalation else False

    def to_output(self) -> RAGResponseOutput:
        """Convert to output format for API."""
        return RAGResponseOutput(
            answer=self.answer,
            query=self.query,
            citations=[c.to_dict() for c in self.citations],
            sources_used=self.sources_used,
            confidence=self.confidence,
            grounded=self.grounded,
            validation_passed=self.is_valid,
            validation_issues=[str(i) for i in (self.validation.issues if self.validation else [])],
            model=self.model,
            latency_ms=self.latency_ms,
            tokens_used=self.total_tokens,
            classification=self.output_classification.value if self.classification else None,
            escalation_request_id=self.escalation_request_id,
        )


class SAGERAGPipeline:
    """
    Complete RAG pipeline for SAGE NCS.

    Orchestrates the full RAG workflow:
    1. Query processing
    2. Document retrieval
    3. Context building
    4. LLM generation
    5. Response validation
    6. Output formatting

    Example:
        config = RAGConfig(
            qdrant_mode=QdrantMode.MEMORY,
            openai_api_key="sk-...",
            anthropic_api_key="sk-ant-...",
        )

        async with SAGERAGPipeline(config) as pipeline:
            # Index documents
            await pipeline.index_document(Path("nureg.pdf"), doc_id="NUREG-1520")

            # Query
            response = await pipeline.query("What is the subcritical limit for U-235?")
            print(response.answer)
            print(f"Sources: {response.sources_used}")
    """

    def __init__(self, config: RAGConfig):
        """Initialize RAG pipeline."""
        self.config = config
        self._search: SAGESearch | None = None
        self._generation_engine: GenerationEngine | None = None
        self._context_builder: ContextBuilder | None = None
        self._validator: ResponseValidator | None = None
        self._formatter: ResponseFormatter | None = None

        # Safety handlers
        self._abstention_handler: AbstentionHandler | None = None
        self._escalation_workflow: EscalationWorkflow | None = None
        self._classifier: OutputClassifier | None = None

        self._initialized = False

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

        logger.info("Initializing SAGE RAG pipeline...")

        # Initialize search
        search_config = self.config.to_search_config()
        self._search = SAGESearch(search_config)
        await self._search.initialize()
        logger.info("Search engine initialized")

        # Initialize generation engine
        llm_config = self.config.to_llm_config()
        self._generation_engine = GenerationEngine(llm_config)
        await self._generation_engine.initialize()
        logger.info("Generation engine initialized")

        # Initialize context builder
        self._context_builder = ContextBuilder(
            max_context_length=self.config.max_context_length,
            score_threshold=self.config.context_score_threshold,
        )

        # Initialize validator
        self._validator = ResponseValidator(
            level=self.config.validation_level,
        )

        # Initialize formatter
        self._formatter = ResponseFormatter(
            output_format=self.config.output_format,
            include_confidence=self.config.include_confidence,
            include_sources_section=self.config.include_sources,
        )

        # Initialize safety handlers
        if self.config.enable_abstention:
            self._abstention_handler = AbstentionHandler(
                config=self.config.abstention_config
            )
            logger.info("Abstention handler initialized")

        if self.config.enable_escalation:
            self._escalation_workflow = EscalationWorkflow(
                config=self.config.escalation_config
            )
            logger.info("Escalation workflow initialized")

        # Always initialize classifier
        self._classifier = OutputClassifier()
        logger.info("Output classifier initialized")

        self._initialized = True
        logger.info("SAGE RAG pipeline ready")

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

    async def __aenter__(self) -> "SAGERAGPipeline":
        """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,
    ) -> dict:
        """
        Index a document into the knowledge base.

        Args:
            file_path: Path to document (PDF)
            doc_id: Document identifier
            source: Document source
            title: Optional title

        Returns:
            Indexing statistics
        """
        if not self._initialized:
            await self.initialize()

        indexed = await self._search.index_document(
            file_path=file_path,
            doc_id=doc_id,
            source=source,
            title=title,
        )

        return {
            "doc_id": indexed.doc_id,
            "chunks": indexed.total_chunks,
            "pages": indexed.total_pages,
        }

    async def query(
        self,
        question: str,
        top_k: int | None = None,
        filters: dict | None = None,
        system_prompt: SystemPrompt | None = None,
        task_type: str = "general",
        validate: bool = True,
        conversation_history: list[dict] | None = None,
    ) -> RAGResponse:
        """
        Execute a RAG query.

        Args:
            question: User's question
            top_k: Number of chunks to retrieve
            filters: Optional search filters
            system_prompt: Optional custom system prompt
            task_type: Task type for prompt selection
            validate: Whether to validate response
            conversation_history: Optional conversation history

        Returns:
            RAGResponse with answer and metadata
        """
        if not self._initialized:
            await self.initialize()

        logger.info(f"Processing query: {question[:100]}...")

        # Step 1: Retrieve relevant documents
        search_response = await self._search.search(
            query=question,
            top_k=top_k or self.config.default_top_k,
            filters=filters,
        )
        logger.info(f"Retrieved {len(search_response.results)} chunks")

        # Step 2: Build context
        context = self._context_builder.build(
            query=question,
            search_results=search_response.results,
        )
        logger.info(
            f"Built context with {len(context.chunks)} chunks, "
            f"confidence: {context.confidence_level}"
        )

        # Step 3: Generate response
        generation_result = await self._generation_engine.generate(
            context=context,
            system_prompt=system_prompt,
            task_type=task_type,
            conversation_history=conversation_history,
        )
        logger.info(
            f"Generated response in {generation_result.latency_ms:.0f}ms "
            f"({generation_result.completion_tokens} tokens)"
        )

        # Step 4: Validate response
        validation_result = None
        if validate:
            validation_result = self._validator.validate(
                generation_result.content,
                context,
            )
            if not validation_result.is_valid:
                logger.warning(
                    f"Validation issues: {[str(i) for i in validation_result.issues]}"
                )

        # Step 5: Build response
        response = RAGResponse(
            answer=generation_result.content,
            query=question,
            citations=context.citations,
            sources_used=context.source_doc_ids,
            confidence=context.avg_score,
            grounded=validation_result.grounded if validation_result else True,
            validation=validation_result,
            context=context,
            model=generation_result.model,
            provider=generation_result.provider,
            latency_ms=generation_result.latency_ms,
            prompt_tokens=generation_result.prompt_tokens,
            completion_tokens=generation_result.completion_tokens,
        )

        return response

    async def query_with_safety(
        self,
        question: str,
        user_id: str | None = None,
        context_metadata: dict | None = None,
        top_k: int | None = None,
        filters: dict | None = None,
        system_prompt: SystemPrompt | None = None,
        task_type: str = "general",
        validate: bool = True,
        conversation_history: list[dict] | None = None,
    ) -> RAGResponse:
        """
        Execute a RAG query with full safety checks (abstention, escalation, classification).

        This is the recommended method for production use as it includes all
        safety gates before returning a response.

        Args:
            question: User's question
            user_id: Optional user identifier for escalation tracking
            context_metadata: Optional context (affects_safety_basis, etc.)
            top_k: Number of chunks to retrieve
            filters: Optional search filters
            system_prompt: Optional custom system prompt
            task_type: Task type for prompt selection
            validate: Whether to validate response
            conversation_history: Optional conversation history

        Returns:
            RAGResponse with safety checks applied

        Example:
            response = await pipeline.query_with_safety(
                question="What about this new configuration?",
                user_id="user123",
                context_metadata={"affects_safety_basis": True}
            )

            if response.should_abstain:
                print("Cannot answer:", response.abstention.reason)
            elif response.requires_escalation:
                print("Pending review:", response.escalation_request_id)
            else:
                print(response.answer)
        """
        if not self._initialized:
            await self.initialize()

        context_metadata = context_metadata or {}
        logger.info(f"Processing query with safety checks: {question[:100]}...")

        # Step 1: Retrieve relevant documents
        search_response = await self._search.search(
            query=question,
            top_k=top_k or self.config.default_top_k,
            filters=filters,
        )
        logger.info(f"Retrieved {len(search_response.results)} chunks")

        # Convert search results to format expected by abstention handler
        rag_results = [
            {
                "doc_id": r.doc_id,
                "doc_title": r.title or r.doc_id,
                "content": r.content,
                "confidence": r.score,
                "section": r.metadata.get("section") if r.metadata else None,
            }
            for r in search_response.results
        ]

        # Step 2: Check for abstention
        abstention_result = None
        if self._abstention_handler:
            abstention_result = self._abstention_handler.evaluate(
                query=question,
                rag_results=rag_results,
                context=context_metadata,
            )

            if abstention_result.should_abstain:
                logger.info(
                    f"Abstention triggered: {[t.value for t in abstention_result.triggers]}"
                )

                # Build abstention response
                abstention_response = format_abstention_response(abstention_result)

                # Classify as RED since we're abstaining
                classification = self._classifier.classify(
                    confidence=0.0,
                    grounded=False,
                    abstention_result=abstention_result,
                )

                return RAGResponse(
                    answer=abstention_response,
                    query=question,
                    confidence=0.0,
                    grounded=False,
                    abstention=abstention_result,
                    classification=classification,
                )

        # Step 3: Build context
        context = self._context_builder.build(
            query=question,
            search_results=search_response.results,
        )
        logger.info(
            f"Built context with {len(context.chunks)} chunks, "
            f"confidence: {context.confidence_level}"
        )

        # Step 4: Generate response
        generation_result = await self._generation_engine.generate(
            context=context,
            system_prompt=system_prompt,
            task_type=task_type,
            conversation_history=conversation_history,
        )
        logger.info(
            f"Generated response in {generation_result.latency_ms:.0f}ms "
            f"({generation_result.completion_tokens} tokens)"
        )

        # Step 5: Validate response
        validation_result = None
        if validate:
            validation_result = self._validator.validate(
                generation_result.content,
                context,
            )
            if not validation_result.is_valid:
                logger.warning(
                    f"Validation issues: {[str(i) for i in validation_result.issues]}"
                )

        # Step 6: Check for escalation
        escalation_result = None
        escalation_request_id = None

        if self._escalation_workflow:
            escalation_result = self._escalation_workflow.evaluate(
                query=question,
                response=generation_result.content,
                confidence=context.avg_score,
                context=context_metadata,
            )

            if escalation_result.should_escalate:
                logger.info(
                    f"Escalation triggered: {[t.value for t in escalation_result.triggers]}, "
                    f"priority: {escalation_result.priority.value}"
                )

                # Create escalation request
                escalation_request_id = self._escalation_workflow.create_escalation(
                    query=question,
                    response=generation_result.content,
                    triggers=escalation_result.triggers,
                    user_id=user_id,
                    context=context_metadata,
                )

        # Step 7: Classify output
        classification_result = self._classifier.classify(
            confidence=context.avg_score,
            grounded=validation_result.grounded if validation_result else True,
            grounding_score=validation_result.grounding_score if validation_result else 1.0,
            citations_valid=len(validation_result.citations_missing) == 0 if validation_result else True,
            citation_count=len(validation_result.citations_verified) if validation_result else 0,
            has_safety_warnings=validation_result.has_safety_warnings if validation_result else False,
            abstention_result=abstention_result,
            escalation_result=escalation_result,
        )
        logger.info(f"Classification: {classification_result.classification.value}")

        # Step 8: Build response
        # If escalated and RED, modify the answer
        answer = generation_result.content
        if (
            classification_result.classification == OutputClassification.RED
            and escalation_request_id
        ):
            # Response is held - provide notice instead
            request = self._escalation_workflow.queue.get(escalation_request_id)
            if request and request.response_held:
                answer = format_escalation_notice(request)

        response = RAGResponse(
            answer=answer,
            query=question,
            citations=context.citations,
            sources_used=context.source_doc_ids,
            confidence=context.avg_score,
            grounded=validation_result.grounded if validation_result else True,
            validation=validation_result,
            abstention=abstention_result,
            escalation=escalation_result,
            classification=classification_result,
            escalation_request_id=escalation_request_id,
            context=context,
            model=generation_result.model,
            provider=generation_result.provider,
            latency_ms=generation_result.latency_ms,
            prompt_tokens=generation_result.prompt_tokens,
            completion_tokens=generation_result.completion_tokens,
        )

        return response

    async def query_formatted(
        self,
        question: str,
        **kwargs,
    ) -> FormattedResponse:
        """
        Execute query and return formatted response.

        Args:
            question: User's question
            **kwargs: Additional arguments for query()

        Returns:
            FormattedResponse ready for display
        """
        response = await self.query(question, **kwargs)

        return self._formatter.format(
            answer=response.answer,
            context=response.context,
            confidence=response.confidence,
            grounded=response.grounded,
        )

    async def search_only(
        self,
        query: str,
        top_k: int | None = None,
        filters: dict | None = None,
    ) -> SearchResponse:
        """
        Perform search without generation.

        Useful for debugging or retrieving raw results.

        Args:
            query: Search query
            top_k: Number of results
            filters: Optional filters

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

        return await self._search.search(
            query=query,
            top_k=top_k or self.config.default_top_k,
            filters=filters,
        )

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

        search_stats = await self._search.get_stats()

        return {
            "search": search_stats,
            "llm": {
                "primary_model": self._generation_engine.primary_model,
                "fallback_model": self._generation_engine.fallback_model,
            },
            "config": {
                "validation_level": self.config.validation_level.value,
                "max_context_length": self.config.max_context_length,
                "score_threshold": self.config.context_score_threshold,
            },
        }


# Convenience alias
SAGERAG = SAGERAGPipeline


async def create_rag_pipeline(
    openai_api_key: str,
    anthropic_api_key: str | None = None,
    qdrant_mode: QdrantMode = QdrantMode.MEMORY,
    **kwargs,
) -> SAGERAGPipeline:
    """
    Factory function to create and initialize a RAG pipeline.

    Args:
        openai_api_key: OpenAI API key (required for embeddings)
        anthropic_api_key: Anthropic API key (optional, uses OpenAI if not provided)
        qdrant_mode: Vector store mode
        **kwargs: Additional RAGConfig parameters

    Returns:
        Initialized SAGERAGPipeline
    """
    # Use OpenAI as primary if no Anthropic key
    primary_provider = LLMProvider.ANTHROPIC if anthropic_api_key else LLMProvider.OPENAI

    config = RAGConfig(
        openai_api_key=openai_api_key,
        anthropic_api_key=anthropic_api_key,
        qdrant_mode=qdrant_mode,
        primary_llm_provider=primary_provider,
        **kwargs,
    )

    pipeline = SAGERAGPipeline(config)
    await pipeline.initialize()
    return pipeline
