"""
Main cheat sheet generation pipeline.

Orchestrates the full RAG pipeline from document ingestion
to final PDF generation.
"""

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
import uuid

from .ingestion import DocumentIngestionService, ParsedDocument
from .chunking import ChunkingService, Chunk
from .indexing import IndexingService, EmbeddingModel, create_embedding_model
from .retrieval import RetrievalService, RetrievalResult
from .models import LLMClient, GenerationConfig, create_llm_client
from .prompts import get_prompt

from ..infra.settings import Settings, get_default_settings
from ..infra.vectordb import create_vector_store, VectorStore
from ..infra.logging import get_logger, LogContext
from ..infra.tracing import trace_function, Tracer

logger = get_logger(__name__)


@dataclass
class CheatSheetContent:
    """Structured content for a cheat sheet."""
    title: str
    overview: str = ""
    concepts: List[Dict[str, Any]] = field(default_factory=list)
    formulas: List[Dict[str, Any]] = field(default_factory=list)
    models: List[Dict[str, Any]] = field(default_factory=list)
    algorithms: List[Dict[str, Any]] = field(default_factory=list)
    examples: List[Dict[str, Any]] = field(default_factory=list)
    connections: List[Dict[str, Any]] = field(default_factory=list)
    recap: str = ""
    references: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "title": self.title,
            "overview": self.overview,
            "concepts": self.concepts,
            "formulas": self.formulas,
            "models": self.models,
            "algorithms": self.algorithms,
            "examples": self.examples,
            "connections": self.connections,
            "recap": self.recap,
            "references": self.references
        }


@dataclass
class PipelineResult:
    """Result from pipeline execution."""
    success: bool
    content: Optional[CheatSheetContent] = None
    output_path: Optional[Path] = None
    latex_source: Optional[str] = None
    errors: List[str] = field(default_factory=list)
    stats: Dict[str, Any] = field(default_factory=dict)


