"""
High-level API for using SkyDiscover as a library
"""

import asyncio
import tempfile
import os
import uuid
from typing import Union, Callable, Optional, List, Dict, Any
from dataclasses import dataclass
from pathlib import Path

from skydiscover.search.main_controller import MainController
from skydiscover.search.base_database import Program
from skydiscover.config import Config, SearchConfig, load_config


@dataclass
class DiscoveryResult:
    """Result of a discovery process run"""

    best_program: Optional[Program]
    best_score: float
    best_solution: str
    metrics: Dict[str, Any]
    output_dir: Optional[str]
    initial_score: Optional[float] = None

    def __repr__(self):
        initial_score = self.initial_score if self.initial_score is not None else "N/A"
        return f"DiscoveryResult(best_score={self.best_score:.4f}, initial_score={initial_score})"


def run_discovery(
    initial_program: Union[str, Path, List[str]],
    evaluator: Union[str, Path, Callable],
    config: Union[str, Path, Config, None] = None,
    iterations: Optional[int] = None,
    output_dir: Optional[str] = None,
    cleanup: bool = True,
    codebase: Optional[Union[str, Path]] = None,
    model: Optional[str] = None,
    search: Optional[str] = None,
    system_prompt: Optional[str] = None,
    api_base: Optional[str] = None,
) -> DiscoveryResult:
    """
    Run discovery with flexible inputs - the main library API

    Args:
        initial_program: Can be:
            - Path to a program file (str or Path)
            - Solution (e.g., code, prompt, etc.)
        evaluator: Can be:
            - Path to an evaluator file (str or Path)
            - Callable function that takes (program_path) and returns metrics dict
        config: Can be:
            - Path to config YAML file (str or Path)
            - Config object
            - None for defaults
        iterations: Number of iterations (overrides config)
        output_dir: Output directory (None for temp directory)
        cleanup: If True, clean up temp files after discovery process
        codebase: Path to codebase directory. When provided, enables agentic
            mode so the LLM can read files and search the codebase before
            generating a solution.
        model: LLM model name (e.g. "gpt-5"). Overrides model in config.
        search: Search algorithm to use. Built-in: "topk", "adaevolve",
            "evox", "beam_search", "best_of_n". External: "openevolve",
            "shinkaevolve", "gepa" (requires the corresponding package
            to be installed).
        system_prompt: Domain-specific context for the LLM. Describes the
            problem, constraints, and what the solution should optimise for.
            Mapped to the appropriate field for each backend (e.g.
            prompt.system_message for built-in, task_sys_msg for ShinkaEvolve,
            background for GEPA). Overrides config.prompt.system_message.
        api_base: Base URL for the LLM API. Use this for non-OpenAI providers
            that expose an OpenAI-compatible endpoint.

    Returns:
        DiscoveryResult with best program and metrics
    """
    return asyncio.run(
        _run_discovery_async(
            initial_program, evaluator, config, iterations, output_dir, cleanup,
            codebase=codebase, model=model, search=search,
            system_prompt=system_prompt, api_base=api_base,
        )
    )


