"""
Models and theorems extraction pass.

Extracts models, theorems, and theoretical frameworks from documents.
"""

import time
from typing import Any, Dict, List, Optional, Tuple

from ..schema import ModelItem, ModelsTheoremsSection, PassResult, ImportanceLevel
from ...core.retrieval import RetrievalService
from ...core.models import LLMClient, GenerationConfig
from ...core.prompts import get_prompt
from ...infra.logging import get_logger
from ...infra.tracing import trace_function, Span

logger = get_logger(__name__)


class ModelsTheoremsPass:
    """
    Pass for extracting models and theorems from documents.
    
    Focuses on:
    - Model architectures and objectives
    - Theorems and proofs
    - Training/application procedures
    - Strengths and limitations
    """
    
    PASS_NAME = "models_theorems"
    
    def __init__(
        self,
        retrieval_service: RetrievalService,
        llm_client: LLMClient,
        max_models: int = 15,
        max_theorems: int = 15,
        top_k_retrieval: int = 8
    ):
        """
        Initialize models and theorems pass.
        
        Args:
            retrieval_service: RAG retrieval service
            llm_client: LLM for generation
            max_models: Maximum models to extract
            max_theorems: Maximum theorems to extract
            top_k_retrieval: Documents to retrieve per query
        """
        self.retrieval_service = retrieval_service
        self.llm_client = llm_client
        self.max_models = max_models
        self.max_theorems = max_theorems
        self.top_k = top_k_retrieval
    
    @trace_function(name="models_theorems_pass")
    def run(
        self,
        document_title: str,
        concept_hints: Optional[List[str]] = None,
        config: Optional[GenerationConfig] = None
    ) -> PassResult:
        """
        Run the models and theorems extraction pass.
        
        Args:
            document_title: Title of source document
            concept_hints: Related concepts to guide extraction
            config: Generation configuration
            
        Returns:
            PassResult with ModelsTheoremsSection output
        """
        start_time = time.perf_counter()
        config = config or GenerationConfig(max_tokens=3000, temperature=0.3)
        
        errors = []
        total_tokens = 0
        
        try:
            with Span("extract_models_theorems"):
                # Extract models
                models, model_tokens, model_errors = self._extract_models(
                    document_title, concept_hints, config
                )
                total_tokens += model_tokens
                errors.extend(model_errors)
                
                # Extract theorems
                theorems, theorem_tokens, theorem_errors = self._extract_theorems(
                    document_title, concept_hints, config
                )
                total_tokens += theorem_tokens
                errors.extend(theorem_errors)
                
                section = ModelsTheoremsSection(
                    models=models[:self.max_models],
                    theorems=theorems[:self.max_theorems]
                )
            
            duration_ms = (time.perf_counter() - start_time) * 1000
            
            return PassResult(
                success=True,
                pass_name=self.PASS_NAME,
                output=section,
                errors=errors,
                metrics={
                    "models_extracted": len(section.models),
                    "theorems_extracted": len(section.theorems),
                },
                duration_ms=duration_ms,
                tokens_used=total_tokens
            )
            
        except Exception as e:
            logger.error(f"Models/theorems pass failed: {e}")
            return PassResult(
                success=False,
                pass_name=self.PASS_NAME,
                errors=[str(e)],
                duration_ms=(time.perf_counter() - start_time) * 1000
            )
    
    def _extract_models(
        self,
        title: str,
        concept_hints: Optional[List[str]],
        config: GenerationConfig
    ) -> Tuple[List[ModelItem], int, List[str]]:
        """Extract models from the document."""
        models = []
        errors = []
        total_tokens = 0
        
        queries = [
            "model architecture training objective loss function",
            f"models in {title}",
            "neural network model design",
            "statistical model assumptions",
        ]
        
        for query in queries:
            try:
                results = self.retrieval_service.retrieve(query, top_k=self.top_k)
                if not results:
                    continue
                
                context = "\n\n".join(r.content for r in results)
                prompt = get_prompt("model_extraction")(
                    context=context,
                    model_name=title
                )
                
                result = self.llm_client.generate(prompt, config)
                parsed = self._parse_models(result.text)
                models.extend(parsed)
                total_tokens += result.tokens_used if hasattr(result, 'tokens_used') else 0
                
            except Exception as e:
                errors.append(f"Model extraction for '{query[:30]}...': {str(e)}")
        
        return self._deduplicate(models), total_tokens, errors
    
    def _extract_theorems(
        self,
        title: str,
        concept_hints: Optional[List[str]],
        config: GenerationConfig
    ) -> Tuple[List[ModelItem], int, List[str]]:
        """Extract theorems from the document."""
        theorems = []
        errors = []
        total_tokens = 0
        
        queries = [
            "theorem proof lemma corollary",
            f"theorems in {title}",
            "mathematical theorem statement proof",
        ]
        
        for query in queries:
            try:
                results = self.retrieval_service.retrieve(query, top_k=self.top_k)
                if not results:
                    continue
                
                context = "\n\n".join(r.content for r in results)
                prompt = self._theorem_prompt(context)
                
                result = self.llm_client.generate(prompt, config)
                parsed = self._parse_theorems(result.text)
                theorems.extend(parsed)
                total_tokens += result.tokens_used if hasattr(result, 'tokens_used') else 0
                
            except Exception as e:
                errors.append(f"Theorem extraction for '{query[:30]}...': {str(e)}")
        
        return self._deduplicate(theorems), total_tokens, errors
    
    def _theorem_prompt(self, context: str) -> str:
        """Build theorem extraction prompt."""
        return f"""Extract theorems, lemmas, and corollaries from this text.

TEXT:
{context}

OUTPUT FORMAT:
THEOREM: [name]
STATEMENT: [formal statement]
PROOF_SKETCH: [brief proof outline]
IMPLICATIONS: [key implications]
---

Extract theorems now:"""
    
    def _parse_models(self, text: str) -> List[ModelItem]:
        """Parse model extraction response."""
        # The response is often a prose summary, convert to ModelItem
        return [ModelItem(
            name="Main Model",
            summary=text.strip()[:1000],
            importance=ImportanceLevel.HIGH
        )] if text.strip() else []
    
    def _parse_theorems(self, text: str) -> List[ModelItem]:
        """Parse theorem extraction response."""
        theorems = []
        current: Dict[str, Any] = {}
        
        for line in text.split("\n"):
            line = line.strip()
            
            if line.startswith("THEOREM:") or line.startswith("LEMMA:") or line.startswith("COROLLARY:"):
                if current and "name" in current:
                    theorems.append(self._build_theorem(current))
                current = {"name": line.split(":", 1)[1].strip()}
            elif line.startswith("STATEMENT:"):
                current["statement"] = line[10:].strip()
            elif line.startswith("PROOF_SKETCH:"):
                current["proof"] = line[13:].strip()
            elif line.startswith("IMPLICATIONS:"):
                current["implications"] = line[13:].strip()
        
        if current and "name" in current:
            theorems.append(self._build_theorem(current))
        
        return theorems
    
    def _build_theorem(self, data: Dict[str, Any]) -> ModelItem:
        """Build ModelItem for a theorem."""
        summary_parts = []
        if data.get("statement"):
            summary_parts.append(f"Statement: {data['statement']}")
        if data.get("proof"):
            summary_parts.append(f"Proof: {data['proof']}")
        if data.get("implications"):
            summary_parts.append(f"Implications: {data['implications']}")
        
        return ModelItem(
            name=data.get("name", "Unknown Theorem"),
            summary="\n".join(summary_parts) if summary_parts else "",
            importance=ImportanceLevel.MEDIUM
        )
    
    def _deduplicate(self, items: List[ModelItem]) -> List[ModelItem]:
        """Remove duplicate items by name."""
        seen = set()
        unique = []
        for item in items:
            key = item.name.lower().strip()
            if key not in seen and key:
                seen.add(key)
                unique.append(item)
        return unique
