"""
High-value edge case tests for benchmarks/optimization/optuna_optimizer.py

Tests cover:
- Edge cases in parameter optimization (_objective with different param types)
- Trial generation and convergence scenarios
- Constraint violation / error handling paths (all try/except blocks)
- Boundary conditions in speed score calculation
- Invalid input handling (zero weights, empty param spaces)
- State management edge cases (_optimization_callback, trials_history)
- Visualization error paths and corner cases
- _save_results with missing/partial study data
- Convenience function argument forwarding
"""

import numpy as np
import pytest
from unittest.mock import Mock, patch


MODULE = "local_deep_research.benchmarks.optimization.optuna_optimizer"


def _make_optimizer(**kwargs):
    """Helper to create an OptunaOptimizer with mocked evaluator."""
    from local_deep_research.benchmarks.optimization.optuna_optimizer import (
        OptunaOptimizer,
    )

    defaults = {"base_query": "test query"}
    defaults.update(kwargs)
    return OptunaOptimizer(**defaults)


class TestObjectiveParameterTypes:
    """Tests for _objective handling of different parameter types (int, float, categorical)."""

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_objective_handles_float_param_type(self, mock_evaluator):
        """Test that _objective correctly suggests float parameters."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()

        mock_trial = Mock()
        mock_trial.number = 0
        mock_trial.suggest_float.return_value = 0.5
        mock_trial.set_user_attr = Mock()

        param_space = {
            "learning_rate": {
                "type": "float",
                "low": 0.01,
                "high": 1.0,
                "step": None,
                "log": True,
            }
        }

        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.6}
            score = optimizer._objective(mock_trial, param_space=param_space)

        mock_trial.suggest_float.assert_called_once_with(
            "learning_rate", 0.01, 1.0, step=None, log=True
        )
        assert score == 0.6

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_objective_handles_unknown_param_type_silently(
        self, mock_evaluator
    ):
        """Test that _objective skips params with unrecognized type strings."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()

        mock_trial = Mock()
        mock_trial.number = 0
        mock_trial.set_user_attr = Mock()

        # A param type that doesn't match any branch
        param_space = {
            "unknown_param": {
                "type": "boolean",
                "default": True,
            }
        }

        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.5}
            score = optimizer._objective(mock_trial, param_space=param_space)

        # The unknown param type should be silently skipped
        assert score == 0.5
        mock_trial.suggest_int.assert_not_called()
        mock_trial.suggest_float.assert_not_called()
        mock_trial.suggest_categorical.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_objective_with_empty_param_space(self, mock_evaluator):
        """Test _objective with an empty parameter space dict."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()

        mock_trial = Mock()
        mock_trial.number = 0
        mock_trial.set_user_attr = Mock()

        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.3}
            score = optimizer._objective(mock_trial, param_space={})

        assert score == 0.3
        # _run_experiment should be called with empty params dict
        mock_run.assert_called_once_with({})


class TestObjectiveErrorAndCallback:
    """Tests for _objective error handling and progress callback invocation."""

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_objective_error_triggers_progress_callback_with_error_status(
        self, mock_evaluator
    ):
        """Test that when _run_experiment raises, the progress callback receives error status."""
        mock_evaluator.return_value = Mock()
        callback = Mock()
        optimizer = _make_optimizer(progress_callback=callback)

        mock_trial = Mock()
        mock_trial.number = 5
        mock_trial.suggest_int.return_value = 2
        mock_trial.suggest_categorical.return_value = "iterdrag"
        mock_trial.set_user_attr = Mock()

        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.side_effect = RuntimeError("GPU OOM")
            param_space = optimizer._get_default_param_space()
            score = optimizer._objective(mock_trial, param_space=param_space)

        assert score == float("-inf")
        # Check that the error callback was invoked
        error_calls = [
            c
            for c in callback.call_args_list
            if c[0][2].get("status") == "error"
        ]
        assert len(error_calls) == 1
        assert "GPU OOM" in error_calls[0][0][2]["error"]

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_objective_success_appends_to_trials_history(self, mock_evaluator):
        """Test that a successful trial is appended to trials_history with correct fields."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()

        mock_trial = Mock()
        mock_trial.number = 7
        mock_trial.suggest_int.return_value = 3
        mock_trial.suggest_categorical.return_value = "standard"
        mock_trial.set_user_attr = Mock()

        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.return_value = {"score": 0.92, "quality_score": 0.95}
            param_space = optimizer._get_default_param_space()
            optimizer._objective(mock_trial, param_space=param_space)

        assert len(optimizer.trials_history) == 1
        entry = optimizer.trials_history[0]
        assert entry["trial_number"] == 7
        assert entry["score"] == 0.92
        assert "params" in entry
        assert "duration" in entry
        assert "timestamp" in entry

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_objective_error_does_not_append_to_trials_history(
        self, mock_evaluator
    ):
        """Test that a failed trial is NOT appended to trials_history."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()

        mock_trial = Mock()
        mock_trial.number = 0
        mock_trial.suggest_int.return_value = 1
        mock_trial.suggest_categorical.return_value = "rapid"
        mock_trial.set_user_attr = Mock()

        with patch.object(optimizer, "_run_experiment") as mock_run:
            mock_run.side_effect = ValueError("bad config")
            param_space = optimizer._get_default_param_space()
            optimizer._objective(mock_trial, param_space=param_space)

        assert len(optimizer.trials_history) == 0


class TestRunExperimentEdgeCases:
    """Tests for _run_experiment edge cases and error paths."""

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_run_experiment_error_returns_failure_dict(
        self, mock_profiler_cls, mock_evaluator
    ):
        """Test that _run_experiment catches exceptions and returns a failure dict."""
        mock_eval_instance = Mock()
        mock_eval_instance.evaluate.side_effect = ConnectionError(
            "network down"
        )
        mock_evaluator.return_value = mock_eval_instance

        mock_profiler = Mock()
        mock_profiler_cls.return_value = mock_profiler

        optimizer = _make_optimizer()
        result = optimizer._run_experiment({"iterations": 1})

        assert result["success"] is False
        assert result["score"] == 0.0
        assert "network down" in result["error"]
        # Profiler stop should still be called on error
        mock_profiler.stop.assert_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_speed_score_boundary_very_fast(
        self, mock_profiler_cls, mock_evaluator
    ):
        """Test speed_score is clamped to 1.0 for very fast durations (<60s)."""
        mock_eval_instance = Mock()
        mock_eval_instance.evaluate.return_value = {
            "quality_score": 0.9,
            "benchmark_results": {},
        }
        mock_evaluator.return_value = mock_eval_instance

        mock_profiler = Mock()
        mock_profiler.get_summary.return_value = {"total_duration": 10.0}
        mock_profiler_cls.return_value = mock_profiler

        optimizer = _make_optimizer(
            metric_weights={"quality": 0.5, "speed": 0.5}
        )
        result = optimizer._run_experiment({"iterations": 1})

        # speed_score = max(0, min(1, 1 - (10-60)/180)) = min(1, 1 + 50/180) = 1.0
        assert result["speed_score"] == 1.0

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_speed_score_boundary_very_slow(
        self, mock_profiler_cls, mock_evaluator
    ):
        """Test speed_score is clamped to 0.0 for very slow durations (>240s)."""
        mock_eval_instance = Mock()
        mock_eval_instance.evaluate.return_value = {
            "quality_score": 0.5,
            "benchmark_results": {},
        }
        mock_evaluator.return_value = mock_eval_instance

        mock_profiler = Mock()
        mock_profiler.get_summary.return_value = {"total_duration": 500.0}
        mock_profiler_cls.return_value = mock_profiler

        optimizer = _make_optimizer()
        result = optimizer._run_experiment({"iterations": 5})

        # speed_score = max(0, min(1, 1 - (500-60)/180)) = max(0, 1 - 2.44) = 0.0
        assert result["speed_score"] == 0.0

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_run_experiment_default_params_when_missing(
        self, mock_profiler_cls, mock_evaluator
    ):
        """Test _run_experiment uses defaults when params dict is sparse."""
        mock_eval_instance = Mock()
        mock_eval_instance.evaluate.return_value = {
            "quality_score": 0.7,
            "benchmark_results": {},
        }
        mock_evaluator.return_value = mock_eval_instance

        mock_profiler = Mock()
        mock_profiler.get_summary.return_value = {"total_duration": 100.0}
        mock_profiler_cls.return_value = mock_profiler

        optimizer = _make_optimizer()
        result = optimizer._run_experiment({})  # empty params

        assert result["success"] is True
        # Should have used defaults and not crashed
        assert "score" in result


class TestSpeedScoreCalculation:
    """Tests for the speed score formula boundary conditions."""

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_speed_score_at_exact_60_seconds(
        self, mock_profiler_cls, mock_evaluator
    ):
        """Test speed_score at exactly 60 seconds (should be 1.0)."""
        mock_eval_instance = Mock()
        mock_eval_instance.evaluate.return_value = {
            "quality_score": 0.5,
            "benchmark_results": {},
        }
        mock_evaluator.return_value = mock_eval_instance

        mock_profiler = Mock()
        mock_profiler.get_summary.return_value = {"total_duration": 60.0}
        mock_profiler_cls.return_value = mock_profiler

        optimizer = _make_optimizer()
        result = optimizer._run_experiment({"iterations": 1})

        # speed_score = max(0, min(1, 1 - (60-60)/180)) = 1.0
        assert result["speed_score"] == pytest.approx(1.0)

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.SpeedProfiler")
    def test_speed_score_at_exact_240_seconds(
        self, mock_profiler_cls, mock_evaluator
    ):
        """Test speed_score at exactly 240 seconds (should be 0.0)."""
        mock_eval_instance = Mock()
        mock_eval_instance.evaluate.return_value = {
            "quality_score": 0.5,
            "benchmark_results": {},
        }
        mock_evaluator.return_value = mock_eval_instance

        mock_profiler = Mock()
        mock_profiler.get_summary.return_value = {"total_duration": 240.0}
        mock_profiler_cls.return_value = mock_profiler

        optimizer = _make_optimizer()
        result = optimizer._run_experiment({"iterations": 1})

        # speed_score = max(0, min(1, 1 - (240-60)/180)) = max(0, 0) = 0.0
        assert result["speed_score"] == pytest.approx(0.0)


class TestWeightNormalizationEdgeCases:
    """Tests for metric weight normalization edge cases."""

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_zero_total_weight_preserved(self, mock_evaluator):
        """Test that zero total weight does not cause division by zero."""
        mock_evaluator.return_value = Mock()
        # All weights are zero - normalization guard: total_weight > 0 is False
        optimizer = _make_optimizer(metric_weights={"quality": 0, "speed": 0})

        # Weights should remain as-is (all zeros) since total is 0
        assert optimizer.metric_weights["quality"] == 0
        assert optimizer.metric_weights["speed"] == 0

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_single_metric_weight_normalizes_to_one(self, mock_evaluator):
        """Test that a single metric weight normalizes to 1.0."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(metric_weights={"quality": 5.0})

        assert optimizer.metric_weights["quality"] == pytest.approx(1.0)

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_many_metrics_normalize_correctly(self, mock_evaluator):
        """Test normalization with many metrics."""
        mock_evaluator.return_value = Mock()
        weights = {"quality": 3.0, "speed": 2.0, "resource": 1.0, "cost": 4.0}
        optimizer = _make_optimizer(metric_weights=weights)

        total = sum(optimizer.metric_weights.values())
        assert total == pytest.approx(1.0)
        assert optimizer.metric_weights["quality"] == pytest.approx(0.3)
        assert optimizer.metric_weights["cost"] == pytest.approx(0.4)


