"""
Algorithm synthesis chain for extracting and formalizing algorithms.
"""

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

from ..core.retrieval import RetrievalService, RetrievalResult
from ..core.models import LLMClient, GenerationConfig
from ..core.prompts import get_prompt
from ..infra.logging import get_logger
from ..infra.tracing import trace_function

logger = get_logger(__name__)


@dataclass
class Algorithm:
    """Represents an extracted algorithm."""
    name: str
    purpose: str = ""
    input_spec: str = ""
    output_spec: str = ""
    pseudocode: str = ""
    time_complexity: str = ""
    space_complexity: str = ""
    key_insight: str = ""
    edge_cases: List[str] = field(default_factory=list)
    source_sections: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "purpose": self.purpose,
            "input": self.input_spec,
            "output": self.output_spec,
            "pseudocode": self.pseudocode,
            "time_complexity": self.time_complexity,
            "space_complexity": self.space_complexity,
            "key_insight": self.key_insight,
            "edge_cases": self.edge_cases,
            "source_sections": self.source_sections
        }
    
    @property
    def complexity_str(self) -> str:
        """Get formatted complexity string."""
        parts = []
        if self.time_complexity:
            parts.append(f"Time: {self.time_complexity}")
        if self.space_complexity:
            parts.append(f"Space: {self.space_complexity}")
        return ", ".join(parts) if parts else ""


