"""
Math formulas extraction pass.

Extracts mathematical formulas, equations, and their contexts.
"""

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

from ..schema import FormulaItem, MathFormulasSection, 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 MathFormulasPass:
    """
    Pass for extracting mathematical formulas from documents.
    
    Specializes in:
    - Identifying equations and formulas
    - Normalizing to LaTeX
    - Extracting variable definitions
    - Documenting conditions and usage
    """
    
    PASS_NAME = "math_formulas"
    
    # Patterns for detecting math content
    MATH_PATTERNS = [
        r'\$\$.+?\$\$',
        r'\$[^$]+\$',
        r'\\begin\{equation\}.+?\\end\{equation\}',
        r'\\begin\{align\}.+?\\end\{align\}',
        r'[A-Za-z]\s*=\s*[^=]+(?:$|\n)',  # Variable assignments
        r'\b(?:equation|formula|theorem)\s+\d+',
    ]
    
    def __init__(
        self,
        retrieval_service: RetrievalService,
        llm_client: LLMClient,
        max_formulas: int = 25,
        top_k_retrieval: int = 10
    ):
        """
        Initialize math formulas pass.
        
        Args:
            retrieval_service: RAG retrieval service
            llm_client: LLM for generation
            max_formulas: Maximum formulas to extract
            top_k_retrieval: Documents to retrieve per query
        """
        self.retrieval_service = retrieval_service
        self.llm_client = llm_client
        self.max_formulas = max_formulas
        self.top_k = top_k_retrieval
        self._compiled_patterns = [re.compile(p, re.DOTALL | re.IGNORECASE) for p in self.MATH_PATTERNS]
    
    @trace_function(name="math_formulas_pass")
    def run(
        self,
        document_title: str,
        concept_hints: Optional[List[str]] = None,
        config: Optional[GenerationConfig] = None
    ) -> PassResult:
        """
        Run the math formulas extraction pass.
        
        Args:
            document_title: Title of source document
            concept_hints: Related concepts to guide extraction
            config: Generation configuration
            
        Returns:
            PassResult with MathFormulasSection output
        """
        start_time = time.perf_counter()
        config = config or GenerationConfig(max_tokens=2048, temperature=0.2)
        
        errors = []
        all_formulas = []
        total_tokens = 0
        
        queries = self._build_queries(document_title, concept_hints)
        
        try:
            with Span("extract_formulas", {"query_count": len(queries)}):
                # Phase 1: Retrieve math-heavy content
                for query in queries:
                    try:
                        formulas, tokens = self._extract_for_query(query, config)
                        all_formulas.extend(formulas)
                        total_tokens += tokens
                    except Exception as e:
                        logger.warning(f"Formula extraction failed for query '{query}': {e}")
                        errors.append(f"Query '{query[:50]}...': {str(e)}")
                
                # Deduplicate and rank
                unique_formulas = self._deduplicate_formulas(all_formulas)
                ranked_formulas = self._rank_formulas(unique_formulas)
                final_formulas = ranked_formulas[:self.max_formulas]
                
                # Build notation guide
                notation_guide = self._build_notation_guide(final_formulas)
                
                section = MathFormulasSection(
                    formulas=final_formulas,
                    notation_guide=notation_guide
                )
            
            duration_ms = (time.perf_counter() - start_time) * 1000
            
            return PassResult(
                success=True,
                pass_name=self.PASS_NAME,
                output=section,
                errors=errors,
                metrics={
                    "formulas_extracted": len(final_formulas),
                    "queries_processed": len(queries),
                    "avg_variables_per_formula": self._avg_variables(final_formulas)
                },
                duration_ms=duration_ms,
                tokens_used=total_tokens
            )
            
        except Exception as e:
            logger.error(f"Math formulas pass failed: {e}")
            return PassResult(
                success=False,
                pass_name=self.PASS_NAME,
                errors=[str(e)],
                duration_ms=(time.perf_counter() - start_time) * 1000
            )
    
    def _build_queries(
        self,
        title: str,
        concept_hints: Optional[List[str]]
    ) -> List[str]:
        """Build retrieval queries for math content."""
        queries = [
            "equations formulas mathematical expressions",
            "derivation proof theorem",
            f"key equations in {title}",
            "formula definition variables",
        ]
        
        if concept_hints:
            for concept in concept_hints[:5]:
                queries.append(f"formula equation for {concept}")
        
        return queries
    
    def _extract_for_query(
        self,
        query: str,
        config: GenerationConfig
    ) -> Tuple[List[FormulaItem], int]:
        """Extract formulas for a single query."""
        # Use math-specific retrieval
        results = self.retrieval_service.retrieve_math(query, top_k=self.top_k)
        
        if not results:
            # Fallback to general retrieval
            results = self.retrieval_service.retrieve(query, top_k=self.top_k)
        
        if not results:
            return [], 0
        
        context = "\n\n".join(r.content for r in results)
        
        # Check if context has math content
        if not self._has_math_content(context):
            return [], 0
        
        # Build and run prompt
        prompt = get_prompt("formula_extraction")(context=context)
        result = self.llm_client.generate(prompt, config)
        
        # Parse response
        formulas = self._parse_formulas(result.text)
        tokens = result.tokens_used if hasattr(result, 'tokens_used') else 0
        
        return formulas, tokens
    
    def _has_math_content(self, text: str) -> bool:
        """Check if text contains mathematical content."""
        for pattern in self._compiled_patterns:
            if pattern.search(text):
                return True
        return False
    
    def _parse_formulas(self, text: str) -> List[FormulaItem]:
        """Parse LLM response into FormulaItem objects."""
        formulas = []
        current: Dict[str, Any] = {}
        
        for line in text.split("\n"):
            line = line.strip()
            
            if line.startswith("FORMULA:"):
                if current and "latex" in current:
                    formulas.append(self._build_formula(current))
                current = {"name": line[8:].strip()}
            elif line.startswith("LATEX:"):
                current["latex"] = line[6:].strip()
            elif line.startswith("VARIABLES:"):
                current["variables"] = self._parse_variables(line[10:])
            elif line.startswith("CONDITIONS:"):
                current["conditions"] = [c.strip() for c in line[11:].split(",")]
            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(self._build_formula(current))
        
        return formulas
    
    def _parse_variables(self, text: str) -> Dict[str, str]:
        """Parse variable definitions from text."""
        variables = {}
        parts = text.split(",")
        for part in parts:
            if ":" in part:
                var, desc = part.split(":", 1)
                variables[var.strip()] = desc.strip()
        return variables
    
    def _build_formula(self, data: Dict[str, Any]) -> FormulaItem:
        """Build FormulaItem from parsed data."""
        return FormulaItem(
            name=data.get("name", "Unknown Formula"),
            latex=data.get("latex", ""),
            variables=data.get("variables", {}),
            conditions=data.get("conditions", []),
            usage_context=data.get("context", ""),
            depends_on=data.get("depends_on", []),
            importance=ImportanceLevel.MEDIUM
        )
    
    def _deduplicate_formulas(self, formulas: List[FormulaItem]) -> List[FormulaItem]:
        """Remove duplicate formulas."""
        seen_latex = set()
        unique = []
        
        for formula in formulas:
            # Normalize LaTeX for comparison
            normalized = re.sub(r'\s+', '', formula.latex.lower())
            if normalized not in seen_latex and normalized:
                seen_latex.add(normalized)
                unique.append(formula)
        
        return unique
    
    def _rank_formulas(self, formulas: List[FormulaItem]) -> List[FormulaItem]:
        """Rank formulas by importance."""
        def score(f: FormulaItem) -> float:
            s = 0.0
            if f.variables:
                s += 2.0
            if f.usage_context:
                s += 1.5
            if f.conditions:
                s += 1.0
            if "=" in f.latex:
                s += 0.5  # Equations over expressions
            return s
        
        return sorted(formulas, key=score, reverse=True)
    
    def _avg_variables(self, formulas: List[FormulaItem]) -> float:
        """Calculate average variables per formula."""
        if not formulas:
            return 0.0
        return sum(len(f.variables) for f in formulas) / len(formulas)
    
    def _build_notation_guide(self, formulas: List[FormulaItem]) -> str:
        """Build a notation guide from extracted formulas."""
        all_vars: Dict[str, str] = {}
        for f in formulas:
            all_vars.update(f.variables)
        
        if not all_vars:
            return ""
        
        # Build compact notation guide
        guide_parts = []
        for var, desc in sorted(all_vars.items())[:20]:  # Limit
            guide_parts.append(f"${var}$: {desc}")
        
        return "; ".join(guide_parts)
