"""
Routing functions for search algorithms.
"""

import os
from typing import Dict, Any, Optional, Tuple, Type
import logging

from skydiscover.search.base_database import ProgramDatabase, Program
from skydiscover.search.default_discovery_manager import DiscoveryManagerInput, DiscoveryManager
from skydiscover.config import DatabaseConfig, Config, load_config, build_output_dir
from skydiscover.search.utils.discovery_utils import load_database_from_file
from skydiscover.utils.code_utils import extract_solution_language

logger = logging.getLogger(__name__)

# Import implementations
from skydiscover.search.best_of_n.database import BestOfNDatabase
from skydiscover.search.beam_search.database import BeamSearchDatabase
from skydiscover.search.topk.database import TopKDatabase
from skydiscover.search.adaevolve.database import AdaEvolveDatabase
from skydiscover.search.adaevolve.manager import AdaEvolveManager
from skydiscover.search.evox.database.search_strategy_db import SearchStrategyDatabase

######################### REGISTRY #########################
_PROGRAM_REGISTRY: Dict[str, Type[Program]] = {}
_DATABASE_REGISTRY: Dict[str, Type[ProgramDatabase]] = {}
_MANAGER_REGISTRY: Dict[str, Type[DiscoveryManager]] = {}

def register_program(search_type: str, program_class: Type[Program]) -> None:
    """Register a program class for a search type."""
    _PROGRAM_REGISTRY[search_type] = program_class
    logger.debug(f"Registered program class '{program_class.__name__}' for search type '{search_type}'")


def register_database(search_type: str, database_class: Type[ProgramDatabase]) -> None:
    """Register a database class for a search type."""
    _DATABASE_REGISTRY[search_type] = database_class
    logger.debug(f"Registered database class '{database_class.__name__}' for search type '{search_type}'")


def register_manager(search_type: str, manager_class: Type[DiscoveryManager]) -> None:
    """Register a discovery manager class for a search type."""
    _MANAGER_REGISTRY[search_type] = manager_class
    logger.debug(f"Registered manager class '{manager_class.__name__}' for search type '{search_type}'")


######################### ROUTING FUNCTIONS #########################

def get_database_class(search_type: str, config: DatabaseConfig) -> ProgramDatabase:
    """
    Get a database instance for a given search type.
    
    Supports both registered search types and dynamic loading for "evox"/"evolve" types
    when a custom database_file_path is specified.
    """
    # Handle dynamic loading for evox/evolve types (evolve = search-side config in co-evolution)
    if search_type in ("evox", "evolve") and getattr(config, "database_file_path", None):
        database_class, program_class = load_database_from_file(config.database_file_path)
        db = database_class(search_type, config)
        db._program_class = program_class
        return db
    
    # Use registry for standard types
    if search_type not in _DATABASE_REGISTRY:
        available_types = ", ".join(sorted(_DATABASE_REGISTRY.keys()))
        raise ValueError(
            f"Unknown search type: '{search_type}'. "
            f"Available types: {available_types}. "
            f"For 'evox'/'evolve' type with custom database, set config.search.database.database_file_path"
        )
    
    database_class = _DATABASE_REGISTRY[search_type]
    return database_class(search_type, config)

def get_discovery_manager(manager_input: DiscoveryManagerInput) -> DiscoveryManager:
    """
    Get the discovery manager for a given search type.
    If search type is "evox", uses CoEvolutionManager.
    Otherwise, uses registered manager or defaults to DiscoveryManager.
    """
    search_type = manager_input.config.search.type

    if search_type == "evox":
        logger.debug("Search type 'evox', using CoEvolutionManager")
        from skydiscover.search.evox.coevolve_manager import CoEvolutionManager
        return CoEvolutionManager(manager_input)

    logger.debug(f"Getting manager for search type: {search_type}")
    
    manager_class = _MANAGER_REGISTRY.get(search_type, DiscoveryManager)
    return manager_class(manager_input)

def get_program(
    config: Config,
    initial_program_solution: str,
    initial_program_id: str,
    initial_metrics: Dict[str, Any],
    start_iteration: int,
) -> Program:
    """
    Create an initial Program instance appropriate to the search type.
    
    Supports both registered program classes and dynamic loading for "evox" type
    when a custom database_file_path is specified.
    """
    search_type = config.search.type
    
    # Handle dynamic loading for evox type
    if search_type == "evox" and getattr(config.search.database, "database_file_path", None):
        logger.info(f"Using search strategy from: {config.search.database.database_file_path}")
        _, program_class = load_database_from_file(config.search.database.database_file_path)
        return program_class(
            id=initial_program_id,
            solution=initial_program_solution,
            language=config.language,
            metrics=initial_metrics,
            iteration_found=start_iteration,
        )
    
    # Use registered program class if available
    program_class = _PROGRAM_REGISTRY.get(search_type, Program)
    
    # Create program instance with appropriate parameters
    # Base Program has iteration_found as a field, so we can always include it
    return program_class(
        id=initial_program_id,
        solution=initial_program_solution,
        language=config.language,
        metrics=initial_metrics,
        iteration_found=start_iteration,
    )


######################### SEARCH SETUP #########################

def setup_search(
    initial_program_path: str,
    evaluation_file: str,
    config_path: str,
    output_dir: Optional[str] = None,
) -> Tuple[DiscoveryManagerInput, str]:
    """
    Load config, create database, and build an DiscoveryManagerInput from a config path.

    This is the lightweight alternative to creating a full MainController when
    you only need the config/database/manager (e.g. for the search side of
    co-evolution).

    Returns:
        Tuple of (manager_input, initial_program_solution)
    """
    config = load_config(config_path)

    with open(initial_program_path, "r") as f:
        initial_program_solution = f.read()

    if not config.language:
        config.language = extract_solution_language(initial_program_solution)

    file_extension = os.path.splitext(initial_program_path)[1] or ".py"
    if not file_extension.startswith("."):
        file_extension = f".{file_extension}"
    if not hasattr(config, "file_suffix") or config.file_suffix == ".py":
        config.file_suffix = file_extension

    database = get_database_class(config.search.type, config.search.database)

    if not output_dir:
        output_dir = build_output_dir(config.search.type, initial_program_path)

    manager_input = DiscoveryManagerInput(
        config=config,
        evaluation_file=evaluation_file,
        database=database,
        file_suffix=config.file_suffix,
        output_dir=output_dir,
    )
    return manager_input, initial_program_solution


######################### AUTO-REGISTRATION #########################
# Base Program is the default, so we don't need to register it
register_database("best_of_n", BestOfNDatabase)
register_database("beam_search", BeamSearchDatabase)
register_database("topk", TopKDatabase)

# AdaEvolve: adaptive multi-island search with UCB selection
register_database("adaevolve", AdaEvolveDatabase)
register_manager("adaevolve", AdaEvolveManager)

# EvoX guide: search strategy database for co-evolution
register_database("evolve_search_strategy", SearchStrategyDatabase)