"""
Discovery manager for running discovery processes.

Provides the default execution loop for discovery processes (sample → prompt → LLM → evaluate).
Subclasses only need to override ``run_discovery`` to change orchestration
(e.g. co-evolution interleaves solution and search-algorithm evolution).
"""

import logging
import multiprocessing as mp
import os
import time
import uuid
from dataclasses import dataclass
from typing import Any, Callable, Dict, Optional, Union

from skydiscover.config import Config
from skydiscover.evaluation.evaluator import Evaluator
from skydiscover.llm.llm_pool import LLMPool
from skydiscover.prompt.default import DefaultContextBuilder
from skydiscover.prompt.evox import EvoxContextBuilder
from skydiscover.search.base_database import Program, ProgramDatabase
from skydiscover.search.utils.discovery_utils import SerializableResult, build_image_content
from skydiscover.utils.code_utils import apply_diff, extract_diffs, format_diff_summary, parse_full_rewrite

logger = logging.getLogger(__name__)



@dataclass
class DiscoveryManagerInput:
    """Input to the discovery manager"""

    config: Config
    evaluation_file: str
    database: ProgramDatabase
    file_suffix: str = ".py"
    output_dir: Optional[str] = None


class DiscoveryManager:
    """
    Discovery manager with a default sequential execution strategy.

    Handles the full generate-evaluate cycle: prompt building, LLM calls,
    response parsing, evaluation, and result processing.

    The default ``run_discovery`` runs iterations sequentially.  Subclasses
    (e.g. CoevolutionManager) can override it for different orchestration
    while reusing the shared iteration primitives.
    """

    def __init__(self, manager_input: DiscoveryManagerInput):
        self.config = manager_input.config
        self.evaluation_file = manager_input.evaluation_file
        self.database = manager_input.database
        self.file_suffix = manager_input.file_suffix
        self.output_dir = manager_input.output_dir

        self.shutdown_event = mp.Event()
        self.early_stopping_triggered = False

        self.llms = LLMPool(self.config.llm.models)
        self.evaluator_llms = LLMPool(self.config.llm.evaluator_models)
        self.guide_llms = LLMPool(self.config.llm.guide_models)

        self._init_context_builder()

        self.evaluator_context_builder = DefaultContextBuilder(self.config.prompt)
        self.evaluator_context_builder.set_templates("evaluator_system_message")

        self.evaluator = Evaluator(
            self.config.evaluator,
            self.evaluation_file,
            self.evaluator_llms,
            self.evaluator_context_builder,
            database=self.database,
            suffix=self.file_suffix,
            is_image_mode=(self.config.language == "image"),
        )

        self.num_inspirations = manager_input.config.search.num_inspirations

        # Optional callbacks set by MainController
        self.monitor_callback: Optional[Callable] = None  # (program, iteration) -> None
        self.feedback_reader: Optional[Any] = None  # HumanFeedbackReader

        # Context for the context builder
        self._prompt_context: Dict[str, Any] = {}
        
        logger.info(f"DiscoveryManager initialized: num_inspirations={self.num_inspirations}")

    # ------------------------------------------------------------------
    # Initialisation helpers
    # ------------------------------------------------------------------

    def _init_context_builder(self):
        """Initialize the appropriate context builder based on config."""
        if getattr(self.config.prompt, 'template', 'default') == "evox":
            self.context_builder = EvoxContextBuilder(
                self.config.prompt,
                summary_llm_config=self.config.llm.guide_models,
            )
            template_name = "search_evolution_user" if self.config.diff_based_generation else "search_evolution_full_rewrite_user"
            self.context_builder.set_templates(user_template=template_name)
        else:
            self.context_builder = DefaultContextBuilder(self.config.prompt)

    # ------------------------------------------------------------------
    # Discovery loop 
    # ------------------------------------------------------------------

    async def run_discovery(
        self,
        start_iteration: int,
        max_iterations: int,
        checkpoint_callback: Optional[Callable[[int], None]] = None,
        post_process_result: Optional[bool] = True,
        retry_times: Optional[int] = 3,
    ) -> Optional[Union[Program, SerializableResult]]:
        """
        Run the discovery process.

        Default implementation: sequential iteration without parallelism.

        Args:
            start_iteration: The iteration to start from.
            max_iterations: The number of iterations to run.
            checkpoint_callback: Optional callback for checkpointing.
            post_process_result: If True, add results to the database and
                return the best Program.  If False, return the raw
                ``SerializableResult`` from the last iteration.
            retry_times: Number of retry attempts per iteration.

        Returns:
            Best ``Program`` found (post_process_result=True) or raw
            ``SerializableResult`` (post_process_result=False).
        """
        total_iterations = start_iteration + max_iterations

        result = None
        for iteration in range(start_iteration, total_iterations):
            if self.shutdown_event.is_set():
                logger.info("Shutdown requested, stopping discovery loop early")
                break

            try:
                result = await self._run_iteration(iteration, retry_times=retry_times)
                if result.error:
                    logger.warning(f"Iteration {iteration} failed: {result.error}")
                    continue

                if post_process_result:
                    self._process_iteration_result(result, iteration, checkpoint_callback)

            except Exception as e:
                logger.exception(f"Error in iteration {iteration}: {e}")

        if not post_process_result:
            return result

        if self.shutdown_event.is_set():
            logger.info(f"✅ Discovery process completed (search strategy = {self.database.name}) - Shutdown requested")
        else:
            logger.info(f"✅ Discovery process completed (search strategy = {self.database.name}) - Maximum iterations reached")

        return self.database.get_best_program()

    # ------------------------------------------------------------------
    # Single-iteration primitives (shared by all managers)
    # ------------------------------------------------------------------

    async def _run_iteration(
        self,
        iteration: int,
        retry_times: int = 1,
    ) -> SerializableResult:
        """Run a single generate-evaluate iteration."""
        try:
            parent_dict, inspiration_programs_dict = self.database.sample(num_inspirations=self.num_inspirations)

            assert len(parent_dict) == 1, "parent_dict should contain exactly one program"
            parent_label = list(parent_dict.keys())[0]
            parent = list(parent_dict.values())[0]

            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]

            logger.debug(
                f"Iteration {iteration}: parent {parent.id} ({parent_label}), "
                f"inspiration_programs keys: {list(inspiration_programs_dict.keys())}"
            )

            iteration_start = time.time()

            failed_attempts = []
            child_solution, child_id, child_metrics, llm_response, changes_summary = None, None, None, None, None

            image_path = None  # set by image mode or evaluator

            for retry in range(retry_times):
                prompt = self._build_prompt(
                    parent_dict=parent_dict,
                    parent=parent,
                    inspiration_programs=inspiration_programs_dict,
                    failed_attempts=failed_attempts,
                )

                if failed_attempts:
                    logger.info(f"Retry {retry + 1}/{retry_times}: rebuilding prompt with {len(failed_attempts)} failed attempt(s)")

                # 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)

                try:
                    if self.config.language == "image":
                        # Image mode: VLM generates image natively via Responses API
                        child_id = str(uuid.uuid4())
                        user_content = build_image_content(
                            prompt["user"], parent, inspiration_programs_dict
                        )
                        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,
                        )
                        llm_response = result.text or ""
                        image_path = result.image_path
                        if image_path:
                            child_solution = result.text or "(image generated)"
                            changes_summary = "Image generation"
                            parse_error = None
                        else:
                            child_solution = None
                            changes_summary = None
                            parse_error = "VLM did not generate an image"
                    else:
                        llm_response = await self.llms.generate(
                            system_message=prompt["system"],
                            messages=[{"role": "user", "content": prompt["user"]}],
                        )
                except Exception as e:
                    logger.error(f"LLM generation failed: {e}")
                    return SerializableResult(error=f"LLM generation failed: {str(e)}", iteration=iteration, attempts_used=retry + 1)

                if self.config.language != "image":
                    # Text/code mode: parse LLM response
                    if llm_response is None:
                        return SerializableResult(error="LLM returned None response", iteration=iteration, attempts_used=retry + 1)

                child_solution, changes_summary, parse_error = self._parse_llm_response(
                    llm_response, parent.solution, iteration, retry + 1, retry_times
                )

                if child_solution and len(child_solution) > self.config.max_solution_length:
                    logger.warning(
                        "Generated solution exceeds maximum length (iteration=%s, attempt %s/%s): %s > %s",
                        iteration,
                        retry + 1,
                        retry_times,
                        len(child_solution),
                        self.config.max_solution_length,
                    )
                    parse_error = f"Generated solution exceeds maximum length ({len(child_solution)} > {self.config.max_solution_length})"
                    child_solution = None

                if parse_error:
                    failed_attempts.append({
                        "solution": child_solution or "",
                        "llm_response": llm_response,
                        "metrics": {},
                        "metadata": {
                            "error": parse_error,
                            "attempt_number": retry + 1,
                        },
                    })
                    if retry < retry_times - 1:
                        continue
                    logger.error(
                        "All %s retry attempts failed due to parse/validation error: %s",
                        retry_times, parse_error,
                    )
                    return SerializableResult(
                        error=f"{parse_error} (after {retry_times} attempts)",
                        iteration=iteration,
                        prompt=prompt,
                        llm_response=llm_response,
                        attempts_used=retry_times,
                    )

                if self.config.language != "image":
                    child_id = str(uuid.uuid4())

                eval_input = image_path if self.config.language == "image" else child_solution
                child_eval_result = await self.evaluator.evaluate_program(
                    eval_input, child_id
                )
                child_metrics = child_eval_result.metrics
                # Extract image_path from evaluator metrics (non-image mode fallback)
                if not image_path:
                    image_path = child_metrics.pop("image_path", None) if isinstance(child_metrics.get("image_path"), str) else None

                if (
                    child_metrics.get("validity") in (0, -1) or
                    (child_metrics.get("timeout") is True and child_metrics.get("validity") is None) or
                    (child_metrics.get("combined_score") == 0 and child_metrics.get("error") is not None)
                ):
                    error_msg = (
                        (child_metrics.get("error") if isinstance(child_metrics.get("error"), str) else None)
                        or child_metrics.get("error_message")
                        or "Evaluation failed (validity=0)"
                    )

                    logger.warning(
                        "Evaluation failed (attempt %s/%s): validity=%s, error=%s\n"
                        "----- Generated Solution -----\n%s\n"
                        "--------------------------",
                        retry + 1, retry_times, child_metrics.get("validity"), error_msg, child_solution,
                    )

                    failed_attempts.append({
                        "solution": child_solution,
                        "metrics": child_metrics,
                        "metadata": {
                            "changes": changes_summary,
                            "parent_metrics": parent.metrics,
                            "error": error_msg,
                            "attempt_number": retry + 1,
                        },
                    })

                    if retry < retry_times - 1:
                        continue
                    logger.error("All %s retry attempts failed. Final error: %s", retry_times, error_msg)
                    iteration_time = time.time() - iteration_start
                    failed_extra = {"failed_attempts": failed_attempts}
                    if image_path:
                        failed_extra["image_path"] = image_path
                    failed_child_program = self._create_child_program(
                        child_id=child_id,
                        child_solution=child_solution,
                        parent=parent,
                        inspiration_program_ids=inspiration_program_ids,
                        parent_info=parent_info,
                        context_info=context_info,
                        child_metrics=child_metrics or {},
                        iteration=iteration,
                        changes_summary=changes_summary,
                        extra_metadata=failed_extra,
                        artifacts=child_eval_result.artifacts,
                    )
                    return SerializableResult(
                        error=f"Evaluator failed after {retry_times} attempts: {error_msg}",
                        iteration=iteration,
                        child_program_dict=failed_child_program.to_dict(),
                        parent_id=parent.id,
                        other_context_ids=inspiration_program_ids,
                        iteration_time=iteration_time,
                        prompt=prompt,
                        llm_response=llm_response,
                        attempts_used=retry_times,
                    )
                break

            extra_meta = {}
            if image_path:
                extra_meta["image_path"] = image_path
            child_program = self._create_child_program(
                child_id=child_id,
                child_solution=child_solution,
                parent=parent,
                inspiration_program_ids=inspiration_program_ids,
                parent_info=parent_info,
                context_info=context_info,
                child_metrics=child_metrics,
                iteration=iteration,
                changes_summary=changes_summary,
                extra_metadata=extra_meta if extra_meta else None,
                artifacts=child_eval_result.artifacts,
            )
            iteration_time = time.time() - iteration_start

            return SerializableResult(
                child_program_dict=child_program.to_dict(),
                parent_id=parent.id,
                other_context_ids=inspiration_program_ids,
                iteration_time=iteration_time,
                prompt=prompt,
                llm_response=llm_response,
                iteration=iteration,
                attempts_used=retry + 1,
            )
        except Exception as e:
            logger.exception(f"Error in iteration {iteration}")
            return SerializableResult(error=str(e), iteration=iteration, attempts_used=1)

    # ------------------------------------------------------------------
    # Prompt / parsing / program creation helpers
    # ------------------------------------------------------------------

    def _build_prompt(
        self,
        parent_dict: Dict[str, Program],
        parent: Program,
        inspiration_programs: Dict[str, list],
        failed_attempts: list,
    ) -> Dict[str, str]:
        """Build the prompt for LLM generation."""
        db_stats = self._prompt_context.get("db_stats") or self.database.get_statistics()

        context = {
            "program_metrics": parent.metrics,
            "inspiration_programs": inspiration_programs,
            "evaluator_timeout": getattr(self.config.evaluator, 'timeout', None),
            "previous_programs": db_stats.get("previous_programs", []),
            "diff_based_generation": self.config.diff_based_generation,
            "language": self.config.language,
            "db_stats": db_stats,
        }
        for k, v in self._prompt_context.items():
            if k not in context:
                context[k] = v

        if failed_attempts:
            context["errors"] = failed_attempts

        return self.context_builder.build_prompt(current_program=parent_dict, context=context)

    def _parse_llm_response(
        self,
        llm_response: str,
        parent_solution: str,
        iteration: int,
        attempt: int,
        retry_times: int,
    ) -> tuple[Optional[str], Optional[str], Optional[str]]:
        """
        Parse LLM response to extract child solution.

        Returns:
            Tuple of (child_solution, changes_summary, parse_error)
        """
        if self.config.diff_based_generation:
            diff_blocks = extract_diffs(llm_response)
            if not diff_blocks:
                preview = llm_response[:2000] + ("\n... (truncated) ..." if len(llm_response) > 2000 else "")
                logger.warning(
                    "No valid diffs found in LLM response (iteration=%s, attempt %s/%s). "
                    "Expected SEARCH/REPLACE blocks. Preview:\n%s",
                    iteration, attempt, retry_times, preview,
                )
                return None, None, "No valid diffs found in response"

            child_solution = apply_diff(parent_solution, llm_response)
            changes_summary = format_diff_summary(diff_blocks)

            if child_solution == parent_solution:
                logger.warning(
                    "Diff blocks found but none matched parent solution (iteration=%s, attempt %s/%s).",
                    iteration, attempt, retry_times,
                )
                return None, None, "Diff SEARCH blocks did not match parent solution - no changes applied"

            return child_solution, changes_summary, None
        else:
            new_solution = parse_full_rewrite(llm_response, self.config.language)
            if not new_solution:
                logger.warning(
                    "No valid solution found in LLM response (iteration=%s, attempt %s/%s).",
                    iteration, attempt, retry_times,
                )
                return None, None, "No valid solution found in response"
            return new_solution, "Full rewrite", None

    def _create_child_program(
        self,
        child_id: str,
        child_solution: str,
        parent: Program,
        inspiration_program_ids: list,
        parent_info: tuple,
        context_info: list,
        child_metrics: Dict[str, Any],
        iteration: int,
        changes_summary: Optional[str],
        extra_metadata: Optional[Dict[str, Any]] = None,
        artifacts: Optional[Dict[str, Any]] = None,
    ) -> Program:
        """Create a child program with the given attributes."""
        metadata = {
            "changes": changes_summary,
            "parent_metrics": parent.metrics,
        }
        if extra_metadata:
            metadata.update(extra_metadata)

        return Program(
            id=child_id,
            solution=child_solution,
            language=self.config.language,
            parent_id=parent.id,
            other_context_ids=inspiration_program_ids,
            parent_info=parent_info,
            context_info=context_info,
            metrics=child_metrics,
            iteration_found=iteration,
            metadata=metadata,
            artifacts=artifacts or {},
        )

    # ------------------------------------------------------------------
    # Post-processing
    # ------------------------------------------------------------------

    async def postprocess_result(self, result: SerializableResult, iteration_number: int, verbose: bool = True):
        """
        Process the iteration result and return the best program from the database.

        Used by co-evolution where evaluation can be delayed.
        """
        self._process_iteration_result(result, iteration_number, checkpoint_callback=None, verbose=verbose)
        return self.database.get_best_program()

    # ------------------------------------------------------------------
    # Utilities
    # ------------------------------------------------------------------

    def _get_image_output_dir(self) -> str:
        """Return the directory for saving VLM-generated images."""
        base = self.output_dir or "."
        d = os.path.join(base, "generated_images")
        os.makedirs(d, exist_ok=True)
        return d

    def request_shutdown(self) -> None:
        """Request graceful shutdown"""
        logger.info("Graceful shutdown requested...")
        self.shutdown_event.set()

    def _process_iteration_result(
        self,
        result: Any,
        iteration: int,
        checkpoint_callback: Optional[Callable[[int], None]] = None,
        verbose: bool = True,
    ) -> bool:
        """
        Process the result from a single iteration.

        Args:
            result: The iteration result to process.
            iteration: Current iteration number.
            checkpoint_callback: Optional callback for checkpoint intervals.
            verbose: If True, log progress and metrics; if False, suppress logging.

        Returns:
            True if should continue, False if should stop.
        """
        if result.error:
            if verbose:
                logger.warning(f"Iteration {iteration} failed: {result.error}")
            return True

        program_class = getattr(self.database, '_program_class', Program)
        child_program = program_class(**result.child_program_dict)

        self.database.add(child_program, iteration=iteration)

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

        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_program.id,
                prompt=result.prompt,
                responses=[result.llm_response] if result.llm_response else [],
            )

        if verbose:
            logger.info(
                f"Iteration {iteration}: "
                f"Program {child_program.id} "
                f"(parent: {result.parent_id}) "
                f"completed in {result.iteration_time:.2f}s"
            )

        if iteration > 0 and iteration % self.config.checkpoint_interval == 0:
            if verbose:
                logger.info(f"[CHECKPOINT] Checkpoint interval reached at iteration {iteration}")

            self.database.log_status()
            if checkpoint_callback:
                checkpoint_callback(iteration)

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

            if not hasattr(self, "_warned_about_combined_score"):
                self._warned_about_combined_score = False

            if (
                "combined_score" not in child_program.metrics
                and not self._warned_about_combined_score
            ):
                if verbose:
                    logger.warning(
                        "⚠️  No 'combined_score' metric found in evaluation results. "
                        "Using 0.0 for discovery process guidance. "
                        "For better solution discovery results, please modify your evaluator to return a 'combined_score' "
                        "metric that properly weights different aspects of program performance."
                    )
                self._warned_about_combined_score = True

        if self.database.best_program_id == child_program.id and verbose:
            logger.info(
                f"🌟 New best solution found at iteration {iteration}"
                # f"{child_program.id}"
            )

        return True