class AlgorithmSynthesizer:
    """
    Chain for extracting and synthesizing algorithms.
    
    Generates:
    - Structured pseudocode
    - Complexity analysis
    - Key insights
    - Edge cases
    """
    
    def __init__(
        self,
        retrieval_service: RetrievalService,
        llm_client: LLMClient,
        include_pseudocode: bool = True,
        include_complexity: bool = True,
        include_examples: bool = True
    ):
        """
        Initialize algorithm synthesizer.
        
        Args:
            retrieval_service: Service for retrieving context
            llm_client: LLM for generation
            include_pseudocode: Generate pseudocode
            include_complexity: Include complexity analysis
            include_examples: Include examples
        """
        self.retrieval_service = retrieval_service
        self.llm_client = llm_client
        self.include_pseudocode = include_pseudocode
        self.include_complexity = include_complexity
        self.include_examples = include_examples
    
    @trace_function(name="extract_algorithms")
    def extract(
        self,
        queries: Optional[List[str]] = None,
        config: Optional[GenerationConfig] = None
    ) -> List[Algorithm]:
        """
        Extract algorithms from the indexed document.
        
        Args:
            queries: Optional specific queries
            config: Generation configuration
            
        Returns:
            List of extracted algorithms
        """
        config = config or GenerationConfig()
        
        # Default queries for algorithm extraction
        if queries is None:
            queries = [
                "algorithm procedure steps",
                "pseudocode implementation",
                "training algorithm optimization",
                "method procedure approach",
            ]
        
        all_algorithms = []
        seen_names = set()
        
        for query in queries:
            # Retrieve code/algorithm content
            results = self.retrieval_service.retrieve_code(query, top_k=5)
            
            if not results:
                results = self.retrieval_service.retrieve(query, top_k=5)
            
            if not results:
                continue
            
            # Extract algorithms from context
            algorithms = self._extract_from_context(results, config)
            
            # Deduplicate
            for algo in algorithms:
                name_key = algo.name.lower().replace(" ", "_")
                if name_key not in seen_names and algo.pseudocode:
                    all_algorithms.append(algo)
                    seen_names.add(name_key)
        
        logger.info(f"Extracted {len(all_algorithms)} algorithms")
        return all_algorithms
    
    def _extract_from_context(
        self,
        results: List[RetrievalResult],
        config: GenerationConfig
    ) -> List[Algorithm]:
        """Extract algorithms from retrieval results."""
        context = "\n\n".join(r.content for r in results)
        
        # Track source sections
        source_sections = list(set(
            section
            for r in results
            for section in r.section_hierarchy
            if section
        ))
        
        prompt = get_prompt("algorithm_extraction")(context=context)
        result = self.llm_client.generate(prompt, config)
        
        # Parse response
        algorithms = self._parse_response(result.text, source_sections)
        
        return algorithms
    
    def _parse_response(
        self,
        text: str,
        source_sections: List[str]
    ) -> List[Algorithm]:
        """Parse algorithm extraction response."""
        algorithms = []
        current: Dict[str, Any] = {}
        in_pseudocode = False
        pseudocode_lines = []
        
        for line in text.split("\n"):
            line_stripped = line.strip()
            
            if line_stripped.startswith("ALGORITHM:"):
                # Save previous algorithm
                if current:
                    if pseudocode_lines:
                        current["pseudocode"] = "\n".join(pseudocode_lines)
                    algorithms.append(self._create_algorithm(current, source_sections))
                
                current = {"name": line_stripped[10:].strip()}
                pseudocode_lines = []
                in_pseudocode = False
                
            elif line_stripped.startswith("PURPOSE:"):
                current["purpose"] = line_stripped[8:].strip()
                
            elif line_stripped.startswith("INPUT:"):
                current["input"] = line_stripped[6:].strip()
                
            elif line_stripped.startswith("OUTPUT:"):
                current["output"] = line_stripped[7:].strip()
                
            elif line_stripped.startswith("PSEUDOCODE:"):
                in_pseudocode = True
                
            elif line_stripped == "```" or line_stripped == "```python":
                in_pseudocode = not in_pseudocode
                
            elif in_pseudocode:
                pseudocode_lines.append(line.rstrip())
                
            elif line_stripped.startswith("COMPLEXITY:"):
                # Parse complexity
                complexity_text = line_stripped[11:].strip()
                current["complexity"] = complexity_text
                in_pseudocode = False
                
            elif line_stripped.startswith("- Time:"):
                current["time_complexity"] = line_stripped[7:].strip()
                
            elif line_stripped.startswith("- Space:"):
                current["space_complexity"] = line_stripped[8:].strip()
                
            elif line_stripped.startswith("KEY INSIGHT:"):
                current["key_insight"] = line_stripped[12:].strip()
                in_pseudocode = False
                
            elif line_stripped.startswith("EDGE CASES:"):
                current["edge_cases"] = line_stripped[11:].strip()
                in_pseudocode = False
        
        # Don't forget the last algorithm
        if current:
            if pseudocode_lines:
                current["pseudocode"] = "\n".join(pseudocode_lines)
            algorithms.append(self._create_algorithm(current, source_sections))
        
        return algorithms
    
    def _create_algorithm(
        self,
        data: Dict[str, Any],
        source_sections: List[str]
    ) -> Algorithm:
        """Create an Algorithm from parsed data."""
        # Parse edge cases if string
        edge_cases = data.get("edge_cases", [])
        if isinstance(edge_cases, str):
            edge_cases = [e.strip() for e in edge_cases.split(",") if e.strip()]
        
        # Parse complexity if combined
        time_complexity = data.get("time_complexity", "")
        space_complexity = data.get("space_complexity", "")
        
        if "complexity" in data and not time_complexity:
            complexity = data["complexity"]
            if "Time:" in complexity:
                time_complexity = complexity.split("Time:")[1].split(",")[0].strip()
            if "Space:" in complexity:
                space_complexity = complexity.split("Space:")[1].strip()
        
        return Algorithm(
            name=data.get("name", ""),
            purpose=data.get("purpose", ""),
            input_spec=data.get("input", ""),
            output_spec=data.get("output", ""),
            pseudocode=data.get("pseudocode", ""),
            time_complexity=time_complexity,
            space_complexity=space_complexity,
            key_insight=data.get("key_insight", ""),
            edge_cases=edge_cases,
            source_sections=source_sections
        )
    
    @trace_function(name="synthesize_algorithm")
    def synthesize(
        self,
        description: str,
        algorithm_name: str,
        config: Optional[GenerationConfig] = None
    ) -> Algorithm:
        """
        Synthesize pseudocode from a description.
        
        Args:
            description: Natural language description
            algorithm_name: Name for the algorithm
            config: Generation configuration
            
        Returns:
            Synthesized Algorithm
        """
        config = config or GenerationConfig()
        
        prompt = get_prompt("algorithm_synthesis")(
            description=description,
            algorithm_name=algorithm_name
        )
        
        result = self.llm_client.generate(prompt, config)
        
        # Parse the synthesized algorithm
        algorithms = self._parse_response(result.text, [])
        
        if algorithms:
            return algorithms[0]
        
        # Fallback: create basic algorithm
        return Algorithm(
            name=algorithm_name,
            purpose=description,
            pseudocode=result.text
        )
    
    def format_for_latex(self, algorithm: Algorithm) -> str:
        """
        Format algorithm for LaTeX document.
        
        Args:
            algorithm: Algorithm to format
            
        Returns:
            LaTeX formatted string
        """
        lines = []
        
        # Use algorithm environment
        lines.append("\\begin{algorithm}[H]")
        lines.append(f"\\caption{{{algorithm.name}}}")
        lines.append("\\begin{algorithmic}[1]")
        
        # Add input/output
        if algorithm.input_spec:
            lines.append(f"\\Require {algorithm.input_spec}")
        if algorithm.output_spec:
            lines.append(f"\\Ensure {algorithm.output_spec}")
        
        # Add pseudocode
        if algorithm.pseudocode:
            # Convert pseudocode to algorithmic format
            for line in algorithm.pseudocode.split("\n"):
                line = line.strip()
                if not line:
                    continue
                
                # Handle common pseudocode patterns
                if line.lower().startswith("for "):
                    lines.append(f"\\For{{{line[4:]}}}")
                elif line.lower().startswith("while "):
                    lines.append(f"\\While{{{line[6:]}}}")
                elif line.lower().startswith("if "):
                    lines.append(f"\\If{{{line[3:]}}}")
                elif line.lower().startswith("else if "):
                    lines.append(f"\\ElsIf{{{line[8:]}}}")
                elif line.lower() == "else":
                    lines.append("\\Else")
                elif line.lower().startswith("return "):
                    lines.append(f"\\Return {line[7:]}")
                elif line.lower() in ["end for", "endfor"]:
                    lines.append("\\EndFor")
                elif line.lower() in ["end while", "endwhile"]:
                    lines.append("\\EndWhile")
                elif line.lower() in ["end if", "endif"]:
                    lines.append("\\EndIf")
                else:
                    lines.append(f"\\State {line}")
        
        lines.append("\\end{algorithmic}")
        lines.append("\\end{algorithm}")
        
        # Add complexity note
        if algorithm.complexity_str:
            lines.append(f"\\textbf{{Complexity:}} {algorithm.complexity_str}")
        
        return "\n".join(lines)
    
    def format_pseudocode_simple(self, algorithm: Algorithm) -> str:
        """
        Format algorithm as simple pseudocode block.
        
        Args:
            algorithm: Algorithm to format
            
        Returns:
            Formatted pseudocode string
        """
        lines = [
            f"Algorithm: {algorithm.name}",
            f"Input: {algorithm.input_spec}" if algorithm.input_spec else None,
            f"Output: {algorithm.output_spec}" if algorithm.output_spec else None,
            "",
            algorithm.pseudocode,
        ]
        
        if algorithm.complexity_str:
            lines.extend(["", f"Complexity: {algorithm.complexity_str}"])
        
        return "\n".join(line for line in lines if line is not None)
