"""
Summarization orchestrator for multi-pass extraction.

Coordinates all extraction passes and assembles final output.
"""

import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from .schema import (
    CheatSheetSchema,
    CoreConceptsSection,
    MathFormulasSection,
    ModelsTheoremsSection,
    AlgorithmSection,
    PassResult,
)
from .passes import (
    CoreConceptsPass,
    MathFormulasPass,
    ModelsTheoremsPass,
    AlgorithmsPass,
    GlobalOutlinePass,
)
from ..core.retrieval import RetrievalService
from ..core.models import LLMClient, GenerationConfig
from ..infra.logging import get_logger
from ..infra.tracing import trace_function, Span

logger = get_logger(__name__)


@dataclass
class OrchestratorConfig:
    """Configuration for summarization orchestrator."""
    max_concepts: int = 30
    max_formulas: int = 25
    max_models: int = 15
    max_theorems: int = 15
    max_algorithms: int = 20
    page_budget: int = 12
    enable_parallel: bool = False
    
    # Pass-specific configs
    generation_config: Optional[GenerationConfig] = None


@dataclass
class OrchestratorResult:
    """Result from orchestrator execution."""
    success: bool
    schema: Optional[CheatSheetSchema] = None
    pass_results: Dict[str, PassResult] = field(default_factory=dict)
    total_duration_ms: float = 0.0
    total_tokens: int = 0
    errors: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "success": self.success,
            "total_duration_ms": self.total_duration_ms,
            "total_tokens": self.total_tokens,
            "errors": self.errors,
            "pass_results": {k: v.to_dict() for k, v in self.pass_results.items()},
            "content_counts": self.schema.get_content_counts() if self.schema else {},
        }