class CheatSheetPipeline:
    """
    Main pipeline for generating cheat sheets from documents.
    
    Orchestrates:
    1. Document ingestion and parsing
    2. Chunking and indexing
    3. RAG-based content generation
    4. LaTeX document assembly
    5. PDF compilation
    """
    
    def __init__(
        self,
        settings: Optional[Settings] = None,
        vector_store: Optional[VectorStore] = None,
        embedding_model: Optional[EmbeddingModel] = None,
        llm_client: Optional[LLMClient] = None
    ):
        """
        Initialize the pipeline.
        
        Args:
            settings: Pipeline settings
            vector_store: Pre-configured vector store
            embedding_model: Pre-configured embedding model
            llm_client: Pre-configured LLM client
        """
        self.settings = settings or get_default_settings()
        
        # Initialize components
        self.ingestion_service = DocumentIngestionService(self.settings.ingestion)
        self.chunking_service = ChunkingService(self.settings.chunking)
        
        # Initialize or use provided embedding model
        self.embedding_model = embedding_model or create_embedding_model(self.settings)
        
        # Initialize or use provided vector store
        self.vector_store = vector_store or create_vector_store(self.settings)
        
        # Initialize indexing service
        self.indexing_service = IndexingService(
            vector_store=self.vector_store,
            embedding_model=self.embedding_model,
            settings=self.settings
        )
        
        # Initialize retrieval service
        self.retrieval_service = RetrievalService(
            vector_store=self.vector_store,
            embedding_model=self.embedding_model,
            settings=self.settings.retrieval
        )
        
        # Initialize or use provided LLM client
        self.llm_client = llm_client or create_llm_client(self.settings)
        
        # Initialize tracer
        if self.settings.observability.tracing.enabled:
            self.tracer = Tracer(
                service_name=self.settings.observability.tracing.service_name,
                endpoint=self.settings.observability.tracing.endpoint,
                enabled=True
            )
        else:
            self.tracer = None
        
        # State
        self._document: Optional[ParsedDocument] = None
        self._chunks: List[Chunk] = []
        self._indexed = False
    
    @trace_function(name="pipeline_ingest")
    def ingest(self, source: Union[str, Path]) -> ParsedDocument:
        """
        Ingest a document from file.
        
        Args:
            source: Path to document file
            
        Returns:
            ParsedDocument
        """
        source = Path(source)
        logger.info(f"Ingesting document: {source}")
        
        self._document = self.ingestion_service.ingest(source)
        
        logger.info(
            f"Ingested document: {self._document.metadata.title or source.name} "
            f"({self._document.metadata.word_count} words, "
            f"{len(self._document.content_blocks)} blocks)"
        )
        
        return self._document
    
    @trace_function(name="pipeline_index")
    def index(
        self,
        document: Optional[ParsedDocument] = None,
        strategy: str = "hierarchical",
        collection_name: Optional[str] = None
    ) -> List[str]:
        """
        Index a document for retrieval.
        
        Args:
            document: Document to index (uses last ingested if not provided)
            strategy: Chunking strategy
            collection_name: Optional name for the vector store collection
            
        Returns:
            List of chunk IDs
        """
        document = document or self._document
        if document is None:
            raise ValueError("No document to index. Call ingest() first.")
        
        # Re-initialize vector store with custom collection name if provided
        if collection_name:
            self.settings.vectordb.collection_name = collection_name
            self.vector_store = create_vector_store(self.settings)
            self.indexing_service = IndexingService(
                vector_store=self.vector_store,
                embedding_model=self.embedding_model,
                settings=self.settings
            )
            self.retrieval_service = RetrievalService(
                vector_store=self.vector_store,
                embedding_model=self.embedding_model,
                settings=self.settings.retrieval
            )
        
        logger.info(f"Indexing document with strategy: {strategy}")
        
        # Chunk the document
        self._chunks = self.chunking_service.chunk(document, strategy=strategy)
        
        # Index chunks
        chunk_ids = self.indexing_service.index_chunks(self._chunks)
        
        # Build sparse index for hybrid retrieval
        documents = [
            {"content": c.content, "metadata": c.to_dict()}
            for c in self._chunks
        ]
        self.retrieval_service.build_sparse_index(documents)
        
        self._indexed = True
        
        logger.info(f"Indexed {len(chunk_ids)} chunks")
        return chunk_ids
    
    @trace_function(name="pipeline_generate_content")
    def generate_content(
        self,
        title: Optional[str] = None,
        config: Optional[GenerationConfig] = None
    ) -> CheatSheetContent:
        """
        Generate cheat sheet content using RAG.
        
        Args:
            title: Title for the cheat sheet
            config: Generation configuration
            
        Returns:
            CheatSheetContent
        """
        if not self._indexed:
            raise ValueError("Document not indexed. Call index() first.")
        
        document = self._document
        title = title or document.metadata.title or "Cheat Sheet"
        config = config or GenerationConfig(
            max_tokens=self.settings.llm.max_tokens,
            temperature=self.settings.llm.temperature,
            top_p=self.settings.llm.top_p
        )
        
        logger.info(f"Generating cheat sheet content: {title}")
        
        content = CheatSheetContent(title=title)
        
        # Generate each section
        with LogContext(document_id=document.metadata.source_path):
            # 1. Generate overview
            content.overview = self._generate_overview(title, config)
            
            # 2. Extract concepts
            content.concepts = self._extract_concepts(config)
            
            # 3. Extract formulas
            content.formulas = self._extract_formulas(config)
            
            # 4. Summarize models
            content.models = self._extract_models(config)
            
            # 5. Extract algorithms
            content.algorithms = self._extract_algorithms(config)
            
            # 6. Generate connections
            content.connections = self._generate_connections(title, config)
            
            # 7. Generate recap
            content.recap = self._generate_recap(content, config)
            
            # 8. Add references
            if document.metadata.authors:
                content.references.append(
                    f"{', '.join(document.metadata.authors)}. {title}."
                )
        
        logger.info("Content generation complete")
        return content
    
    def _generate_overview(self, title: str, config: GenerationConfig) -> str:
        """Generate overview section."""
        logger.info("Generating overview")
        
        # Get high-level context
        results = self.retrieval_service.retrieve(
            f"introduction overview {title}",
            top_k=5
        )
        context = "\n\n".join(r.content for r in results)
        
        # Get subtopics from document structure
        subtopics = []
        if self._document and self._document.table_of_contents:
            subtopics = [item["title"] for item in self._document.table_of_contents[:10]]
        
        prompt = get_prompt("overview")(
            topic=title,
            subtopics=", ".join(subtopics) if subtopics else "various aspects"
        )
        
        result = self.llm_client.generate(prompt, config)
        return result.text
    
    def _extract_concepts(self, config: GenerationConfig) -> List[Dict[str, Any]]:
        """Extract key concepts."""
        logger.info("Extracting concepts")
        
        concepts = []
        
        # Retrieve concept-rich content
        results = self.retrieval_service.retrieve_narrative(
            "key concepts definitions terminology",
            top_k=10
        )
        
        if not results:
            return concepts
        
        context = "\n\n".join(r.content for r in results)
        
        prompt = get_prompt("concept_extraction")(context=context)
        result = self.llm_client.generate(prompt, config)
        
        # Parse concepts from response
        concepts = self._parse_concepts(result.text)
        
        # Limit to max concepts
        max_concepts = self.settings.generation.concepts.max_concepts
        return concepts[:max_concepts]
    
    def _parse_concepts(self, text: str) -> List[Dict[str, Any]]:
        """Parse concept extraction response."""
        concepts = []
        current = {}
        
        for line in text.split("\n"):
            line = line.strip()
            
            if line.startswith("CONCEPT:"):
                if current:
                    concepts.append(current)
                current = {"name": line[8:].strip()}
            elif line.startswith("DEFINITION:"):
                current["definition"] = line[11:].strip()
            elif line.startswith("INTUITION:"):
                current["intuition"] = line[10:].strip()
            elif line.startswith("RELATED:"):
                current["related"] = [r.strip() for r in line[8:].split(",")]
        
        if current:
            concepts.append(current)
        
        return concepts
    
    def _extract_formulas(self, config: GenerationConfig) -> List[Dict[str, Any]]:
        """Extract mathematical formulas."""
        logger.info("Extracting formulas")
        
        formulas = []
        
        # Retrieve math content
        results = self.retrieval_service.retrieve_math(
            "equations formulas mathematical",
            top_k=15
        )
        
        if not results:
            # Try general retrieval
            results = self.retrieval_service.retrieve(
                "equation formula derivation",
                top_k=10
            )
        
        if not results:
            return formulas
        
        context = "\n\n".join(r.content for r in results)
        
        prompt = get_prompt("formula_extraction")(context=context)
        result = self.llm_client.generate(prompt, config)
        
        # Parse formulas
        formulas = self._parse_formulas(result.text)
        
        # Limit
        max_formulas = self.settings.generation.formulas.max_formulas
        return formulas[:max_formulas]
    
    def _parse_formulas(self, text: str) -> List[Dict[str, Any]]:
        """Parse formula extraction response."""
        formulas = []
        current = {}
        
        for line in text.split("\n"):
            line = line.strip()
            
            if line.startswith("FORMULA:"):
                if current and "latex" in current:
                    formulas.append(current)
                current = {"name": line[8:].strip()}
            elif line.startswith("LATEX:"):
                current["latex"] = line[6:].strip()
            elif line.startswith("CONTEXT:"):
                current["context"] = line[8:].strip()
            elif line.startswith("DEPENDS_ON:"):
                current["depends_on"] = [d.strip() for d in line[11:].split(",")]
        
        if current and "latex" in current:
            formulas.append(current)
        
        return formulas
    
    def _extract_models(self, config: GenerationConfig) -> List[Dict[str, Any]]:
        """Extract and summarize models/theories."""
        logger.info("Extracting models")
        
        models = []
        
        # Retrieve model descriptions
        results = self.retrieval_service.retrieve(
            "model architecture training objective loss function",
            top_k=10
        )
        
        if not results:
            return models
        
        context = "\n\n".join(r.content for r in results)
        
        # Try to identify model names from context
        model_name = self._document.metadata.title or "Main Model"
        
        prompt = get_prompt("model_extraction")(
            context=context,
            model_name=model_name
        )
        result = self.llm_client.generate(prompt, config)
        
        # Parse as a single model summary
        models.append({
            "name": model_name,
            "summary": result.text
        })
        
        return models
    
    def _extract_algorithms(self, config: GenerationConfig) -> List[Dict[str, Any]]:
        """Extract algorithms and pseudocode."""
        logger.info("Extracting algorithms")
        
        algorithms = []
        
        # Retrieve algorithm content
        results = self.retrieval_service.retrieve_code(
            "algorithm procedure steps",
            top_k=10
        )
        
        if not results:
            results = self.retrieval_service.retrieve(
                "algorithm procedure method steps",
                top_k=8
            )
        
        if not results:
            return algorithms
        
        context = "\n\n".join(r.content for r in results)
        
        prompt = get_prompt("algorithm_extraction")(context=context)
        result = self.llm_client.generate(prompt, config)
        
        # Parse algorithms
        algorithms = self._parse_algorithms(result.text)
        
        return algorithms
    
    def _parse_algorithms(self, text: str) -> List[Dict[str, Any]]:
        """Parse algorithm extraction response."""
        algorithms = []
        current = {}
        in_pseudocode = False
        pseudocode_lines = []
        
        for line in text.split("\n"):
            if line.startswith("ALGORITHM:"):
                if current:
                    if pseudocode_lines:
                        current["pseudocode"] = "\n".join(pseudocode_lines)
                    algorithms.append(current)
                current = {"name": line[10:].strip()}
                pseudocode_lines = []
                in_pseudocode = False
            elif line.startswith("PURPOSE:"):
                current["purpose"] = line[8:].strip()
            elif line.startswith("INPUT:"):
                current["input"] = line[6:].strip()
            elif line.startswith("OUTPUT:"):
                current["output"] = line[7:].strip()
            elif line.startswith("PSEUDOCODE:") or line.strip() == "```":
                in_pseudocode = not in_pseudocode
            elif in_pseudocode:
                pseudocode_lines.append(line)
            elif line.startswith("COMPLEXITY:"):
                current["complexity"] = line[11:].strip()
            elif line.startswith("KEY INSIGHT:"):
                current["insight"] = line[12:].strip()
        
        if current:
            if pseudocode_lines:
                current["pseudocode"] = "\n".join(pseudocode_lines)
            algorithms.append(current)
        
        return algorithms
    
    def _generate_connections(
        self,
        title: str,
        config: GenerationConfig
    ) -> List[Dict[str, Any]]:
        """Generate connections to related topics."""
        logger.info("Generating connections")
        
        # Get summary of content
        results = self.retrieval_service.retrieve(
            "related work connections applications",
            top_k=5
        )
        summary = "\n".join(r.content[:500] for r in results)
        
        prompt = get_prompt("connections")(
            topic=title,
            summary=summary,
            related_areas="machine learning, statistics, optimization, linear algebra"
        )
        
        result = self.llm_client.generate(prompt, config)
        
        # Parse connections
        connections = []
        for line in result.text.split("\n"):
            line = line.strip()
            if line.startswith("- ") and ":" in line:
                parts = line[2:].split(":", 1)
                connections.append({
                    "topic": parts[0].strip(),
                    "description": parts[1].strip() if len(parts) > 1 else ""
                })
        
        return connections
    
    def _generate_recap(
        self,
        content: CheatSheetContent,
        config: GenerationConfig
    ) -> str:
        """Generate one-page recap."""
        logger.info("Generating recap")
        
        # Compile summary of content
        summary_parts = []
        
        if content.concepts:
            summary_parts.append("Key Concepts: " + ", ".join(
                c.get("name", "") for c in content.concepts[:10]
            ))
        
        if content.formulas:
            summary_parts.append("Key Formulas: " + ", ".join(
                f.get("name", "") for f in content.formulas[:10]
            ))
        
        if content.algorithms:
            summary_parts.append("Algorithms: " + ", ".join(
                a.get("name", "") for a in content.algorithms[:5]
            ))
        
        prompt = get_prompt("recap")(content="\n".join(summary_parts))
        result = self.llm_client.generate(prompt, config)
        
        return result.text
    
    @trace_function(name="pipeline_build")
    def build(
        self,
        source: Union[str, Path],
        output_path: Optional[Union[str, Path]] = None,
        title: Optional[str] = None,
        compile_pdf: bool = True
    ) -> PipelineResult:
        """
        Build a complete cheat sheet from a document.
        
        Args:
            source: Path to source document
            output_path: Output path for PDF/LaTeX
            title: Optional title override
            compile_pdf: Whether to compile to PDF
            
        Returns:
            PipelineResult
        """
        run_id = str(uuid.uuid4())[:8]
        errors = []
        stats = {}
        
        logger.info(f"Starting cheat sheet build [run={run_id}]")
        
        try:
            with LogContext(request_id=run_id):
                # 1. Ingest
                document = self.ingest(source)
                stats["word_count"] = document.metadata.word_count
                stats["block_count"] = len(document.content_blocks)
                
                # 2. Index
                chunk_ids = self.index(document)
                stats["chunk_count"] = len(chunk_ids)
                
                # 3. Generate content
                content = self.generate_content(title=title)
                stats["concept_count"] = len(content.concepts)
                stats["formula_count"] = len(content.formulas)
                stats["algorithm_count"] = len(content.algorithms)
                
                # 4. Build LaTeX
                from ..generation.latex_builder import LaTeXBuilder
                builder = LaTeXBuilder(self.settings.latex)
                latex_source = builder.build(content)
                
                # 5. Compile PDF if requested
                output_path = Path(output_path) if output_path else Path(f"{content.title.replace(' ', '_')}_cheatsheet.pdf")
                
                if compile_pdf:
                    pdf_path = builder.compile_pdf(latex_source, output_path)
                    logger.info(f"Generated PDF: {pdf_path}")
                else:
                    # Save LaTeX source
                    tex_path = output_path.with_suffix(".tex")
                    tex_path.write_text(latex_source)
                    pdf_path = None
                    logger.info(f"Generated LaTeX: {tex_path}")
                
                return PipelineResult(
                    success=True,
                    content=content,
                    output_path=pdf_path or tex_path,
                    latex_source=latex_source,
                    stats=stats
                )
                
        except Exception as e:
            logger.error(f"Pipeline error: {e}")
            errors.append(str(e))
            return PipelineResult(
                success=False,
                errors=errors,
                stats=stats
            )
    
    def clear(self) -> None:
        """Clear all state and indexed data."""
        self.vector_store.clear()
        self._document = None
        self._chunks = []
        self._indexed = False
        logger.info("Pipeline state cleared")


