"""
Default context builder for SkyDiscover.

Builds LLM prompts for solution discovery by formatting:
- Program metrics (all metrics from program_metrics dict)
- Inspiration programs (from sample() dict with search-specific keys)
- Feature dimensions (for diversity guidance)

Supports:
- Diff-based and full-rewrite generation modes
- Template overrides (system/user message customization)
- Stochastic template variations
- Backward compatibility with old context format
"""

import logging
import random
from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING

from skydiscover.config import PromptConfig
from skydiscover.prompt.base import ContextBuilder
from skydiscover.prompt.templates import TemplateManager

if TYPE_CHECKING:
    from skydiscover.search.base_database import Program

logger = logging.getLogger(__name__)

def _filter_other_metrics(metrics: dict) -> dict:
    """Filter metrics for display: exclude combined_score/error."""
    return {k: v for k, v in metrics.items() if k not in {"combined_score", "error"}}


class DefaultContextBuilder(ContextBuilder):
    """
    Default context builder for discovery process.
    
    Takes:
        - current_program: str - current solution program to evolve
        - context: Dict containing:
            * program_metrics: Dict[str, float] - All performance metrics
            * inspiration_programs: Dict[str, List[Program]] - Programs from sample() with search-specific keys
            * diff_based_generation: bool - Generation mode (defaults to True)
            * language: str - Programming language (defaults to "python")
            * template_key: str - Template override (optional)
    
    Does:
        - Formats all metrics from program_metrics dict
        - Formats inspiration programs generically (iterates over dict keys)
        - Selects template (diff_user or full_rewrite_user)
        - Applies stochastic template variations if enabled
        - Handles backward compatibility with old context format
    
    Returns:
        Dict[str, str] with keys:
            - "system": System message for LLM
            - "user": User message with formatted prompt
    """

    def __init__(self, config: PromptConfig):
        """
        Initialize the default context builder.

        Args:
            config: The prompt configuration object
        """
        self.config = config
        self.template_manager = TemplateManager(custom_template_dir=config.template_dir)

        # Store custom template mappings
        self.system_template_override = None
        self.user_template_override = None

        # Only log once to reduce duplication
        if not hasattr(logger, "_context_builder_logged"):
            logger.info("Initialized default context builder")
            logger._context_builder_logged = True

    def set_templates(
        self, system_template: Optional[str] = None, user_template: Optional[str] = None
    ) -> None:
        """
        Set custom templates to use for this generator

        Args:
            system_template: Template name for system message
            user_template: Template name for user message
        """
        self.system_template_override = system_template
        self.user_template_override = user_template
        logger.info(f"Set custom templates: system={system_template}, user={user_template}")

    def build_prompt(
        self,
        current_program: Union["Program", Dict[str, "Program"]],
        context: Dict[str, Any] = None,
        **kwargs: Any,
    ) -> Dict[str, str]:
        """
        Build prompt using the new interface (current_program + context).

        Args:
            current_program: Dict[str, Program] with label as key (e.g., {"Current Program": program})
            context: Additional context dictionary containing:
                - 'program_metrics': Dictionary of metric names to values (required)
                - 'inspiration_programs': Dict[str, List[Program]] with labels as keys
                - 'previous_programs': List[Program] of previous attempts
                - 'template_key': Optional override for template key
                - 'diff_based_generation': Whether to use diff-based generation (defaults to True)
                - 'language': Programming language (defaults to "python")
            **kwargs: Additional parameters passed to template formatting

        Returns:
            Dictionary with 'system' and 'user' keys containing the prompt messages
        """
        context = context or {}

        # Extract context
        program_metrics = context.get("program_metrics", {})
        inspiration_programs = context.get("inspiration_programs", {})
        previous_programs = context.get("previous_programs", [])
        language = context.get("language", "python")
        diff_based_generation = context.get("diff_based_generation", True)
        template_key = context.get("template_key")
        task_description = context.get("task_description", "")

        # Format sections
        metrics_str = self._format_metrics(program_metrics)
        
        # Format previous attempts
        previous_attempts_section = self._format_previous_attempts(previous_programs)
        
        # Format inspiration programs (now expects Dict[str, List[Program]])
        inspiration_section = self._format_inspiration_programs(inspiration_programs, language)
        
        # Format current program (now expects Dict[str, Program])
        current_program_section = self._format_current_program(current_program, language)
        
        # Extract actual program for solution access
        if isinstance(current_program, dict):
            actual_program = list(current_program.values())[0]
            current_solution = actual_program.solution if hasattr(actual_program, 'solution') else ""
        else:
            current_solution = current_program.solution if hasattr(current_program, 'solution') else current_program.get("solution", "")
        improvement_areas = self._identify_improvement_areas(
            current_solution, program_metrics, previous_programs
        )
        
        # Format failed attempts from retries (if any)
        if context.get("errors"):
            inspiration_section += "\n## ❌ Previous Failed Attempts (this retry):\n"
            inspiration_section += "The following attempts failed. Avoid these errors:\n\n"
            for attempt in context["errors"]:
                err_msg = attempt.get("metadata", {}).get("error", "Unknown error")
                attempt_num = attempt.get("metadata", {}).get("attempt_number", "?")
                inspiration_section += f"### Attempt {attempt_num}:\n"
                inspiration_section += f"**Error:** {err_msg}\n"
                
                # Show the failed solution/response so LLM can see what went wrong
                failed_solution = attempt.get("solution", "")
                llm_response = attempt.get("llm_response", "")
                
                # For parse errors (SEARCH didn't match), show the LLM response with the bad SEARCH blocks
                if "SEARCH" in err_msg and llm_response:
                    # Truncate if too long
                    if len(llm_response) > 1500:
                        llm_response = llm_response[:1500] + "\n... (truncated)"
                    inspiration_section += f"**Your response that failed:**\n```\n{llm_response}\n```\n\n"
                # For evaluation errors, show the generated solution and traceback
                elif failed_solution:
                    if len(failed_solution) > 1500:
                        failed_solution = failed_solution[:1500] + "\n... (truncated)"
                    inspiration_section += f"**Generated solution that failed:**\n```{language}\n{failed_solution}\n```\n"
                    
                    # Show traceback if available (helps LLM understand the exact error)
                    traceback_str = attempt.get("metadata", {}).get("traceback", "")
                    if traceback_str:
                        if len(traceback_str) > 800:
                            traceback_str = traceback_str[-800:]  # Show last 800 chars (most relevant)
                            traceback_str = "... (truncated)\n" + traceback_str
                        inspiration_section += f"**Traceback:**\n```\n{traceback_str}\n```\n\n"
                    else:
                        inspiration_section += "\n"

        # Build timeout warning if evaluator_timeout is provided
        evaluator_timeout = context.get("evaluator_timeout")
        if evaluator_timeout:
            timeout_warning = f"- Time limit: Programs should complete execution within {evaluator_timeout} seconds; otherwise, they will timeout."
        else:
            timeout_warning = ""

        # Select template
        if template_key:
            user_template_key = template_key
        elif self.user_template_override:
            user_template_key = self.user_template_override
        elif language == "image":
            user_template_key = "image_user"
        else:
            if diff_based_generation:
                user_template_key = "diff_user"
            elif language.lower() in ("text", "prompt"):
                user_template_key = "full_rewrite_user_prompt_opt"
            else:
                user_template_key = "full_rewrite_user"

        user_template = self.template_manager.get_template(user_template_key)

        # Get system message
        system_message = self._get_system_message()

        # Format prompt
        user_message = user_template.format(
            current_program=current_program_section,
            metrics=metrics_str,
            previous_attempts=previous_attempts_section,
            inspiration_programs=inspiration_section,
            improvement_areas=improvement_areas,
            language=language,
            timeout_warning=timeout_warning,
            task_description=task_description,
            **kwargs,
        )

        return {
            "system": system_message,
            "user": user_message
        }

    @staticmethod
    def _format_artifact_feedback(program) -> str:
        """Extract and format evaluator artifacts for the prompt.

        Returns a prompt section string, or empty string if no artifacts.
        """
        artifacts = program.artifacts if hasattr(program, 'artifacts') else (program.get("artifacts") if isinstance(program, dict) else None)
        if not artifacts:
            return ""
        max_len = 2000
        sections = []
        for key, value in artifacts.items():
            if value is None:
                continue
            text = str(value)
            if len(text) > max_len:
                text = text[:max_len] + "\n... (truncated)"
            if key == "feedback":
                sections.append(
                    "## Evaluator Feedback\n"
                    "The evaluator produced the following diagnostic feedback. "
                    "Use this to make targeted improvements:\n\n"
                    f"{text}"
                )
            else:
                sections.append(f"## {key}\n{text}")
        if not sections:
            return ""
        return "\n" + "\n\n".join(sections) + "\n"

    def _format_current_program(self, current_program: Union["Program", Dict[str, "Program"]], language: str) -> str:
        """Format current_program with label and solution (e.g., code, prompt).
        
        Args:
            current_program: Either a Program or Dict[str, Program] with label as key
            language: Programming language for the solution block
        """
        if not current_program:
            return ""

        # Handle labeled dict format (key = "" or DIVERGE_LABEL or REFINE_LABEL)
        if isinstance(current_program, dict):
            label = list(current_program.keys())[0]
            program = list(current_program.values())[0]
        else:
            label = ""
            program = current_program
        
        solution = program.solution if hasattr(program, 'solution') else program.get("solution", "")
        metrics = program.metrics if hasattr(program, 'metrics') else program.get("metrics", {})

        lines = []
        # lines.append(f"\n# Current Program\n")
        # Only inject DIVERGE/REFINE instruction when search returned a non-empty label
        if label:
            lines.append(f"\n{label}\n")

        lines.append("\n## Program Information\n")
        # Show metrics if available (with SP-specific filtering)
        if metrics:
            combined = metrics.get("combined_score")
            if combined is not None and isinstance(combined, (int, float)):
                lines.append(f"combined_score: {combined:.4f}\n")

            error = metrics.get("error")
            if error:
                lines.append(f"error: {error}\n")
            
            other_metrics = _filter_other_metrics(metrics)
            if other_metrics:
                lines.append("Score breakdown:")
                for key, value in other_metrics.items():
                    if isinstance(value, float):
                        lines.append(f"\n  - {key}: {value:.4f}")
                    elif isinstance(value, (int, str, bool)):
                        lines.append(f"\n  - {key}: {value}")
                lines.append("\n")
        
        # In image mode, skip the code block — the VLM sees actual images
        # via multimodal content, not recycled text reasoning.
        if language != "image":
            lines.append(f"\n```{language}\n{solution}\n```\n")

        # Append evaluator feedback from artifacts (if available)
        feedback_section = self._format_artifact_feedback(program)
        if feedback_section:
            lines.append(feedback_section)

        return "".join(lines)

    def _identify_improvement_areas(
        self,
        current_program: str,
        metrics: Dict[str, float],
        previous_programs: List["Program"],
    ) -> str:
        """Identify improvement areas based on combined_score and solution length"""
        improvement_areas = []

        # Get current combined_score
        current_score = metrics.get("combined_score")
        if current_score is None:
            current_score = 0.0
        elif not isinstance(current_score, (int, float)):
            try:
                current_score = float(current_score)
            except (ValueError, TypeError):
                current_score = 0.0

        # Track combined_score changes (not individual metrics)
        if previous_programs:
            # Get the most recent previous program's metrics
            prev_program = previous_programs[-1]
            if hasattr(prev_program, 'metrics'):
                prev_metrics = prev_program.metrics or {}
            else:
                prev_metrics = prev_program.get("metrics", {})
            
            prev_score = prev_metrics.get("combined_score")
            if prev_score is None:
                prev_score = 0.0
            elif not isinstance(prev_score, (int, float)):
                try:
                    prev_score = float(prev_score)
                except (ValueError, TypeError):
                    prev_score = 0.0

            if current_score > prev_score:
                improvement_areas.append(
                    f"Combined score improved: {prev_score:.4f} → {current_score:.4f}"
                )
            elif current_score < prev_score:
                improvement_areas.append(
                    f"Combined score declined: {prev_score:.4f} → {current_score:.4f}. Consider revising recent changes."
                )
            elif abs(current_score - prev_score) < 1e-6:
                improvement_areas.append(
                    f"Combined score unchanged at {current_score:.4f}"
                )

        # Solution length check (configurable threshold)
        threshold = self.config.suggest_simplification_after_chars
        if threshold and len(current_program) > threshold:
            improvement_areas.append(
                f"Consider simplifying - solution length exceeds {threshold} characters"
            )

        # Default guidance if nothing specific
        if not improvement_areas:
            improvement_areas.append(
                "Focus on improving the combined_score"
            )

        return "\n".join(f"- {area}" for area in improvement_areas)

    def _format_single_inspiration_program(
        self, program: "Program", index: int, language: str, lines: list
    ) -> None:
        """Format a single inspiration program and append to lines."""
        if program is None:
            return

        solution = program.solution if hasattr(program, 'solution') else program.get("solution", "")
        metrics = program.metrics if hasattr(program, 'metrics') else program.get("metrics", {})

        combined = metrics.get("combined_score") if metrics else None
        if combined is not None and isinstance(combined, (int, float)):
            lines.append(f"### Program {index} (combined_score: {combined:.4f})\n")
        else:
            lines.append(f"### Program {index}\n")

        if metrics:
            error = metrics.get("error")
            if error:
                lines.append(f"- error: {error}\n")

            other_metrics = _filter_other_metrics(metrics)
            if other_metrics:
                lines.append("Score breakdown:")
                for key, value in other_metrics.items():
                    if isinstance(value, float):
                        lines.append(f"  - {key}: {value:.4f}")
                    elif isinstance(value, (int, str, bool)):
                        lines.append(f"  - {key}: {value}")
                lines.append("\n")
            if language != "image":
                lines.append(f"\n```{language}\n{solution}\n```\n")
        lines.append("\n")

    def _format_inspiration_programs(
        self,
        inspiration_programs: Union[List["Program"], Dict[str, List["Program"]]],
        language: str,
    ) -> str:
        """Format inspiration programs with labels as section headers."""
        if not inspiration_programs:
            return ""

        lines = []

        if isinstance(inspiration_programs, dict):
            for label, programs in inspiration_programs.items():
                if not label:
                    label = "Other Inspirational Programs"
                if not programs:
                    continue
                lines.append(f"\n## {label}\n")
                lines.append("These programs represent diverse approaches and creative solutions that may inspire new ideas:\n\n")
                for i, program in enumerate(programs, start=1):
                    self._format_single_inspiration_program(program, i, language, lines)
        else:
            lines.append("These programs represent diverse approaches and creative solutions that may inspire new ideas:\n")
            for i, program in enumerate(inspiration_programs, start=1):
                self._format_single_inspiration_program(program, i, language, lines)

        return "".join(lines)

    def _format_previous_attempts(
        self,
        previous_programs: List["Program"],
        num_previous_attempts: int = 3,
    ) -> str:
        """Format previous attempts section"""
        if not previous_programs:
            return "No previous attempts yet."
        
        # Get template
        try:
            previous_attempt_template = self.template_manager.get_template("previous_attempt")
        except (ValueError, KeyError):
            # Fallback if template not found
            previous_attempt_template = "### Attempt {attempt_number}\n- Changes: {changes}\n- Metrics: {performance}\n- Outcome: {outcome}"
        
        lines = []
        
        # Select most recent attempts (most recent first)
        # Sort the previous programs by combined_score, from highest to lowest
        # Select the top attempts (considering the length of previous_programs) in mind

        previous_programs.sort(key=lambda p: p.metrics.get("combined_score", 0.0), reverse=True)
        selected_previous = previous_programs[:min(num_previous_attempts, len(previous_programs))]
        
        for i, program in enumerate(reversed(selected_previous)):
            attempt_number = len(selected_previous) - i
            
            # Extract program info (handle both Program objects and dicts)
            if hasattr(program, 'metadata'):
                metadata = program.metadata or {}
            else:
                metadata = program.get("metadata", {})
            
            if hasattr(program, 'metrics'):
                metrics = program.metrics or {}
            else:
                metrics = program.get("metrics", {})
            
            changes = metadata.get("changes", "Unknown changes")
            
            # Format performance metrics
            performance_parts = []
            for name, value in metrics.items():
                if isinstance(value, (int, float)):
                    try:
                        performance_parts.append(f"{name}: {value:.4f}")
                    except (ValueError, TypeError):
                        performance_parts.append(f"{name}: {value}")
                else:
                    performance_parts.append(f"{name}: {value}")
            performance_str = ", ".join(performance_parts) if performance_parts else "No metrics"
            
            # Determine outcome based on comparison with parent
            parent_metrics = metadata.get("parent_metrics", {})
            outcome = self._determine_outcome(metrics, parent_metrics)
            
            lines.append(
                previous_attempt_template.format(
                    attempt_number=attempt_number,
                    changes=changes,
                    performance=performance_str,
                    outcome=outcome,
                )
                + "\n\n"
            )
        
        return "".join(lines)

    def _determine_outcome(
        self,
        program_metrics: Dict[str, Any],
        parent_metrics: Dict[str, Any]
    ) -> str:
        """Determine the outcome of an attempt based on combined_score comparison."""
        prog_value = program_metrics.get("combined_score")
        parent_value = parent_metrics.get("combined_score", 0)

        if isinstance(prog_value, (int, float)) and isinstance(parent_value, (int, float)):
            if prog_value > parent_value:
                return "Improvement in combined_score"
            elif prog_value < parent_value:
                return "Regression in combined_score"

        return "No change in combined_score"

    def _format_metrics(self, metrics: Dict[str, float]) -> str:
        """Format metrics - combined_score first, then breakdown of others."""
        if not metrics:
            return "No metrics available"
        
        lines = []
        
        # Main score first
        combined_score = metrics.get("combined_score")
        if combined_score is not None:
            lines.append(f"- combined_score: {combined_score:.4f}")
        
        # Error if present
        error = metrics.get("error")
        if error:
            lines.append(f"- error: {error}")
        
        # All other metrics (score breakdown)
        other_metrics = _filter_other_metrics(metrics)
        
        if other_metrics:
            lines.append("")
            lines.append("Metrics:")
            for key, value in other_metrics.items():
                if isinstance(value, float):
                    lines.append(f"  - {key}: {value:.4f}")
                elif isinstance(value, (int, str, bool)):
                    lines.append(f"  - {key}: {value}")
        
        return "\n".join(lines) if lines else "No metrics available"

    def _get_system_message(self) -> str:
        """Get system message"""
        if self.system_template_override:
            return self.template_manager.get_template(self.system_template_override)
        system_msg = self.config.system_message
        if system_msg in self.template_manager.templates:
            return self.template_manager.get_template(system_msg)
        return system_msg
