"""
API functions for optimization tasks in Local Deep Research.

This module provides a simplified interface for parameter optimization
without having to directly work with the optimizer classes.
"""

from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple

# No metrics imports needed here, they're used in the OptunaOptimizer
from .optuna_optimizer import OptunaOptimizer


def optimize_parameters(
    query: str,
    param_space: Optional[Dict[str, Any]] = None,
    output_dir: str = str(Path("data") / "optimization_results"),
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    search_tool: Optional[str] = None,
    temperature: float = 0.7,
    n_trials: int = 30,
    timeout: Optional[int] = None,
    n_jobs: int = 1,
    study_name: Optional[str] = None,
    optimization_metrics: Optional[List[str]] = None,
    metric_weights: Optional[Dict[str, float]] = None,
    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
    benchmark_weights: Optional[Dict[str, float]] = None,
) -> Tuple[Dict[str, Any], float]:
    """
    Optimize parameters for Local Deep Research.

    Args:
        query: The research query to use for all experiments
        param_space: Dictionary defining parameter search spaces (optional)
        output_dir: Directory to save optimization results
        model_name: Name of the LLM model to use
        provider: LLM provider
        search_tool: Search engine to use
        temperature: LLM temperature
        n_trials: Number of parameter combinations to try
        timeout: Maximum seconds to run optimization (None for no limit)
        n_jobs: Number of parallel jobs for optimization
        study_name: Name of the Optuna study
        optimization_metrics: List of metrics to optimize (default: ["quality", "speed"])
        metric_weights: Dictionary of weights for each metric
        progress_callback: Optional callback for progress updates
        benchmark_weights: Dictionary mapping benchmark types to weights
            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
            If None, only SimpleQA is used with weight 1.0

    Returns:
        Tuple of (best_parameters, best_score)
    """
    # Create optimizer
    optimizer = OptunaOptimizer(
        base_query=query,
        output_dir=output_dir,
        model_name=model_name,
        provider=provider,
        search_tool=search_tool,
        temperature=temperature,
        n_trials=n_trials,
        timeout=timeout,
        n_jobs=n_jobs,
        study_name=study_name,
        optimization_metrics=optimization_metrics,
        metric_weights=metric_weights,
        progress_callback=progress_callback,
        benchmark_weights=benchmark_weights,
    )

    # Run optimization
    return optimizer.optimize(param_space)


def optimize_for_speed(
    query: str,
    n_trials: int = 20,
    output_dir: str = str(Path("data") / "optimization_results"),
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    search_tool: Optional[str] = None,
    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
    benchmark_weights: Optional[Dict[str, float]] = None,
) -> Tuple[Dict[str, Any], float]:
    """
    Optimize parameters with a focus on speed performance.

    Args:
        query: The research query to use for all experiments
        n_trials: Number of parameter combinations to try
        output_dir: Directory to save optimization results
        model_name: Name of the LLM model to use
        provider: LLM provider
        search_tool: Search engine to use
        progress_callback: Optional callback for progress updates
        benchmark_weights: Dictionary mapping benchmark types to weights
            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
            If None, only SimpleQA is used with weight 1.0

    Returns:
        Tuple of (best_parameters, best_score)
    """
    # Focus on speed with reduced parameter space
    param_space = {
        "iterations": {
            "type": "int",
            "low": 1,
            "high": 3,
            "step": 1,
        },
        "questions_per_iteration": {
            "type": "int",
            "low": 1,
            "high": 3,
            "step": 1,
        },
        "search_strategy": {
            "type": "categorical",
            "choices": ["rapid", "parallel", "source_based"],
        },
    }

    # Speed-focused weights
    metric_weights = {"speed": 0.8, "quality": 0.2, "resource": 0.0}

    return optimize_parameters(
        query=query,
        param_space=param_space,
        output_dir=output_dir,
        model_name=model_name,
        provider=provider,
        search_tool=search_tool,
        n_trials=n_trials,
        metric_weights=metric_weights,
        optimization_metrics=["speed", "quality"],
        progress_callback=progress_callback,
        benchmark_weights=benchmark_weights,
    )


def optimize_for_quality(
    query: str,
    n_trials: int = 30,
    output_dir: str = str(Path("data") / "optimization_results"),
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    search_tool: Optional[str] = None,
    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
    benchmark_weights: Optional[Dict[str, float]] = None,
) -> Tuple[Dict[str, Any], float]:
    """
    Optimize parameters with a focus on result quality.

    Args:
        query: The research query to use for all experiments
        n_trials: Number of parameter combinations to try
        output_dir: Directory to save optimization results
        model_name: Name of the LLM model to use
        provider: LLM provider
        search_tool: Search engine to use
        progress_callback: Optional callback for progress updates
        benchmark_weights: Dictionary mapping benchmark types to weights
            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
            If None, only SimpleQA is used with weight 1.0

    Returns:
        Tuple of (best_parameters, best_score)
    """
    # Quality-focused weights
    metric_weights = {"quality": 0.9, "speed": 0.1, "resource": 0.0}

    return optimize_parameters(
        query=query,
        output_dir=output_dir,
        model_name=model_name,
        provider=provider,
        search_tool=search_tool,
        n_trials=n_trials,
        metric_weights=metric_weights,
        optimization_metrics=["quality", "speed"],
        progress_callback=progress_callback,
        benchmark_weights=benchmark_weights,
    )


def optimize_for_efficiency(
    query: str,
    n_trials: int = 25,
    output_dir: str = str(Path("data") / "optimization_results"),
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    search_tool: Optional[str] = None,
    progress_callback: Optional[Callable[[int, int, Dict], None]] = None,
    benchmark_weights: Optional[Dict[str, float]] = None,
) -> Tuple[Dict[str, Any], float]:
    """
    Optimize parameters with a focus on resource efficiency.

    Args:
        query: The research query to use for all experiments
        n_trials: Number of parameter combinations to try
        output_dir: Directory to save optimization results
        model_name: Name of the LLM model to use
        provider: LLM provider
        search_tool: Search engine to use
        progress_callback: Optional callback for progress updates
        benchmark_weights: Dictionary mapping benchmark types to weights
            (e.g., {"simpleqa": 0.6, "browsecomp": 0.4})
            If None, only SimpleQA is used with weight 1.0

    Returns:
        Tuple of (best_parameters, best_score)
    """
    # Balance of quality, speed and resource usage
    metric_weights = {"quality": 0.4, "speed": 0.3, "resource": 0.3}

    return optimize_parameters(
        query=query,
        output_dir=output_dir,
        model_name=model_name,
        provider=provider,
        search_tool=search_tool,
        n_trials=n_trials,
        metric_weights=metric_weights,
        optimization_metrics=["quality", "speed", "resource"],
        progress_callback=progress_callback,
        benchmark_weights=benchmark_weights,
    )


def get_default_param_space() -> Dict[str, Any]:
    """
    Get the default parameter search space for optimization.

    Returns:
        Dictionary defining the default parameter search spaces
    """
    return {
        "iterations": {
            "type": "int",
            "low": 1,
            "high": 5,
            "step": 1,
        },
        "questions_per_iteration": {
            "type": "int",
            "low": 1,
            "high": 5,
            "step": 1,
        },
        "search_strategy": {
            "type": "categorical",
            "choices": [
                "iterdrag",
                "standard",
                "rapid",
                "parallel",
                "source_based",
            ],
        },
        "max_results": {
            "type": "int",
            "low": 10,
            "high": 100,
            "step": 10,
        },
        "max_filtered_results": {
            "type": "int",
            "low": 5,
            "high": 50,
            "step": 5,
        },
    }