class SummarizationOrchestrator:
    """
    Orchestrates multi-pass content extraction.
    
    Coordinates:
    1. Global outline pass
    2. Core concepts pass
    3. Math formulas pass
    4. Models/theorems pass
    5. Algorithms pass
    
    Assembles results into a complete CheatSheetSchema.
    """
    
    def __init__(
        self,
        retrieval_service: RetrievalService,
        llm_client: LLMClient,
        config: Optional[OrchestratorConfig] = None
    ):
        """
        Initialize the orchestrator.
        
        Args:
            retrieval_service: RAG retrieval service
            llm_client: LLM for generation
            config: Orchestrator configuration
        """
        self.retrieval_service = retrieval_service
        self.llm_client = llm_client
        self.config = config or OrchestratorConfig()
        
        # Initialize passes
        self.outline_pass = GlobalOutlinePass(
            retrieval_service, llm_client
        )
        self.concepts_pass = CoreConceptsPass(
            retrieval_service, llm_client,
            max_concepts=self.config.max_concepts
        )
        self.formulas_pass = MathFormulasPass(
            retrieval_service, llm_client,
            max_formulas=self.config.max_formulas
        )
        self.models_pass = ModelsTheoremsPass(
            retrieval_service, llm_client,
            max_models=self.config.max_models,
            max_theorems=self.config.max_theorems
        )
        self.algorithms_pass = AlgorithmsPass(
            retrieval_service, llm_client,
            max_algorithms=self.config.max_algorithms
        )
    
    @trace_function(name="orchestrator_run")
    def run(
        self,
        document_title: str,
        source_document: str = "",
        table_of_contents: Optional[List[Dict[str, Any]]] = None
    ) -> OrchestratorResult:
        """
        Run all extraction passes and assemble cheat sheet.
        
        Args:
            document_title: Title for the cheat sheet
            source_document: Source document name/path
            table_of_contents: Pre-extracted TOC if available
            
        Returns:
            OrchestratorResult with complete CheatSheetSchema
        """
        start_time = time.perf_counter()
        pass_results: Dict[str, PassResult] = {}
        errors = []
        total_tokens = 0
        gen_config = self.config.generation_config
        
        logger.info(f"Starting summarization orchestrator for: {document_title}")
        
        try:
            with Span("orchestrator_run", {"title": document_title}):
                # Phase 1: Global outline
                logger.info("Running global outline pass...")
                outline_result = self.outline_pass.run(
                    document_title, table_of_contents, gen_config
                )
                pass_results["global_outline"] = outline_result
                total_tokens += outline_result.tokens_used
                
                # Extract hints for subsequent passes
                chapter_hints = []
                concept_hints = []
                if outline_result.success and outline_result.output:
                    chapter_hints = [c["title"] for c in outline_result.output.get("chapters", [])[:10]]
                    concept_hints = outline_result.output.get("themes", [])[:5]
                
                # Phase 2: Core concepts
                logger.info("Running core concepts pass...")
                concepts_result = self.concepts_pass.run(
                    document_title, chapter_hints, gen_config
                )
                pass_results["core_concepts"] = concepts_result
                total_tokens += concepts_result.tokens_used
                
                # Update hints from concepts
                if concepts_result.success and concepts_result.output:
                    concept_hints.extend([
                        c.name for c in concepts_result.output.concepts[:5]
                    ])
                
                # Phase 3: Math formulas
                logger.info("Running math formulas pass...")
                formulas_result = self.formulas_pass.run(
                    document_title, concept_hints, gen_config
                )
                pass_results["math_formulas"] = formulas_result
                total_tokens += formulas_result.tokens_used
                
                # Phase 4: Models and theorems
                logger.info("Running models/theorems pass...")
                models_result = self.models_pass.run(
                    document_title, concept_hints, gen_config
                )
                pass_results["models_theorems"] = models_result
                total_tokens += models_result.tokens_used
                
                # Phase 5: Algorithms
                logger.info("Running algorithms pass...")
                algorithms_result = self.algorithms_pass.run(
                    document_title, concept_hints, gen_config
                )
                pass_results["algorithms"] = algorithms_result
                total_tokens += algorithms_result.tokens_used
                
                # Assemble final schema
                schema = self._assemble_schema(
                    document_title,
                    source_document,
                    pass_results
                )
                
                # Collect errors
                for name, result in pass_results.items():
                    if not result.success:
                        errors.append(f"Pass '{name}' failed")
                    errors.extend(result.errors)
            
            duration_ms = (time.perf_counter() - start_time) * 1000
            
            logger.info(
                f"Orchestrator completed in {duration_ms:.1f}ms, "
                f"{total_tokens} tokens used"
            )
            
            return OrchestratorResult(
                success=True,
                schema=schema,
                pass_results=pass_results,
                total_duration_ms=duration_ms,
                total_tokens=total_tokens,
                errors=errors
            )
            
        except Exception as e:
            logger.error(f"Orchestrator failed: {e}")
            return OrchestratorResult(
                success=False,
                pass_results=pass_results,
                total_duration_ms=(time.perf_counter() - start_time) * 1000,
                total_tokens=total_tokens,
                errors=[str(e)] + errors
            )
    
    def _assemble_schema(
        self,
        title: str,
        source: str,
        results: Dict[str, PassResult]
    ) -> CheatSheetSchema:
        """Assemble final CheatSheetSchema from pass results."""
        
        # Get overview from outline pass
        overview = ""
        if results.get("global_outline") and results["global_outline"].output:
            overview = results["global_outline"].output.get("overview", "")
        
        # Get concepts section
        core_concepts = CoreConceptsSection()
        if results.get("core_concepts") and results["core_concepts"].output:
            core_concepts = results["core_concepts"].output
        
        # Get formulas section
        math_formulas = MathFormulasSection()
        if results.get("math_formulas") and results["math_formulas"].output:
            math_formulas = results["math_formulas"].output
        
        # Get models/theorems section
        models_theorems = ModelsTheoremsSection()
        if results.get("models_theorems") and results["models_theorems"].output:
            models_theorems = results["models_theorems"].output
        
        # Get algorithms section
        algorithms = AlgorithmSection()
        if results.get("algorithms") and results["algorithms"].output:
            algorithms = results["algorithms"].output
        
        # Build generation metadata
        metadata = {
            "pass_durations": {
                name: r.duration_ms for name, r in results.items()
            },
            "tokens_by_pass": {
                name: r.tokens_used for name, r in results.items()
            },
        }
        
        return CheatSheetSchema(
            title=title,
            source_document=source,
            overview=overview,
            core_concepts=core_concepts,
            math_formulas=math_formulas,
            models_theorems=models_theorems,
            algorithms=algorithms,
            page_budget=self.config.page_budget,
            generation_metadata=metadata
        )