class TestOptimizationCallbackEdgeCases:
    """Tests for _optimization_callback periodic save behavior."""

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_callback_saves_at_trial_10(self, mock_evaluator):
        """Test that _optimization_callback triggers save at trial 10."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()

        mock_study = Mock()
        mock_trial = Mock()
        mock_trial.number = 10

        with (
            patch.object(optimizer, "_save_results") as mock_save,
            patch.object(optimizer, "_create_quick_visualizations") as mock_viz,
        ):
            optimizer._optimization_callback(mock_study, mock_trial)

        mock_save.assert_called_once()
        mock_viz.assert_called_once()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_callback_does_not_save_at_trial_0(self, mock_evaluator):
        """Test that _optimization_callback does NOT trigger save at trial 0."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()

        mock_study = Mock()
        mock_trial = Mock()
        mock_trial.number = 0

        with (
            patch.object(optimizer, "_save_results") as mock_save,
            patch.object(optimizer, "_create_quick_visualizations") as mock_viz,
        ):
            optimizer._optimization_callback(mock_study, mock_trial)

        mock_save.assert_not_called()
        mock_viz.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_callback_does_not_save_at_non_multiple_of_10(self, mock_evaluator):
        """Test that _optimization_callback does NOT save at trial 7."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()

        mock_study = Mock()
        mock_trial = Mock()
        mock_trial.number = 7

        with patch.object(optimizer, "_save_results") as mock_save:
            optimizer._optimization_callback(mock_study, mock_trial)

        mock_save.assert_not_called()


class TestSaveResultsEdgeCases:
    """Tests for _save_results with edge case data."""

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_save_results_with_numpy_values_in_nested_dicts(
        self, mock_write_json, mock_joblib, mock_evaluator, tmp_path
    ):
        """Test _save_results converts numpy types in nested trial dicts."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))

        mock_study = Mock()
        mock_study.best_params = {"iterations": 2}
        mock_study.best_value = np.float64(0.85)
        mock_study.best_trial = Mock()
        mock_study.best_trial.user_attrs = {}
        mock_study.trials = [Mock()]
        optimizer.study = mock_study
        optimizer.best_params = {"iterations": 2}

        # Include numpy values in trials_history
        optimizer.trials_history = [
            {
                "trial_number": 0,
                "params": {"iterations": np.int64(2)},
                "score": np.float64(0.85),
                "result": {"quality_score": np.float32(0.9)},
            }
        ]

        optimizer._save_results()

        # Verify write_json_verified was called (history + best params)
        assert mock_write_json.call_count == 2

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_save_results_with_no_study_skips_best_params(
        self, mock_write_json, mock_joblib, mock_evaluator, tmp_path
    ):
        """Test _save_results when study is None skips best_params and study save."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        optimizer.study = None
        optimizer.best_params = None
        optimizer.trials_history = []

        optimizer._save_results()

        # Only history file should be written, not best_params
        assert mock_write_json.call_count == 1
        mock_joblib.dump.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_save_results_study_without_best_params(
        self, mock_write_json, mock_joblib, mock_evaluator, tmp_path
    ):
        """Test _save_results when study exists but best_params is empty/falsy."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))

        mock_study = Mock()
        mock_study.best_params = {}  # empty dict is falsy
        mock_study.best_value = 0.0
        mock_study.trials = []
        optimizer.study = mock_study
        optimizer.best_params = {}
        optimizer.trials_history = []

        optimizer._save_results()

        # history is written + study.pkl is saved, but best_params JSON is skipped
        # because self.study.best_params is {} which is falsy
        assert mock_write_json.call_count == 1
        mock_joblib.dump.assert_called_once()