async def _run_discovery_async(
    initial_program: Union[str, Path, List[str]],
    evaluator: Union[str, Path, Callable],
    config: Union[str, Path, Config, None],
    iterations: Optional[int],
    output_dir: Optional[str],
    cleanup: bool,
    codebase: Optional[Union[str, Path]] = None,
    model: Optional[str] = None,
    search: Optional[str] = None,
    system_prompt: Optional[str] = None,
    api_base: Optional[str] = None,
) -> DiscoveryResult:
    """Async implementation of run_discovery"""

    temp_dir = None
    temp_files = []

    try:
        # Handle configuration
        if config is None:
            config_obj = load_config(None)
        elif isinstance(config, Config):
            config_obj = config
        else:
            config_obj = load_config(str(config))

        # Apply api_base before model rebuild so models inherit the correct URL.
        if api_base:
            config_obj.llm.api_base = api_base

        # Apply parameter overrides — replace the entire pool with a single model.
        if model:
            from skydiscover.config import LLMModelConfig
            config_obj.llm.models = [LLMModelConfig(name=model)]
            config_obj.llm.evaluator_models = config_obj.llm.models.copy()
            config_obj.llm.guide_models = config_obj.llm.models.copy()
        elif api_base:
            # No model change, but api_base changed — push to existing models.
            config_obj.llm.update_model_params({"api_base": api_base}, overwrite=True)

        # Re-resolve API key if api_base changed the provider after load_config
        if api_base:
            from skydiscover.config import _resolve_api_key_from_env
            resolved_key = _resolve_api_key_from_env(config_obj.llm.api_base)
            if resolved_key:
                config_obj.llm.api_key = resolved_key
                config_obj.llm.update_model_params({"api_key": resolved_key}, overwrite=True)

        # Ensure model overrides inherit shared LLM request settings.
        # Without this, overridden models can end up with retries/timeout=None.
        if model or api_base:
            config_obj.llm.update_model_params(
                {
                    "api_base": config_obj.llm.api_base,
                    "api_key": config_obj.llm.api_key,
                    "temperature": config_obj.llm.temperature,
                    "top_p": config_obj.llm.top_p,
                    "max_tokens": config_obj.llm.max_tokens,
                    "timeout": config_obj.llm.timeout,
                    "retries": config_obj.llm.retries,
                    "retry_delay": config_obj.llm.retry_delay,
                    "random_seed": config_obj.llm.random_seed,
                    "reasoning_effort": config_obj.llm.reasoning_effort,
                },
                overwrite=True,
            )

        if codebase:
            codebase_path = os.path.abspath(str(codebase))
            if not os.path.isdir(codebase_path):
                raise ValueError(f"Codebase directory not found: {codebase_path}")
            config_obj.agentic.enabled = True
            config_obj.agentic.codebase_root = codebase_path

        if search:
            if not hasattr(config_obj, "search"):
                config_obj.search = SearchConfig()
            config_obj.search.type = search
            # Re-create database config so the correct subclass is used
            from skydiscover.config import _DB_CONFIG_BY_TYPE
            new_db_cls = _DB_CONFIG_BY_TYPE.get(search)
            if new_db_cls and not isinstance(config_obj.search.database, new_db_cls):
                config_obj.search.database = new_db_cls()

        # Apply system_prompt — stored on config so backends can read it
        if system_prompt:
            config_obj.prompt.system_message = system_prompt
            config_obj._system_prompt = system_prompt

        # Process initial program and evaluator
        program_path = _prepare_program(initial_program, temp_dir, temp_files)
        evaluator_path = _prepare_evaluator(evaluator, temp_dir, temp_files)

        # Set up output directory (after program_path is ready for problem name derivation)
        search_type = getattr(config_obj.search, "type", None) if hasattr(config_obj, "search") else None
        if output_dir is None and cleanup:
            temp_dir = tempfile.mkdtemp(prefix="skydiscover_")
            actual_output_dir = temp_dir
        else:
            from skydiscover.config import build_output_dir
            actual_output_dir = output_dir or build_output_dir(
                search_type or "default", program_path
            )
            os.makedirs(actual_output_dir, exist_ok=True)
        if search_type:
            from skydiscover.extras.external import is_external, get_runner, KNOWN_EXTERNAL
            if is_external(search_type):
                try:
                    result = await get_runner(search_type)(
                        program_path=program_path,
                        evaluator_path=evaluator_path,
                        config_obj=config_obj,
                        iterations=iterations or config_obj.max_iterations,
                        output_dir=actual_output_dir,
                    )
                except ModuleNotFoundError as e:
                    from skydiscover.extras.external import get_package_name
                    pkg = get_package_name(search_type)
                    raise ImportError(
                        f"{e}\n\nThe '{search_type}' backend requires its package to be installed. "
                        f"Install it with: pip install {pkg}"
                    ) from e
                result.output_dir = actual_output_dir if not cleanup else None
                return result
            if search_type in KNOWN_EXTERNAL:
                raise ImportError(
                    f"Search type '{search_type}' requires the '{search_type}' package. "
                    f"Install it with: pip install {search_type}"
                )

        # Validate that LLM models are configured (only for built-in algorithms)
        if not config_obj.llm.models:
            raise ValueError(
                "No LLM models configured. Please provide a config with LLM models, or set up "
                "your configuration with models. For example:\n\n"
                "from skydiscover.config import Config, LLMModelConfig\n"
                "config = Config()\n"
                "config.llm.models = [LLMModelConfig(name='gpt-5', api_key='your-key')]\n"
                "result = run_discovery(program, evaluator, config=config)\n\n"
                "Or pass model= directly:\n"
                "result = run_discovery(program, evaluator, model='gpt-5')"
            )

        # Create and run controller
        controller = MainController(
            initial_program_path=program_path,
            evaluation_file=evaluator_path,
            config=config_obj,
            output_dir=actual_output_dir,
        )

        best_program = await controller.run(iterations=iterations)

        # Prepare result
        best_score = 0.0
        metrics = {}
        best_solution = ""

        if best_program:
            best_solution = best_program.solution
            metrics = best_program.metrics or {}

            if "combined_score" in metrics:
                best_score = metrics["combined_score"]
            elif metrics:
                numeric_metrics = [v for v in metrics.values() if isinstance(v, (int, float))]
                if numeric_metrics:
                    best_score = sum(numeric_metrics) / len(numeric_metrics)

        # Extract initial score from the database
        initial_score = None
        if controller.database and controller.database.programs:
            # Find the initial program in the database
            initial_program_solution = controller.initial_program_solution
            initial_program = None
            
            # Search for the initial program by advancement (iteration 0) or by matching code
            for program in controller.database.programs.values():
                if program.solution == initial_program_solution:
                    initial_program = program
                    break
            
            # If not found by code match, try to find the one with iteration 0
            if initial_program is None:
                for program in controller.database.programs.values():
                    if program.iteration_found == 0:
                        initial_program = program
                        break
            
            if initial_program and initial_program.metrics:
                initial_metrics = initial_program.metrics
                if "combined_score" in initial_metrics:
                    initial_score = initial_metrics["combined_score"]
                else:
                    # Calculate average of numeric metrics
                    numeric_metrics = [
                        v for v in initial_metrics.values() 
                        if isinstance(v, (int, float)) and not isinstance(v, bool)
                    ]
                    if numeric_metrics:
                        initial_score = sum(numeric_metrics) / len(numeric_metrics)

        return DiscoveryResult(
            best_program=best_program,
            best_score=best_score,
            best_solution=best_solution,
            metrics=metrics,
            output_dir=actual_output_dir if not cleanup else None,
            initial_score=initial_score,
        )

    finally:
        # Cleanup temporary files if requested
        if cleanup:
            for temp_file in temp_files:
                try:
                    os.unlink(temp_file)
                except Exception as e:
                    print(f"Error deleting temp file: {e}")
                    pass
            if temp_dir and os.path.exists(temp_dir):
                import shutil

                try:
                    shutil.rmtree(temp_dir)
                except Exception as e:
                    print(f"Error deleting temp directory: {e}")
                    pass