def build_cheatsheet(
    book_path: Union[str, Path],
    output_path: Optional[Union[str, Path]] = None,
    config: Optional[Settings] = None,
    profile: Optional[str] = None,
    llm_backend: Optional[str] = None,
    llm_model: Optional[str] = None,
    latex_only: bool = False,
) -> Path:
    """
    Convenience function to build a cheat sheet.
    
    Args:
        book_path: Path to the book/document
        output_path: Output path for PDF (or .tex if latex_only=True)
        config: Optional Settings
        profile: Optional profile name ("technical_book", "math_paper", "nontechnical_book")
        llm_backend: LLM backend to use ("ollama", "huggingface", "openai", "vllm")
        llm_model: Model name for the LLM backend
        latex_only: If True, output LaTeX source only without compiling to PDF
        
    Returns:
        Path to generated PDF or LaTeX file
    """
    # Load settings
    if config is None:
        if profile:
            config = Settings.from_profile(profile)
        else:
            config = get_default_settings()
    
    # Override LLM settings if provided
    if llm_backend:
        config.llm.backend = llm_backend
    if llm_model:
        config.llm.model_name = llm_model
    
    # Create and run pipeline
    pipeline = CheatSheetPipeline(settings=config)
    result = pipeline.build(book_path, output_path, compile_pdf=not latex_only)
    
    if not result.success:
        raise RuntimeError(f"Cheat sheet generation failed: {result.errors}")
    
    return result.output_path
