"""
AdaEvolve Manager - Evolution loop with adaptive search intensity.

A clean implementation that uses the adaptive database for all
exploration/exploitation decisions. No explicit stagnation tracking -
search intensity handles exploration automatically.

Features:
- Adaptive sampling based on accumulated improvement signal
- Mode-aware prompting (exploration vs exploitation)
- Paradigm breakthrough for high-level strategy shifts
- Sibling context for learning from previous attempts
- Comprehensive JSON logging of all AdaEvolve signals
"""

import json
import logging
import os
import time
import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

from skydiscover.search.base_database import Program
from skydiscover.search.default_discovery_manager import DiscoveryManager, DiscoveryManagerInput
from skydiscover.search.utils.discovery_utils import SerializableResult
from skydiscover.search.adaevolve.paradigm import ParadigmGenerator
from skydiscover.llm.llm_pool import LLMPool
from skydiscover.prompt.default import DefaultContextBuilder
from skydiscover.evaluation.evaluator import Evaluator
from skydiscover.utils.code_utils import (
    extract_diffs,
    apply_diff,
    format_diff_summary,
    parse_full_rewrite,
)
from skydiscover.utils.metrics_utils import get_fitness_score

logger = logging.getLogger(__name__)


class AdaEvolveManager(DiscoveryManager):
    """
    AdaEvolve evolution manager with adaptive search intensity.

    Key Features:
    1. Adaptive sampling: Search intensity per island determines exploration/exploitation
    2. Mode-aware prompting: Different guidance for exploration vs exploitation
    3. Sibling context: Shows previous mutations for learning
    4. Error retry: Retries failed generations with error context
    5. Island rotation: UCB-based selection via database.end_iteration()
    6. Paradigm breakthrough: High-level strategy shifts when globally stuck

    No explicit stagnation tracking - search intensity handles exploration
    automatically based on accumulated improvement signal.
    """

    def __init__(self, manager_input: DiscoveryManagerInput):
        super().__init__(manager_input)

        # Configuration
        db_config = self.config.search.database
        self.enable_retry = getattr(db_config, "enable_error_retry", True)
        self.max_retries = getattr(db_config, "max_error_retries", 2)
        self.num_inspirations = self.config.search.num_inspirations

        # Components
        self.llms = LLMPool(self.config.llm.models)
        self.context_builder = DefaultContextBuilder(self.config.prompt)

        eval_prompt = DefaultContextBuilder(self.config.prompt)
        eval_prompt.set_templates("evaluator_system_message")
        self.evaluator = Evaluator(
            self.config.evaluator,
            self.evaluation_file,
            LLMPool(self.config.llm.evaluator_models),
            eval_prompt,
            database=self.database,
            suffix=self.file_suffix,
            is_image_mode=(self.config.language == "image"),
        )

        # Paradigm generator (if paradigm breakthrough is enabled)
        # Note: We check database.use_paradigm_breakthrough at runtime, not this init-time flag
        # This ensures correct behavior after checkpoint load
        if self.database.use_paradigm_breakthrough:
            model_names = ", ".join(m.name for m in self.guide_llms.models_cfg)
            logger.info(f"Paradigm LLM: using guide_models [{model_names}]")

            self.paradigm_generator = ParadigmGenerator(
                llm_pool=self.guide_llms,
                system_message=self.config.prompt.system_message or "",
                evaluator_code=self._load_evaluator_code(),
                num_paradigms=self.database.get_paradigm_num_to_generate(),
                eval_timeout=self.config.evaluator.timeout,
                language=self.config.language or "python",
            )
        else:
            self.paradigm_generator = None

        # JSON logging for comprehensive AdaEvolve stats
        self._iteration_stats_log_path: Optional[str] = None
        self._iteration_stats_file = None
        self._last_sampling_mode: Optional[str] = None
        self._last_sampling_intensity: Optional[float] = None

        logger.info(
            f"AdaEvolveManager initialized "
            f"(language={self.config.language}, "
            f"paradigm_breakthrough={self.database.use_paradigm_breakthrough})"
        )

    def _load_evaluator_code(self) -> str:
        """Load evaluator source code for paradigm generation context."""
        try:
            if self.evaluation_file and Path(self.evaluation_file).exists():
                return Path(self.evaluation_file).read_text()
        except Exception as e:
            logger.warning(f"Could not load evaluator code: {e}")
        return ""

    # =========================================================================
    # JSON Logging for AdaEvolve Stats
    # =========================================================================

    def _setup_iteration_stats_logging(self, output_dir: Optional[str] = None) -> None:
        """
        Set up JSON logging for comprehensive iteration statistics.

        Creates a JSONL file that records all AdaEvolve signals at each iteration.
        This enables detailed post-hoc analysis of the discovery process.

        Args:
            output_dir: Directory to write the log file. If None, uses database.config.db_path
        """
        # Determine output directory
        if output_dir is None:
            output_dir = self.output_dir
        if output_dir is None:
            output_dir = getattr(self.database.config, "db_path", None)
        if output_dir is None:
            output_dir = "."

        os.makedirs(output_dir, exist_ok=True)

        # Create log file with timestamp
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self._iteration_stats_log_path = os.path.join(
            output_dir, f"adaevolve_iteration_stats_{timestamp}.jsonl"
        )

        logger.info(f"AdaEvolve iteration stats will be logged to: {self._iteration_stats_log_path}")

    def _log_iteration_stats(
        self,
        iteration: int,
        sampling_mode: Optional[str] = None,
        sampling_intensity: Optional[float] = None,
        child_program: Optional[Dict] = None,
        iteration_time: Optional[float] = None,
        error: Optional[str] = None,
    ) -> None:
        """
        Log comprehensive iteration statistics to JSON file.

        This method collects all AdaEvolve signals and writes them as a single
        JSON line to the log file for easy post-processing.

        Args:
            iteration: Current iteration number
            sampling_mode: The mode used for sampling (exploration/exploitation/balanced)
            sampling_intensity: The search intensity value used
            child_program: The child program dict if successfully generated
            iteration_time: Time taken for this iteration
            error: Error message if iteration failed
        """
        if self._iteration_stats_log_path is None:
            return

        try:
            # Get comprehensive stats from database
            stats = self.database.get_comprehensive_iteration_stats(
                iteration=iteration,
                sampling_mode=sampling_mode if sampling_mode is not None else self._last_sampling_mode,
                sampling_intensity=sampling_intensity if sampling_intensity is not None else self._last_sampling_intensity,
            )

            # Add timestamp
            stats["timestamp"] = datetime.now().isoformat()

            # Add iteration-specific info
            stats["iteration_result"] = {
                "success": error is None,
                "error": error,
                "iteration_time_seconds": iteration_time,
            }

            # Add child program info if available
            if child_program:
                stats["iteration_result"]["child_program"] = {
                    "id": child_program.get("id"),
                    "metrics": child_program.get("metrics"),
                    "generation": child_program.get("generation"),
                    "parent_id": child_program.get("parent_id"),
                }

            # Write to JSONL file
            with open(self._iteration_stats_log_path, "a") as f:
                f.write(json.dumps(stats, default=str) + "\n")

        except Exception as e:
            logger.warning(f"Failed to log iteration stats: {e}")

    def get_iteration_stats_log_path(self) -> Optional[str]:
        """Get the path to the iteration stats log file."""
        return self._iteration_stats_log_path

    # =========================================================================
    # Main Evolution Loop
    # =========================================================================

    async def run_discovery(
        self,
        start_iteration: int,
        max_iterations: int,
        checkpoint_callback=None,
    ) -> Optional[Program]:
        """Run evolution with adaptive search intensity and island rotation."""
        total = start_iteration + max_iterations
        logger.info(
            f"AdaEvolve: Running {max_iterations} iterations "
            f"across {self.database.num_islands} islands"
        )

        # Set up comprehensive JSON logging for iteration stats
        self._setup_iteration_stats_logging()

        # Ensure all islands are seeded
        self._ensure_all_islands_seeded()

        for iteration in range(start_iteration, total):
            if self.shutdown_event.is_set():
                logger.info("Shutdown requested")
                break

            try:
                await self._run_iteration(iteration, checkpoint_callback)
            except Exception as e:
                logger.exception(f"Iteration {iteration} failed: {e}")
            finally:
                # CRITICAL: Tell database iteration is complete
                # This handles island rotation (UCB) and migration
                self.database.end_iteration(iteration)

        logger.info("AdaEvolve completed")
        self.database.log_status()

        # Log final summary and stats file location
        if self._iteration_stats_log_path:
            logger.info(f"AdaEvolve iteration stats saved to: {self._iteration_stats_log_path}")

        return self.database.get_best_program()

    def _ensure_all_islands_seeded(self) -> None:
        """Ensure all islands have at least one program."""
        # Find a seed program
        seed_program = None
        for i in range(self.database.num_islands):
            size = self.database.get_island_size(i)
            if size > 0 and seed_program is None:
                population = self.database.get_island_population(i)
                if population:
                    seed_program = population[0]
                    break

        if seed_program is None:
            logger.warning("No seed program found")
            return

        # Seed empty islands
        for i in range(self.database.num_islands):
            if self.database.get_island_size(i) == 0:
                copy = Program(
                    id=str(uuid.uuid4()),
                    solution=seed_program.solution,
                    language=seed_program.language,
                    metrics=seed_program.metrics.copy() if seed_program.metrics else {},
                    iteration_found=seed_program.iteration_found,
                    parent_id=None,
                    generation=0,
                    metadata={"seeded_to_island": i},
                )
                self.database.add(copy, iteration=0, target_island=i)
                logger.info(f"Seeded island {i}")

    async def _run_iteration(self, iteration: int, checkpoint_callback) -> None:
        """Execute one evolution iteration."""
        iteration_start_time = time.time()

        # Check for global paradigm stagnation
        # Use database flag directly to stay in sync after checkpoint load
        if self.database.use_paradigm_breakthrough and self.database.is_paradigm_stagnating():
            await self._generate_paradigms_if_needed()

        result = await self._run_normal_step(iteration)

        iteration_time = time.time() - iteration_start_time

        if result.error:
            logger.warning(f"Iteration {iteration}: {result.error}")
            # Log failed iteration stats
            self._log_iteration_stats(
                iteration=iteration,
                sampling_mode=self._last_sampling_mode,
                sampling_intensity=self._last_sampling_intensity,
                child_program=None,
                iteration_time=iteration_time,
                error=result.error,
            )
        else:
            self._process_result(result, iteration, checkpoint_callback)
            # Log successful iteration stats
            self._log_iteration_stats(
                iteration=iteration,
                sampling_mode=self._last_sampling_mode,
                sampling_intensity=self._last_sampling_intensity,
                child_program=result.child_program_dict,
                iteration_time=result.iteration_time,
                error=None,
            )

    async def _generate_paradigms_if_needed(self) -> None:
        """Generate new paradigms if stagnating and none active."""
        if self.paradigm_generator is None:
            return

        if self.database.has_active_paradigm():
            return  # Already have paradigms to use

        logger.info("Global paradigm stagnation detected, generating breakthrough ideas...")

        # Get current best program for context
        best_program = self.database.get_best_program()
        best_solution = best_program.solution if best_program else ""
        best_score = get_fitness_score(
            best_program.metrics if best_program else {}
        )

        # Extract evaluator feedback from the best program's artifacts
        evaluator_feedback = None
        if best_program and best_program.artifacts:
            fb = best_program.artifacts.get("feedback")
            if fb and isinstance(fb, str):
                evaluator_feedback = fb

        # Get previously tried ideas for feedback
        previously_tried = self.database.get_previously_tried_ideas()

        # Generate new paradigms
        paradigms = await self.paradigm_generator.generate(
            current_program_solution=best_solution,
            current_best_score=best_score,
            previously_tried_ideas=previously_tried,
            evaluator_feedback=evaluator_feedback,
        )

        if paradigms:
            self.database.set_paradigms(paradigms)
            logger.info(f"Generated {len(paradigms)} breakthrough paradigms")
        else:
            logger.warning("Failed to generate paradigms")

    async def _run_normal_step(self, iteration: int) -> SerializableResult:
        """Run a normal iteration with optional retry."""
        last_error = None
        attempts = 1 + (self.max_retries if self.enable_retry else 0)

        for attempt in range(attempts):
            result = await self._generate_child(iteration, error_context=last_error)
            if not result.error:
                return result
            last_error = result.error
            logger.debug(f"Attempt {attempt + 1}/{attempts} failed: {last_error}")

        return SerializableResult(
            error=f"All {attempts} attempts failed: {last_error}",
            iteration=iteration,
        )

    def _process_result(
        self,
        result: SerializableResult,
        iteration: int,
        checkpoint_callback,
    ) -> None:
        """Process a successful result by adding to database."""
        child = Program(**result.child_program_dict)

        # Add to database (database handles which island)
        self.database.add(child, iteration=iteration, parent_id=result.parent_id)

        # Fire monitor callback (live dashboard)
        if self.monitor_callback:
            try:
                self.monitor_callback(child, iteration)
            except Exception:
                pass  # Never crash discovery process due to monitor

        # Log prompt
        if result.prompt:
            self.database.log_prompt(
                template_key=(
                    "full_rewrite_user"
                    if not self.config.diff_based_generation
                    else "diff_user"
                ),
                program_id=child.id,
                prompt=result.prompt,
                responses=[result.llm_response] if result.llm_response else [],
            )

        # Log progress
        logger.info(
            f"Iteration {iteration}: Program {child.id[:8]} "
            f"(parent: {result.parent_id[:8] if result.parent_id else 'None'}) "
            f"completed in {result.iteration_time:.2f}s"
        )

        # Log metrics
        if child.metrics:
            metrics_str = ", ".join(
                f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}"
                for k, v in child.metrics.items()
            )
            logger.info(f"Metrics: {metrics_str}")

        # Check for new best
        if self.database.best_program_id == child.id:
            logger.info(f"New best solution found at iteration {iteration}")

        # Checkpoint callback
        if iteration > 0 and iteration % self.config.checkpoint_interval == 0:
            logger.info(f"Checkpoint interval reached at iteration {iteration}")
            self.database.log_status()
            if checkpoint_callback:
                checkpoint_callback(iteration)

    # =========================================================================
    # Child Generation
    # =========================================================================

    async def _generate_child(
        self,
        iteration: int,
        error_context: Optional[str] = None,
        force_exploration: bool = False,
    ) -> SerializableResult:
        """Generate and evaluate a single child program."""
        try:
            # Sample parent and inspiration programs (database returns standard framework dicts)
            parent_dict, inspiration_programs_dict = self.database.sample(
                self.num_inspirations,
                force_exploration=force_exploration,
            )

            # Unpack parent dict (standard framework pattern)
            if not parent_dict:
                logger.error("sample() returned empty parent dict")
                return SerializableResult(error="Empty parent dict from sample()", iteration=iteration)
            parent_label = list(parent_dict.keys())[0]
            parent = list(parent_dict.values())[0]

            # Read sampling mode stashed by database.sample()
            sampling_mode = getattr(self.database, "_last_sampling_mode", None) or "balanced"

            # Capture sampling mode and intensity for logging
            self._last_sampling_mode = sampling_mode
            current_island = self.database.current_island
            if self.database.use_adaptive_search:
                self._last_sampling_intensity = self.database.adapter.get_search_intensity(current_island)
            else:
                self._last_sampling_intensity = self.database.fixed_intensity

            # When paradigm is active, use best program as parent
            # This ensures paradigm (designed from best) is applied to best, not random parent
            paradigm = self.database.get_current_paradigm() if self.database.use_paradigm_breakthrough else None
            if paradigm:
                best_program = self.database.get_best_program()
                if best_program:
                    parent_dict = {parent_label: best_program}
                    parent = best_program
                    # Keep inspiration_programs_dict from sampling for diversity

            # Build context for prompt generation (standard framework keys)
            context = {
                "program_metrics": parent.metrics,
                "inspiration_programs": inspiration_programs_dict,
                "diff_based_generation": self.config.diff_based_generation,
                "language": self.config.language,
                "evaluator_timeout": self.config.evaluator.timeout,
            }
            # Include any extra prompt context (e.g. task_description for image mode)
            for k, v in self._prompt_context.items():
                if k not in context:
                    context[k] = v

            # Build prompt using framework path (label injected by _format_current_program)
            prompt = self.context_builder.build_prompt(parent_dict, context)

            # Augment with AdaEvolve-specific context (paradigm, sibling, error)
            prompt["user"] = self._augment_prompt(
                prompt["user"],
                sampling_mode,
                parent,
                error_context,
            )

            # Build tracking info for child program
            parent_info = (parent_label, parent.id)
            context_info = [
                (label, p.id)
                for label, programs in inspiration_programs_dict.items()
                for p in programs
            ]
            inspiration_program_ids = [
                p.id for programs in inspiration_programs_dict.values() for p in programs
            ]

            # Apply human feedback (append or replace mode)
            if self.feedback_reader:
                self.feedback_reader.set_current_prompt(prompt["system"])
                feedback = self.feedback_reader.read()
                if feedback:
                    prompt = self.feedback_reader.apply_feedback(prompt)
                    self.feedback_reader.log_usage(iteration, feedback, self.feedback_reader.mode)

            # Generate and evaluate
            return await self._execute_generation(
                parent, prompt, iteration,
                parent_info=parent_info,
                context_info=context_info,
                inspiration_program_ids=inspiration_program_ids,
                inspiration_programs=inspiration_programs_dict,
            )

        except Exception as e:
            logger.exception(f"Generation failed: {e}")
            return SerializableResult(error=str(e), iteration=iteration)

    # =========================================================================
    # Prompt Augmentation
    # =========================================================================

    def _augment_prompt(
        self,
        base_prompt: str,
        sampling_mode: str,
        parent: Program,
        error_context: Optional[str],
    ) -> str:
        """Add AdaEvolve-specific sections to the prompt.

        Note: exploration/exploitation mode guidance is now handled by the
        label system — the label text is injected into the prompt by the
        framework's _format_current_program() before the parent code.
        This method adds evaluator feedback, paradigm, sibling, and error context.
        """
        sections = []

        # Evaluator feedback from parent (highest value — shows why parent fails)
        feedback_section = self._format_evaluator_feedback(parent)
        if feedback_section:
            sections.append(feedback_section)

        # Paradigm guidance (high priority - added if active)
        paradigm = self.database.get_current_paradigm() if self.database.use_paradigm_breakthrough else None
        if paradigm:
            sections.append(self._format_paradigm_guidance(paradigm))
            self.database.use_paradigm()

        # Sibling context (if available)
        if hasattr(self.database, "get_children"):
            sibling_section = self._format_sibling_context(parent)
            if sibling_section:
                sections.append(sibling_section)

        # Error context (for retry)
        if error_context:
            sections.append(
                f"## RETRY CONTEXT\n"
                f"Previous attempt failed with error:\n```\n{error_context}\n```\n"
                f"Please fix this issue in your response."
            )

        if not sections:
            return base_prompt

        return base_prompt + "\n\n" + "\n\n".join(sections)

    def _format_evaluator_feedback(self, parent: Program) -> Optional[str]:
        """Format evaluator feedback from parent's artifacts.

        The evaluator may return diagnostic feedback (e.g. analysis of failed
        examples) in artifacts["feedback"]. This is injected into the prompt
        so the LLM can make targeted improvements instead of guessing.
        """
        feedback = parent.artifacts.get("feedback") if parent.artifacts else None
        if not feedback or not isinstance(feedback, str):
            return None

        # Truncate very long feedback to keep prompt focused
        max_len = 2000
        if len(feedback) > max_len:
            feedback = feedback[:max_len] + "\n... (truncated)"

        return (
            "## EVALUATOR FEEDBACK ON CURRENT PROGRAM\n"
            "The evaluator analyzed cases where the current program failed "
            "and produced the following diagnostic feedback. "
            "Use this to make targeted improvements:\n\n"
            f"{feedback}"
        )

    def _format_sibling_context(self, parent: Program) -> Optional[str]:
        """Format sibling context showing previous mutations."""
        try:
            children = self.database.get_children(parent.id)
        except (AttributeError, NotImplementedError):
            return None

        if not children:
            return None

        parent_fitness = get_fitness_score(parent.metrics)

        lines = ["## PREVIOUS ATTEMPTS ON THIS PARENT"]
        improved, regressed, unchanged = 0, 0, 0

        for i, child in enumerate(children, 1):
            child_fitness = get_fitness_score(child.metrics)
            delta = child_fitness - parent_fitness

            if delta > 0.001:
                status, improved = "IMPROVED", improved + 1
            elif delta < -0.001:
                status, regressed = "REGRESSED", regressed + 1
            else:
                status, unchanged = "NO CHANGE", unchanged + 1

            lines.append(
                f"  {i}. {parent_fitness:.4f} -> {child_fitness:.4f} "
                f"({delta:+.4f}) [{status}]"
            )

        lines.insert(
            1,
            f"Summary: {improved} improved, {unchanged} unchanged, {regressed} regressed",
        )
        lines.append("Avoid repeating approaches that didn't work.")

        return "\n".join(lines)

    def _format_paradigm_guidance(self, paradigm: Dict[str, Any]) -> str:
        """Format paradigm breakthrough guidance for LLM."""
        is_prompt_opt = (self.config.language or "").lower() in ("text", "prompt")

        if is_prompt_opt:
            return f"""## BREAKTHROUGH STRATEGY - APPLY THIS

The search has stagnated globally. You MUST apply this breakthrough prompt strategy:

**STRATEGY:** {paradigm.get('idea', 'N/A')}

**HOW TO APPLY:**
{paradigm.get('description', 'N/A')}

**TARGET:** {paradigm.get('what_to_optimize', 'score')}

**CAUTIONS:** {paradigm.get('cautions', 'N/A')}

**APPROACH TYPE:** {paradigm.get('approach_type', 'N/A')}

**CRITICAL:**
- You MUST rewrite the prompt using this strategy
- The strategy must be reflected in the actual prompt structure and content
- Keep the prompt clear and well-structured
- Do not add unnecessary verbosity — every sentence should serve a purpose
- Ensure the prompt still addresses the core task"""

        return f"""## BREAKTHROUGH IDEA - IMPLEMENT THIS

The search has stagnated globally. You MUST implement this breakthrough idea:

**IDEA:** {paradigm.get('idea', 'N/A')}

**HOW TO IMPLEMENT:**
{paradigm.get('description', 'N/A')}

**TARGET METRIC:** {paradigm.get('what_to_optimize', 'score')}

**CAUTIONS:** {paradigm.get('cautions', 'N/A')}

**APPROACH TYPE:** {paradigm.get('approach_type', 'N/A')}

**CRITICAL:**
- You MUST implement the breakthrough idea
- Ensure the paradigm is actually used in your implementation (not just mentioned in comments)
- Correctness is essential - your implementation must be correct and functional
- Verify output format matches evaluator requirements
- Make purposeful changes that implement the idea
- Test your implementation logic carefully"""

    # =========================================================================
    # LLM Generation
    # =========================================================================

    async def _execute_generation(
        self,
        parent: Program,
        prompt: Dict[str, str],
        iteration: int,
        parent_info: Optional[tuple] = None,
        context_info: Optional[List[tuple]] = None,
        inspiration_program_ids: Optional[List[str]] = None,
        inspiration_programs: Optional[Dict] = None,
    ) -> SerializableResult:
        """Execute LLM generation and evaluation."""
        start_time = time.time()

        image_path = None
        child_id = str(uuid.uuid4())

        # Generate
        try:
            if self.config.language == "image":
                # Image mode: VLM generates image natively via Responses API
                from skydiscover.search.utils.discovery_utils import build_image_content
                user_content = build_image_content(
                    prompt["user"], parent, inspiration_programs or {}
                )
                result = await self.llms.generate_image(
                    system_message=prompt["system"],
                    messages=[{"role": "user", "content": user_content}],
                    output_dir=self._get_image_output_dir(),
                    program_id=child_id,
                )
                response = result.text or ""
                image_path = result.image_path
                if not image_path:
                    return SerializableResult(
                        error="VLM did not generate an image", iteration=iteration
                    )
            else:
                response = await self.llms.generate(
                    system_message=prompt["system"],
                    messages=[{"role": "user", "content": prompt["user"]}],
                )
        except Exception as e:
            return SerializableResult(error=f"LLM error: {e}", iteration=iteration)

        if not response and self.config.language != "image":
            return SerializableResult(error="Empty LLM response", iteration=iteration)

        # Parse code from response
        if self.config.language == "image":
            child_solution = response or "(image generated)"
            changes = "Image generation"
        elif self.config.diff_based_generation:
            diffs = extract_diffs(response)
            if diffs:
                child_solution = apply_diff(parent.solution, response)
                changes = format_diff_summary(diffs)
            else:
                # No diffs found, try full rewrite
                child_solution = parse_full_rewrite(response, self.config.language)
                changes = "Full rewrite"
        else:
            child_solution = parse_full_rewrite(response, self.config.language)
            changes = "Full rewrite"

        if not child_solution:
            return SerializableResult(
                error="No valid solution in response", iteration=iteration
            )

        # Evaluate
        try:
            eval_input = image_path if self.config.language == "image" else child_solution
            eval_result = await self.evaluator.evaluate_program(
                eval_input, child_id
            )
        except Exception as e:
            return SerializableResult(error=f"Evaluation error: {e}", iteration=iteration)

        metrics = eval_result.metrics
        artifacts = eval_result.artifacts

        # Extract image_path from evaluator metrics (non-image mode fallback)
        if not image_path:
            image_path = (
                metrics.pop("image_path", None)
                if isinstance(metrics.get("image_path"), str)
                else None
            )

        # Build child program with full tracking info
        child_metadata = {"changes": changes, "parent_metrics": parent.metrics}
        if image_path:
            child_metadata["image_path"] = image_path
        child = Program(
            id=child_id,
            solution=child_solution,
            language=self.config.language,
            metrics=metrics,
            iteration_found=iteration,
            parent_id=parent.id,
            other_context_ids=inspiration_program_ids,
            parent_info=parent_info,
            context_info=context_info,
            generation=parent.generation + 1,
            metadata=child_metadata,
            artifacts=artifacts,
        )

        iteration_time = time.time() - start_time

        return SerializableResult(
            child_program_dict=child.to_dict(),
            parent_id=parent.id,
            other_context_ids=inspiration_program_ids,
            iteration_time=iteration_time,
            prompt=prompt,
            llm_response=response,
            iteration=iteration,
        )
