import logging
import os
import signal
import time
import uuid
from typing import Optional

from skydiscover.config import Config, load_config, build_output_dir
from skydiscover.search.base_database import Program
from skydiscover.search.default_discovery_manager import DiscoveryManagerInput, DiscoveryManager
from skydiscover.search.route import get_database_class, get_discovery_manager, get_program
from skydiscover.search.utils.logging_utils import setup_search_logging
from skydiscover.utils.code_utils import extract_solution_language
from skydiscover.utils.format_utils import format_metrics_safe, format_improvement_safe
from skydiscover.utils.metrics_utils import get_fitness_score

logger = logging.getLogger(__name__)


class MainController:
    """
    Class for main controllers.

    Based on the search algorithm class, it will route to the right `Program` and `ProgramDatabase`, 
    and the corresponding manager for each of the task. 

    This interface captures the essential operations for orchestrating search:
    - key function: run() - runs the search process
    - others:save_checkpoint / load_checkpoint: Save and load the database
    - Other utility methods for saving and loading the database and best program
    """

    def __init__(
        self,
        initial_program_path: str,
        evaluation_file: str,
        config_path: Optional[str] = None,
        config: Optional[Config] = None,
        output_dir: Optional[str] = None,
    ):
        # Load configuration
        if config is not None:
            # Use provided Config object directly
            self.config = config
        else:
            # Load from file or use defaults
            self.config = load_config(config_path)

        # Set up the search type and name of the search
        self.name = self.config.search.type

        # Set up output directory
        self.output_dir = output_dir or build_output_dir(self.name, initial_program_path)
        os.makedirs(self.output_dir, exist_ok=True)
        self._setup_logging()

        # Set random seed for reproducibility if specified
        if self.config.random_seed is not None:
            import random
            import hashlib

            # Set global random seeds
            random.seed(self.config.random_seed)
            try:
                import numpy as np
                np.random.seed(self.config.random_seed)
            except ImportError:
                pass

            # Create hash-based seeds for different components
            base_seed = str(self.config.random_seed).encode("utf-8")
            llm_seed = int(hashlib.md5(base_seed + b"llm").hexdigest()[:8], 16) % (2**31)

            # Propagate seed to LLM configurations
            self.config.llm.random_seed = llm_seed
            for model_cfg in self.config.llm.models:
                if not hasattr(model_cfg, "random_seed") or model_cfg.random_seed is None:
                    model_cfg.random_seed = llm_seed
            for model_cfg in self.config.llm.evaluator_models:
                if not hasattr(model_cfg, "random_seed") or model_cfg.random_seed is None:
                    model_cfg.random_seed = llm_seed

            logger.info(f"Set random seed to {self.config.random_seed} for reproducibility")
            logger.debug(f"Generated LLM seed: {llm_seed}")

        # Load initial program
        self.initial_program_path = initial_program_path
        self.initial_program_solution = self._load_initial_program()
        if not self.config.language:
            self.config.language = extract_solution_language(self.initial_program_solution)

        # Extract file extension from initial program
        self.file_extension = os.path.splitext(initial_program_path)[1]
        if not self.file_extension:
            # Default to .py if no extension found
            self.file_extension = ".py"
        else:
            # Make sure it starts with a dot
            if not self.file_extension.startswith("."):
                self.file_extension = f".{self.file_extension}"

        # Set the file_suffix in config (can be overridden in YAML)
        if not hasattr(self.config, "file_suffix") or self.config.file_suffix == ".py":
            self.config.file_suffix = self.file_extension

        # Pass random seed to database if specified
        if self.config.random_seed is not None:
            self.config.search.database.random_seed = self.config.random_seed

        # Initialize the database based on the search type
        self.database = get_database_class(self.config.search.type, self.config.search.database)
        # Propagate language to the database so it can adapt prompts/labels
        self.database.language = self.config.language or "python"
        self.evaluation_file = evaluation_file

        self.discovery_manager: Optional[DiscoveryManager] = None

        logger.info(f"Initialized discovery process with search strategy = {self.name} and initial program = {self.initial_program_path}")

    async def run(
        self,
        iterations: Optional[int] = None,
        checkpoint_path: Optional[str] = None,
    ) -> Optional[Program]:
        """
        Run the discovery process.

        Args:
            iterations: Maximum number of iterations (uses config if None)
            checkpoint_path: Path to resume from checkpoint

        Returns:
            Best program found
        """
        max_iterations = iterations if iterations is not None else self.config.max_iterations

        start_iteration = 0
        if checkpoint_path and os.path.exists(checkpoint_path):
            self._load_checkpoint(checkpoint_path)
            start_iteration = self.database.last_iteration + 1
            logger.info(f"Resuming from checkpoint at iteration {start_iteration}")
        else:
            start_iteration = self.database.last_iteration

        # Create discovery manager (owns LLMs, evaluator, context builders)
        manager_input = DiscoveryManagerInput(
            config=self.config,
            evaluation_file=self.evaluation_file,
            database=self.database,
            file_suffix=self.config.file_suffix,
            output_dir=self.output_dir,
        )
        self.discovery_manager = get_discovery_manager(manager_input)

        # For image mode, pass the full task description to every discovery process prompt
        if self.config.language == "image":
            self.discovery_manager._prompt_context["task_description"] = self.initial_program_solution

        # Add initial program if starting fresh (not resuming from checkpoint)
        should_add_initial = (
            start_iteration == 0
            and len(self.database.programs) == 0
        )

        if should_add_initial:
            logger.info("Adding initial program to database")
            initial_program_id = str(uuid.uuid4())

            initial_image_path = None
            initial_code = self.initial_program_solution

            if self.config.language == "image":
                # Image mode: generate the initial image from seed text using VLM
                logger.info("Image mode: generating initial image from seed text...")
                img_dir = os.path.join(self.output_dir, "generated_images")
                try:
                    result = await self.discovery_manager.llms.generate_image(
                        system_message=(
                            "Generate an image based on the following description. "
                            "Also provide brief reasoning about your creative choices."
                        ),
                        messages=[{"role": "user", "content": self.initial_program_solution}],
                        output_dir=img_dir,
                        program_id=initial_program_id,
                    )
                    initial_image_path = result.image_path
                    initial_code = result.text or self.initial_program_solution
                    logger.info(f"Initial image generated: {initial_image_path}")
                except Exception as e:
                    logger.warning(f"Failed to generate initial image: {e}")

            eval_input = initial_image_path if self.config.language == "image" else self.initial_program_solution
            initial_eval_result = await self.discovery_manager.evaluator.evaluate_program(
                eval_input, initial_program_id
            )
            initial_metrics = initial_eval_result.metrics



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

            initial_program = get_program(
                self.config,
                self.initial_program_solution,
                initial_program_id,
                initial_metrics,
                start_iteration,
            )
            initial_program.artifacts = initial_eval_result.artifacts

            # Store image_path in metadata so it's available for VLM prompts
            if initial_image_path:
                if not initial_program.metadata:
                    initial_program.metadata = {}
                initial_program.metadata["image_path"] = initial_image_path

            self.database.add(initial_program)

            try:
                self.database.initial_program_id = initial_program.id
                self.database.initial_program_score = get_fitness_score(
                    initial_program.metrics or {}
                )
            except Exception as e:
                logger.warning(f"Failed to set initial program score: {e}")
        else:
            logger.info(
                f"Skipping initial program addition (resuming from iteration {start_iteration} "
                f"with {len(self.database.programs)} existing programs)"
            )

        # Run discovery process
        early_stopped = False
        monitor_server = None
        try:
            # ── Start live monitor if enabled ────────────────────────
            if self.config.monitor.enabled:
                try:
                    from skydiscover.extras.monitor import MonitorServer, create_monitor_callback

                    monitor_server = MonitorServer(
                        host=self.config.monitor.host,
                        port=self.config.monitor.port,
                        max_solution_length=self.config.monitor.max_solution_length,
                    )
                    monitor_server.set_config_summary(
                        f"{self.name} | max_iter={max_iterations} | "
                        f"seed={self.config.random_seed}"
                    )
                    monitor_server.start()

                    monitor_callback = create_monitor_callback(
                        monitor_server, self.database, time.time()
                    )
                    self.discovery_manager.monitor_callback = monitor_callback

                    url = f"http://localhost:{monitor_server.port}/"
                    print(f"\n{'─' * 50}\n  Live monitor: {url}\n{'─' * 50}\n", flush=True)
                    logger.info(f"Live monitor: {url}")
                except Exception as e:
                    logger.warning(f"Failed to start monitor: {e}")

            # ── Set up human feedback reader if enabled or monitor is live ─
            if self.config.human_feedback_enabled or monitor_server is not None:
                try:
                    from skydiscover.prompt import HumanFeedbackReader

                    feedback_path = self.config.human_feedback_file or os.path.join(
                        self.output_dir, "human_feedback.md"
                    )
                    feedback_mode = getattr(self.config, "human_feedback_mode", "append")
                    feedback_reader = HumanFeedbackReader(feedback_path, mode=feedback_mode)
                    self.discovery_manager.feedback_reader = feedback_reader

                    if monitor_server:
                        monitor_server.set_feedback_reader(feedback_reader)

                    logger.info(f"Human feedback enabled. Feedback file: {feedback_path}")
                except Exception as e:
                    logger.warning(f"Failed to set up human feedback: {e}")

            # ── Configure AI summary on monitor ──────────────────────
            if monitor_server and self.config.monitor.summary_model:
                try:
                    monitor_server.configure_summary(
                        model=self.config.monitor.summary_model,
                        api_key=self.config.monitor.summary_api_key or "",
                        api_base=self.config.monitor.summary_api_base,
                        top_k=self.config.monitor.summary_top_k,
                        interval=self.config.monitor.summary_interval,
                    )
                except Exception as e:
                    logger.warning(f"Failed to configure AI summary: {e}")

            # ── Push existing programs to monitor ────────────────────
            if self.discovery_manager.monitor_callback and self.database.programs:
                for prog in self.database.programs.values():
                    try:
                        self.discovery_manager.monitor_callback(
                            prog,
                            getattr(prog, "iteration_found", 0),
                        )
                    except Exception:
                        pass
                logger.info(
                    f"Pushed {len(self.database.programs)} existing program(s) to monitor"
                )

            # Set up signal handlers for graceful shutdown
            def signal_handler(signum, frame):
                logger.info(f"Received signal {signum}, initiating graceful shutdown...")
                self.discovery_manager.request_shutdown()

                def force_exit_handler(signum, frame):
                    logger.info("Force exit requested - terminating immediately")
                    import sys
                    sys.exit(0)

                signal.signal(signal.SIGINT, force_exit_handler)

            signal.signal(signal.SIGINT, signal_handler)
            signal.signal(signal.SIGTERM, signal_handler)

            discovery_start = start_iteration
            discovery_iterations = max_iterations
            if should_add_initial and start_iteration == 0:
                discovery_start = 1

            self.database.log_status()

            def _checkpoint_callback(iteration: int) -> None:
                try:
                    em_db = getattr(self.discovery_manager, "database", None)
                    if em_db is not None and em_db is not self.database:
                        self.database = em_db
                except Exception:
                    pass
                self._save_checkpoint(iteration)

            await self.discovery_manager.run_discovery(
                discovery_start,
                discovery_iterations,
                checkpoint_callback=_checkpoint_callback,
            )

            # Ensure controller points at the manager's final database (evox can swap it)
            em_db = getattr(self.discovery_manager, "database", None)
            if em_db is not None and em_db is not self.database:
                self.database = em_db

            # Always save a final checkpoint
            final_iteration = discovery_start + discovery_iterations - 1
            if final_iteration > 0:
                self._save_checkpoint(final_iteration)

        finally:
            # Stop monitor server
            if monitor_server:
                try:
                    monitor_server.push_event({
                        "type": "discovery_complete",
                        "reason": (
                            "early_stopping"
                            if self.discovery_manager and self.discovery_manager.early_stopping_triggered
                            else "completed"
                        ),
                    })
                    time.sleep(0.5)
                    monitor_server.stop()
                except Exception:
                    pass

            early_stopped = (
                self.discovery_manager is not None
                and self.discovery_manager.early_stopping_triggered
            )
            self.discovery_manager = None

        # Get the best program
        best_program = None
        if self.database.best_program_id:
            best_program = self.database.get(self.database.best_program_id)
            logger.info(f"Using tracked best program: {self.database.best_program_id}")

        if best_program is None:
            best_program = self.database.get_best_program()
            logger.info("Using calculated best program (tracked program not found)")

        if best_program:
            if early_stopped:
                logger.info(
                    f"Discovery process complete via early stopping. Best program has metrics: "
                    f"{format_metrics_safe(best_program.metrics)}"
                )
            else:
                logger.info(
                    f"Discovery process complete. Best program has metrics: "
                    f"{format_metrics_safe(best_program.metrics)}"
                )
            self._save_best_program(best_program)
            return best_program
        else:
            logger.warning("No valid programs found during the solution discovery process")
            return None

    ############################### IGNORE: UTILITY METHODS FOR CHECKPOINTING ###############################
    def _setup_logging(self) -> None:
        """Set up logging"""
        log_dir = self.config.log_dir or os.path.join(self.output_dir, "logs")
        setup_search_logging(
            log_level=self.config.log_level,
            log_dir=log_dir,
            name=self.name,
        )

    def _load_initial_program(self) -> str:
        """Load the initial program from file"""
        with open(self.initial_program_path, "r") as f:
            return f.read()

    def _save_checkpoint(self, iteration: int) -> None:
        """
        Save a checkpoint

        Args:
            iteration: Current iteration number
        """
        checkpoint_dir = os.path.join(self.output_dir, "checkpoints")
        os.makedirs(checkpoint_dir, exist_ok=True)

        # Create specific checkpoint directory
        checkpoint_path = os.path.join(checkpoint_dir, f"checkpoint_{iteration}")
        os.makedirs(checkpoint_path, exist_ok=True)

        # Save the database
        self.database.save(checkpoint_path, iteration)

        # Save the best program found so far
        best_program = None
        if self.database.best_program_id:
            best_program = self.database.get(self.database.best_program_id)
        else:
            best_program = self.database.get_best_program()

        if best_program:
            # Save the best program at this checkpoint
            best_program_path = os.path.join(checkpoint_path, f"best_program{self.file_extension}")
            with open(best_program_path, "w") as f:
                f.write(best_program.solution)

            # Save metrics
            best_program_info_path = os.path.join(checkpoint_path, "best_program_info.json")
            with open(best_program_info_path, "w") as f:
                import json

                json.dump(
                    {
                        "id": best_program.id,
                        "generation": best_program.generation,
                        "iteration": best_program.iteration_found,
                        "current_iteration": iteration,
                        "metrics": best_program.metrics,
                        "language": best_program.language,
                        "timestamp": best_program.timestamp,
                        "saved_at": time.time(),
                    },
                    f,
                    indent=2,
                )

            logger.info(
                f"[CHECKPOINT] Saved best program at checkpoint {iteration} with metrics: "
                f"{format_metrics_safe(best_program.metrics)}"
            )

        logger.info(f"Saved checkpoint at iteration {iteration} to {checkpoint_path}")

    def _load_checkpoint(self, checkpoint_path: str) -> None:
        """Load state from a checkpoint directory."""
        if not os.path.exists(checkpoint_path):
            raise FileNotFoundError(f"Checkpoint directory {checkpoint_path} not found")
        logger.info(f"Loading checkpoint from {checkpoint_path}")
        self.database.load(checkpoint_path)
        logger.info(f"Checkpoint loaded successfully (iteration {self.database.last_iteration})")

    def _save_best_program(self, program: Optional[Program] = None) -> None:
        """
        Save the best program

        Args:
            program: Best program (if None, uses the tracked best program)
        """
        # If no program is provided, use the tracked best program from the database
        if program is None:
            if self.database.best_program_id:
                program = self.database.get(self.database.best_program_id)
            else:
                # Fallback to calculating best program if no tracked best program
                program = self.database.get_best_program()

        if not program:
            logger.warning("No best program found to save")
            return

        best_dir = os.path.join(self.output_dir, "best")
        os.makedirs(best_dir, exist_ok=True)

        # Use the extension from the initial program file
        filename = f"best_program{self.file_extension}"
        code_path = os.path.join(best_dir, filename)

        with open(code_path, "w") as f:
            f.write(program.solution)

        # Save complete program info including metrics
        info_path = os.path.join(best_dir, "best_program_info.json")
        with open(info_path, "w") as f:
            import json

            json.dump(
                {
                    "id": program.id,
                    "generation": program.generation,
                    "iteration": program.iteration_found,
                    "timestamp": program.timestamp,
                    "parent_id": program.parent_id,
                    "metrics": program.metrics,
                    "language": program.language,
                    "saved_at": time.time(),
                },
                f,
                indent=2,
            )

        # Image mode: also copy the best image file
        if self.config.language == "image" and program.metadata:
            img = program.metadata.get("image_path")
            if img and os.path.exists(img):
                import shutil
                best_img = os.path.join(best_dir, "best_image" + os.path.splitext(img)[1])
                shutil.copy2(img, best_img)
                logger.info(f"Saved best image to {best_img}")

        logger.info(f"Saved best program to {code_path} with program info to {info_path}")

    def _log_iteration(
        self,
        iteration: int,
        parent: Program,
        child: Program,
        elapsed_time: float,
    ) -> None:
        """
        Log iteration progress

        Args:
            iteration: Iteration number
            parent: Parent program
            child: Child program
            elapsed_time: Elapsed time in seconds
        """
        # Calculate improvement using safe formatting
        improvement_str = format_improvement_safe(parent.metrics, child.metrics)

        logger.info(
            f"Iteration {iteration+1}: Child {child.id} from parent {parent.id} "
            f"in {elapsed_time:.2f}s. Metrics: "
            f"{format_metrics_safe(child.metrics)} "
            f"(Δ: {improvement_str})"
        )