class TestVisualizationEdgeCases:
    """Tests for visualization methods handling edge cases."""

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", False)
    def test_create_visualizations_returns_early_without_matplotlib(
        self, mock_evaluator
    ):
        """Test _create_visualizations returns early when matplotlib is unavailable."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        optimizer.study = Mock()
        optimizer.study.trials = [Mock(), Mock()]

        # Should not raise and should return early
        optimizer._create_visualizations()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", True)
    def test_create_visualizations_returns_early_with_less_than_2_trials(
        self, mock_evaluator, tmp_path
    ):
        """Test _create_visualizations returns early with fewer than 2 trials."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        optimizer.study = Mock()
        optimizer.study.trials = [Mock()]  # Only 1 trial
        optimizer.trials_history = []

        # Should not raise - returns early due to insufficient trials
        with patch.object(
            optimizer, "_create_optuna_visualizations"
        ) as mock_optuna_viz:
            optimizer._create_visualizations()
            mock_optuna_viz.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_create_custom_visualizations_returns_early_with_no_history(
        self, mock_evaluator, tmp_path
    ):
        """Test _create_custom_visualizations returns early with empty trials_history."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        optimizer.trials_history = []

        # Should not raise or try to create any plots
        with patch(f"{MODULE}.plt") as mock_plt:
            optimizer._create_custom_visualizations(str(tmp_path))
            mock_plt.figure.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_quality_vs_speed_plot_no_successful_trials(
        self, mock_evaluator, tmp_path
    ):
        """Test _create_quality_vs_speed_plot with only failed trials."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        optimizer.trials_history = [
            {"result": {"success": False}, "params": {}, "score": 0.0},
        ]

        with patch(f"{MODULE}.plt") as mock_plt:
            optimizer._create_quality_vs_speed_plot(str(tmp_path), "20260304")
            # Should return early because no successful trials
            mock_plt.figure.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", True)
    def test_quick_visualizations_handles_plot_error(
        self, mock_evaluator, tmp_path
    ):
        """Test _create_quick_visualizations gracefully handles plot errors."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        optimizer.study = Mock()
        optimizer.study.trials = [Mock(), Mock()]

        with patch(f"{MODULE}.plot_optimization_history") as mock_plot:
            mock_plot.side_effect = RuntimeError("plot failed")
            # Should not raise
            optimizer._create_quick_visualizations()


class TestOptimizeMethodEdgeCases:
    """Tests for optimize() method edge cases."""

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimize_calls_progress_callback_on_start(
        self, mock_optuna, mock_evaluator
    ):
        """Test optimize() calls progress_callback with 'starting' status."""
        mock_evaluator.return_value = Mock()
        callback = Mock()

        mock_study = Mock()
        mock_study.best_params = {"iterations": 1}
        mock_study.best_value = 0.5
        mock_study.best_trial = Mock()
        mock_study.best_trial.user_attrs = {}
        mock_study.trials = []
        mock_optuna.create_study.return_value = mock_study

        optimizer = _make_optimizer(n_trials=1, progress_callback=callback)

        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            optimizer.optimize()

        # First callback call should have status "starting"
        first_call = callback.call_args_list[0]
        assert first_call[0][0] == 0  # trial_num = 0
        assert first_call[0][2]["status"] == "starting"

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimize_calls_progress_callback_on_completion(
        self, mock_optuna, mock_evaluator
    ):
        """Test optimize() calls progress_callback with 'completed' status at the end."""
        mock_evaluator.return_value = Mock()
        callback = Mock()

        mock_study = Mock()
        mock_study.best_params = {"iterations": 2}
        mock_study.best_value = 0.9
        mock_study.best_trial = Mock()
        mock_study.best_trial.user_attrs = {}
        mock_study.trials = [Mock()]
        mock_optuna.create_study.return_value = mock_study

        optimizer = _make_optimizer(n_trials=1, progress_callback=callback)

        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            optimizer.optimize()

        # Last callback call should have status "completed"
        last_call = callback.call_args_list[-1]
        assert last_call[0][2]["status"] == "completed"
        assert last_call[0][2]["best_value"] == 0.9

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimize_handles_keyboard_interrupt(
        self, mock_optuna, mock_evaluator
    ):
        """Test optimize() handles KeyboardInterrupt, saves results, and returns best."""
        mock_evaluator.return_value = Mock()

        mock_study = Mock()
        mock_study.best_params = {"iterations": 1}
        mock_study.best_value = 0.3
        mock_study.best_trial = Mock()
        mock_study.best_trial.user_attrs = {}
        mock_study.trials = [Mock()]
        mock_study.optimize.side_effect = KeyboardInterrupt()
        mock_optuna.create_study.return_value = mock_study

        optimizer = _make_optimizer(n_trials=10)

        with (
            patch.object(optimizer, "_save_results") as mock_save,
            patch.object(optimizer, "_create_visualizations") as mock_viz,
        ):
            best_params, best_value = optimizer.optimize()

        # Should still save and visualize
        mock_save.assert_called_once()
        mock_viz.assert_called_once()
        assert best_params == {"iterations": 1}
        assert best_value == 0.3

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimize_keyboard_interrupt_with_callback(
        self, mock_optuna, mock_evaluator
    ):
        """Test optimize() invokes callback with 'interrupted' on KeyboardInterrupt."""
        mock_evaluator.return_value = Mock()
        callback = Mock()

        mock_study = Mock()
        mock_study.best_params = {"iterations": 1}
        mock_study.best_value = 0.2
        mock_study.best_trial = Mock()
        mock_study.best_trial.user_attrs = {}
        mock_study.trials = [Mock(), Mock()]
        mock_study.optimize.side_effect = KeyboardInterrupt()
        mock_optuna.create_study.return_value = mock_study

        optimizer = _make_optimizer(n_trials=10, progress_callback=callback)

        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            optimizer.optimize()

        # Find the interrupted callback
        interrupted_calls = [
            c
            for c in callback.call_args_list
            if c[0][2].get("status") == "interrupted"
        ]
        assert len(interrupted_calls) == 1
        assert interrupted_calls[0][0][2]["trials_completed"] == 2

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimize_uses_custom_param_space(
        self, mock_optuna, mock_evaluator
    ):
        """Test optimize() passes custom param_space through to _objective."""
        mock_evaluator.return_value = Mock()

        mock_study = Mock()
        mock_study.best_params = {"custom_param": 5}
        mock_study.best_value = 0.7
        mock_study.best_trial = Mock()
        mock_study.best_trial.user_attrs = {}
        mock_study.trials = []
        mock_optuna.create_study.return_value = mock_study

        optimizer = _make_optimizer(n_trials=1)

        custom_space = {
            "custom_param": {"type": "int", "low": 1, "high": 10, "step": 1}
        }

        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            best_params, best_value = optimizer.optimize(
                param_space=custom_space
            )

        # Verify study.optimize was called (the partial wrapping the custom space)
        mock_study.optimize.assert_called_once()
        assert best_params == {"custom_param": 5}


class TestConvenienceFunctionForwarding:
    """Tests for convenience functions forwarding arguments correctly."""

    @patch(f"{MODULE}.OptunaOptimizer")
    def test_optimize_parameters_forwards_all_kwargs(self, mock_cls):
        """Test optimize_parameters forwards all keyword arguments."""
        from local_deep_research.benchmarks.optimization.optuna_optimizer import (
            optimize_parameters,
        )

        mock_instance = Mock()
        mock_instance.optimize.return_value = ({"a": 1}, 0.5)
        mock_cls.return_value = mock_instance

        callback = Mock()
        optimize_parameters(
            query="my query",
            output_dir="/tmp/out",
            model_name="gpt-4",
            provider="openai",
            search_tool="google",
            temperature=0.3,
            n_trials=15,
            timeout=600,
            n_jobs=2,
            study_name="my_study",
            optimization_metrics=["quality"],
            metric_weights={"quality": 1.0},
            progress_callback=callback,
            benchmark_weights={"simpleqa": 0.5, "browsecomp": 0.5},
        )

        call_kwargs = mock_cls.call_args[1]
        assert call_kwargs["base_query"] == "my query"
        assert call_kwargs["output_dir"] == "/tmp/out"
        assert call_kwargs["model_name"] == "gpt-4"
        assert call_kwargs["provider"] == "openai"
        assert call_kwargs["temperature"] == 0.3
        assert call_kwargs["n_trials"] == 15
        assert call_kwargs["timeout"] == 600
        assert call_kwargs["n_jobs"] == 2
        assert call_kwargs["study_name"] == "my_study"
        assert call_kwargs["progress_callback"] is callback
        assert call_kwargs["benchmark_weights"] == {
            "simpleqa": 0.5,
            "browsecomp": 0.5,
        }

    @patch(f"{MODULE}.OptunaOptimizer")
    def test_optimize_for_speed_passes_reduced_param_space(self, mock_cls):
        """Test optimize_for_speed provides a reduced param space (max iterations=3)."""
        from local_deep_research.benchmarks.optimization.optuna_optimizer import (
            optimize_for_speed,
        )

        mock_instance = Mock()
        mock_instance.optimize.return_value = ({}, 0.0)
        mock_cls.return_value = mock_instance

        optimize_for_speed(query="test", n_trials=5)

        # Check the param_space passed to optimize()
        optimize_call = mock_instance.optimize.call_args
        param_space = optimize_call[1].get("param_space") or optimize_call[0][0]
        assert param_space["iterations"]["high"] == 3
        assert "focused-iteration" in param_space["search_strategy"]["choices"]