"""
Algorithms extraction pass.

Extracts algorithms and generates pseudocode from documents.
"""

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

from ..schema import AlgorithmItem, AlgorithmSection, 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 AlgorithmsPass:
    """
    Pass for extracting algorithms and generating pseudocode.
    
    Focuses on:
    - Algorithm identification
    - Pseudocode generation
    - Complexity analysis
    - Key insights
    """
    
    PASS_NAME = "algorithms"
    
    def __init__(
        self,
        retrieval_service: RetrievalService,
        llm_client: LLMClient,
        max_algorithms: int = 20,
        top_k_retrieval: int = 8
    ):
        """
        Initialize algorithms pass.
        
        Args:
            retrieval_service: RAG retrieval service
            llm_client: LLM for generation
            max_algorithms: Maximum algorithms to extract
            top_k_retrieval: Documents to retrieve per query
        """
        self.retrieval_service = retrieval_service
        self.llm_client = llm_client
        self.max_algorithms = max_algorithms
        self.top_k = top_k_retrieval
    
    @trace_function(name="algorithms_pass")
    def run(
        self,
        document_title: str,
        concept_hints: Optional[List[str]] = None,
        config: Optional[GenerationConfig] = None
    ) -> PassResult:
        """
        Run the algorithms extraction pass.
        
        Args:
            document_title: Title of source document
            concept_hints: Related concepts to guide extraction
            config: Generation configuration
            
        Returns:
            PassResult with AlgorithmSection output
        """
        start_time = time.perf_counter()
        config = config or GenerationConfig(max_tokens=3000, temperature=0.3)
        
        errors = []
        all_algorithms = []
        total_tokens = 0
        
        queries = self._build_queries(document_title, concept_hints)
        
        try:
            with Span("extract_algorithms", {"query_count": len(queries)}):
                for query in queries:
                    try:
                        algorithms, tokens = self._extract_for_query(query, config)
                        all_algorithms.extend(algorithms)
                        total_tokens += tokens
                    except Exception as e:
                        logger.warning(f"Algorithm extraction failed for query '{query}': {e}")
                        errors.append(f"Query '{query[:50]}...': {str(e)}")
                
                # Deduplicate and rank
                unique = self._deduplicate(all_algorithms)
                ranked = self._rank_algorithms(unique)
                final = ranked[:self.max_algorithms]
                
                section = AlgorithmSection(algorithms=final)
            
            duration_ms = (time.perf_counter() - start_time) * 1000
            
            return PassResult(
                success=True,
                pass_name=self.PASS_NAME,
                output=section,
                errors=errors,
                metrics={
                    "algorithms_extracted": len(final),
                    "queries_processed": len(queries),
                    "with_pseudocode": sum(1 for a in final if a.pseudocode),
                    "with_complexity": sum(1 for a in final if a.complexity),
                },
                duration_ms=duration_ms,
                tokens_used=total_tokens
            )
            
        except Exception as e:
            logger.error(f"Algorithms 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 algorithms."""
        queries = [
            "algorithm procedure steps method",
            f"algorithms in {title}",
            "pseudocode implementation",
            "training algorithm optimization",
            "procedure method approach",
        ]
        
        if concept_hints:
            for concept in concept_hints[:3]:
                queries.append(f"algorithm for {concept}")
        
        return queries
    
    def _extract_for_query(
        self,
        query: str,
        config: GenerationConfig
    ) -> Tuple[List[AlgorithmItem], int]:
        """Extract algorithms for a single query."""
        # Use code-specific retrieval if available
        results = self.retrieval_service.retrieve_code(query, top_k=self.top_k)
        
        if not results:
            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)
        
        # Build and run prompt
        prompt = get_prompt("algorithm_extraction")(context=context)
        result = self.llm_client.generate(prompt, config)
        
        # Parse response
        algorithms = self._parse_algorithms(result.text)
        tokens = result.tokens_used if hasattr(result, 'tokens_used') else 0
        
        return algorithms, tokens
    
    def _parse_algorithms(self, text: str) -> List[AlgorithmItem]:
        """Parse LLM response into AlgorithmItem objects."""
        algorithms = []
        current: Dict[str, Any] = {}
        in_pseudocode = False
        pseudocode_lines = []
        
        for line in text.split("\n"):
            if line.startswith("ALGORITHM:"):
                if current and "name" in current:
                    if pseudocode_lines:
                        current["pseudocode"] = "\n".join(pseudocode_lines)
                    algorithms.append(self._build_algorithm(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["inputs"] = [i.strip() for i in line[6:].split(",")]
            elif line.startswith("OUTPUT:"):
                current["outputs"] = [o.strip() for o in line[7:].split(",")]
            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 and "name" in current:
            if pseudocode_lines:
                current["pseudocode"] = "\n".join(pseudocode_lines)
            algorithms.append(self._build_algorithm(current))
        
        return algorithms
    
    def _build_algorithm(self, data: Dict[str, Any]) -> AlgorithmItem:
        """Build AlgorithmItem from parsed data."""
        return AlgorithmItem(
            name=data.get("name", "Unknown Algorithm"),
            purpose=data.get("purpose", ""),
            pseudocode=data.get("pseudocode", ""),
            inputs=data.get("inputs", []),
            outputs=data.get("outputs", []),
            complexity=data.get("complexity", ""),
            key_insight=data.get("insight", ""),
            importance=ImportanceLevel.MEDIUM
        )
    
    def _deduplicate(self, algorithms: List[AlgorithmItem]) -> List[AlgorithmItem]:
        """Remove duplicate algorithms by name."""
        seen = set()
        unique = []
        for algo in algorithms:
            key = algo.name.lower().strip()
            if key not in seen and key:
                seen.add(key)
                unique.append(algo)
        return unique
    
    def _rank_algorithms(self, algorithms: List[AlgorithmItem]) -> List[AlgorithmItem]:
        """Rank algorithms by completeness."""
        def score(a: AlgorithmItem) -> float:
            s = 0.0
            if a.pseudocode:
                s += 3.0
            if a.purpose:
                s += 2.0
            if a.complexity:
                s += 1.5
            if a.inputs:
                s += 1.0
            if a.outputs:
                s += 1.0
            if a.key_insight:
                s += 1.0
            return s
        
        return sorted(algorithms, key=score, reverse=True)