def _prepare_program(
    initial_program: Union[str, Path, List[str]], temp_dir: Optional[str], temp_files: List[str]
) -> str:
    """Convert program input to a file path"""

    # If already a file path, use it directly
    if isinstance(initial_program, (str, Path)):
        if os.path.exists(str(initial_program)):
            return str(initial_program)

    # Otherwise, treat as solution string and write to temp file
    if isinstance(initial_program, list):
        solution = "\n".join(initial_program)
    else:
        solution = str(initial_program)

    # Ensure solution has evolution markers if it doesn't already
    if "EVOLVE-BLOCK-START" not in solution:
        # Wrap entire solution in evolution block
        solution = f"""# EVOLVE-BLOCK-START
{solution}
# EVOLVE-BLOCK-END"""

    # Write to temp file
    if temp_dir is None:
        temp_dir = tempfile.gettempdir()

    program_file = os.path.join(temp_dir, f"program_{uuid.uuid4().hex[:8]}.py")
    with open(program_file, "w") as f:
        f.write(solution)
    temp_files.append(program_file)

    return program_file


def _prepare_evaluator(
    evaluator: Union[str, Path, Callable], temp_dir: Optional[str], temp_files: List[str]
) -> str:
    """Convert evaluator input to a file path"""

    # If already a file path, use it directly
    if isinstance(evaluator, (str, Path)):
        if os.path.exists(str(evaluator)):
            return str(evaluator)

    # If it's a callable, create a wrapper module
    if callable(evaluator):
        # Create a unique global name for this evaluator
        evaluator_id = f"_skydiscover_evaluator_{uuid.uuid4().hex[:8]}"

        # Store in globals so the wrapper can find it
        globals()[evaluator_id] = evaluator

        evaluator_code = f"""
# Wrapper for user-provided evaluator function
import {__name__} as api_module

def evaluate(program_path):
    '''Wrapper for user-provided evaluator function'''
    user_evaluator = getattr(api_module, '{evaluator_id}')
    return user_evaluator(program_path)
"""
    else:
        # Treat as code string
        evaluator_code = str(evaluator)

        # Ensure it has an evaluate function
        if "def evaluate" not in evaluator_code:
            raise ValueError("Evaluator code must contain an 'evaluate(program_path)' function")

    # Write to temp file
    if temp_dir is None:
        temp_dir = tempfile.gettempdir()

    eval_file = os.path.join(temp_dir, f"evaluator_{uuid.uuid4().hex[:8]}.py")
    with open(eval_file, "w") as f:
        f.write(evaluator_code)
    temp_files.append(eval_file)

    return eval_file


def discover_solution(
    initial_solution: str,
    evaluator: Callable[[str], Dict[str, Any]],
    iterations: int = 100,
    search: Optional[str] = None,
    model: Optional[str] = None,
    **kwargs,
) -> DiscoveryResult:
    """
    Discover an arbitrary solution (code, prompt, etc.) with a custom evaluator and initial solution.

    Args:
        initial_solution: Initial solution string (code, prompt, etc.) to evolve.
        evaluator: Callable that takes the solution string and returns a metrics dict
            (must include "combined_score").
        iterations: Number of iterations.
        search: Search algorithm (e.g. "adaevolve", "evox", "topk", "openevolve", "gepa").
        model: LLM model (e.g. "gpt-5").
        **kwargs: Additional arguments passed to run_discovery.
    """
    return run_discovery(
        initial_program=initial_solution,
        evaluator=evaluator,
        iterations=iterations,
        search=search,
        model=model,
        **kwargs,
    )
