"""
Branch-coverage tests for benchmarks/optimization/optuna_optimizer.py.

Targets branches not fully exercised by the existing test files:
- _get_default_param_space structure and types
- optimize() raising KeyboardInterrupt
- progress_callback invocation during optimize()
- _save_results: joblib.dump called for study
- _create_visualizations with PLOTTING_AVAILABLE=False
- metric_weights normalisation when sum != 1.0
"""

from unittest.mock import Mock, patch

MODULE = "local_deep_research.benchmarks.optimization.optuna_optimizer"


def _make_optimizer(**kwargs):
    from local_deep_research.benchmarks.optimization.optuna_optimizer import (
        OptunaOptimizer,
    )

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


# ---------------------------------------------------------------------------
# _get_default_param_space
# ---------------------------------------------------------------------------


class TestGetDefaultParamSpace:
    # test_get_default_param_space_iterations_is_int_type is defined first so it
    # runs first and warms up the expensive module import within the pytest-timeout
    # window before the other tests (including the bare test_get_default_param_space)
    # are collected.
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_get_default_param_space_iterations_is_int_type(
        self, mock_evaluator
    ):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        space = optimizer._get_default_param_space()
        assert space["iterations"]["type"] == "int"
        assert space["iterations"]["low"] >= 1
        assert space["iterations"]["high"] >= space["iterations"]["low"]

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_get_default_param_space(self, mock_evaluator):
        """_get_default_param_space returns a dict with the four expected keys."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        space = optimizer._get_default_param_space()
        assert isinstance(space, dict)
        assert "iterations" in space
        assert "questions_per_iteration" in space
        assert "search_strategy" in space
        assert "max_results" in space

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_get_default_param_space_search_strategy_is_categorical(
        self, mock_evaluator
    ):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        space = optimizer._get_default_param_space()
        assert space["search_strategy"]["type"] == "categorical"
        choices = space["search_strategy"]["choices"]
        assert isinstance(choices, list)
        assert len(choices) > 0
        assert "source-based" in choices

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_get_default_param_space_max_results_step(self, mock_evaluator):
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        space = optimizer._get_default_param_space()
        mr = space["max_results"]
        assert mr["type"] == "int"
        assert mr["low"] > 0
        assert mr["high"] > mr["low"]


# ---------------------------------------------------------------------------
# optimize() – KeyboardInterrupt
# ---------------------------------------------------------------------------


class TestOptimizeKeyboardInterrupt:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimize_keyboard_interrupt(self, mock_optuna, mock_evaluator):
        """When study.optimize raises KeyboardInterrupt, best_params and value are still returned."""
        mock_evaluator.return_value = Mock()
        mock_study = Mock()
        mock_study.best_params = {"iterations": 2, "max_results": 50}
        mock_study.best_value = 0.55
        mock_study.trials = [Mock(), Mock()]
        mock_study.optimize.side_effect = KeyboardInterrupt()
        mock_optuna.create_study.return_value = mock_study
        mock_optuna.samplers.TPESampler.return_value = Mock()

        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()

        assert best_params == {"iterations": 2, "max_results": 50}
        assert best_value == 0.55
        mock_save.assert_called_once()
        mock_viz.assert_called_once()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimize_keyboard_interrupt_with_callback(
        self, mock_optuna, mock_evaluator
    ):
        """KeyboardInterrupt fires an 'interrupted' status callback."""
        mock_evaluator.return_value = Mock()
        callback = Mock()
        mock_study = Mock()
        mock_study.best_params = {"iterations": 1}
        mock_study.best_value = 0.3
        mock_study.trials = [Mock(), Mock(), Mock()]
        mock_study.optimize.side_effect = KeyboardInterrupt()
        mock_optuna.create_study.return_value = mock_study
        mock_optuna.samplers.TPESampler.return_value = Mock()

        optimizer = _make_optimizer(n_trials=5, progress_callback=callback)
        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            optimizer.optimize()

        interrupted_calls = [
            c
            for c in callback.call_args_list
            if c[0][2].get("status") == "interrupted"
        ]
        assert len(interrupted_calls) == 1
        info = interrupted_calls[0][0][2]
        assert info["stage"] == "interrupted"
        assert info["trials_completed"] == 3

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimize_keyboard_interrupt_no_callback(
        self, mock_optuna, mock_evaluator
    ):
        """KeyboardInterrupt without a callback does not raise."""
        mock_evaluator.return_value = Mock()
        mock_study = Mock()
        mock_study.best_params = {"iterations": 1}
        mock_study.best_value = 0.1
        mock_study.trials = []
        mock_study.optimize.side_effect = KeyboardInterrupt()
        mock_optuna.create_study.return_value = mock_study
        mock_optuna.samplers.TPESampler.return_value = Mock()

        optimizer = _make_optimizer(n_trials=3)
        assert optimizer.progress_callback is None
        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            params, value = optimizer.optimize()
        assert params == {"iterations": 1}


# ---------------------------------------------------------------------------
# optimize() – progress_callback invoked
# ---------------------------------------------------------------------------


class TestOptimizationCallbackInvoked:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimization_callback_invoked(self, mock_optuna, mock_evaluator):
        """progress_callback is called with 'starting' status before study.optimize."""
        mock_evaluator.return_value = Mock()
        callback = Mock()
        mock_study = Mock()
        mock_study.best_params = {"iterations": 3}
        mock_study.best_value = 0.7
        mock_study.trials = [Mock()]
        mock_optuna.create_study.return_value = mock_study
        mock_optuna.samplers.TPESampler.return_value = Mock()

        optimizer = _make_optimizer(n_trials=1, progress_callback=callback)
        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            optimizer.optimize()

        # At least one call should have status 'starting'
        all_statuses = [c[0][2].get("status") for c in callback.call_args_list]
        assert "starting" in all_statuses

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.optuna")
    def test_optimization_callback_invoked_on_completion(
        self, mock_optuna, mock_evaluator
    ):
        """progress_callback is called with 'completed' status after study.optimize."""
        mock_evaluator.return_value = Mock()
        callback = Mock()
        mock_study = Mock()
        mock_study.best_params = {"max_results": 40}
        mock_study.best_value = 0.82
        mock_study.trials = [Mock(), Mock()]
        mock_optuna.create_study.return_value = mock_study
        mock_optuna.samplers.TPESampler.return_value = Mock()

        optimizer = _make_optimizer(n_trials=2, progress_callback=callback)
        with (
            patch.object(optimizer, "_save_results"),
            patch.object(optimizer, "_create_visualizations"),
        ):
            optimizer.optimize()

        completed_calls = [
            c
            for c in callback.call_args_list
            if c[0][2].get("status") == "completed"
        ]
        assert len(completed_calls) == 1
        info = completed_calls[0][0][2]
        assert info["best_value"] == 0.82


# ---------------------------------------------------------------------------
# _save_results – joblib.dump called
# ---------------------------------------------------------------------------


class TestSaveResultsJoblib:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_save_results_joblib(
        self, mock_write_json, mock_joblib, mock_evaluator, tmp_path
    ):
        """_save_results calls joblib.dump to persist the study object."""
        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 = 0.75
        mock_study.trials = [Mock()]
        optimizer.study = mock_study
        optimizer.trials_history = []

        optimizer._save_results()

        mock_joblib.dump.assert_called_once()
        # First arg to dump should be the study, second arg should be the file path
        call_args = mock_joblib.dump.call_args
        assert call_args[0][0] is mock_study
        assert str(call_args[0][1]).endswith(".pkl")

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.joblib")
    @patch(
        "local_deep_research.security.file_write_verifier.write_json_verified"
    )
    def test_save_results_joblib_not_called_without_study(
        self, mock_write_json, mock_joblib, mock_evaluator, tmp_path
    ):
        """joblib.dump is NOT called when study is None."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(output_dir=str(tmp_path))
        optimizer.study = None
        optimizer.trials_history = []

        optimizer._save_results()

        mock_joblib.dump.assert_not_called()


# ---------------------------------------------------------------------------
# _create_visualizations – PLOTTING_AVAILABLE=False
# ---------------------------------------------------------------------------


class TestCreateVisualizationsNoMatplotlib:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", False)
    def test_create_visualizations_no_matplotlib(self, mock_evaluator):
        """_create_visualizations returns early and never calls plot functions when PLOTTING_AVAILABLE=False."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_study = Mock()
        mock_study.trials = [Mock(), Mock(), Mock()]
        optimizer.study = mock_study

        with patch(f"{MODULE}.plot_optimization_history") as mock_plot:
            optimizer._create_visualizations()
            mock_plot.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", False)
    def test_create_visualizations_no_matplotlib_does_not_raise(
        self, mock_evaluator
    ):
        """Calling _create_visualizations without matplotlib available does not raise."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        optimizer.study = Mock()
        optimizer.study.trials = [Mock(), Mock()]
        # Should complete without exception
        optimizer._create_visualizations()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", True)
    def test_create_visualizations_skips_when_no_study(self, mock_evaluator):
        """_create_visualizations returns early when study is None."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        optimizer.study = None

        with patch(f"{MODULE}.plot_optimization_history") as mock_plot:
            optimizer._create_visualizations()
            mock_plot.assert_not_called()

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    @patch(f"{MODULE}.PLOTTING_AVAILABLE", True)
    def test_create_visualizations_skips_with_only_one_trial(
        self, mock_evaluator
    ):
        """_create_visualizations returns early when fewer than 2 trials are present."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        mock_study = Mock()
        mock_study.trials = [Mock()]  # only 1 trial
        optimizer.study = mock_study

        with patch(f"{MODULE}.plot_optimization_history") as mock_plot:
            optimizer._create_visualizations()
            mock_plot.assert_not_called()


# ---------------------------------------------------------------------------
# metric_weights normalisation
# ---------------------------------------------------------------------------


class TestWeightNormalization:
    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_weight_normalization(self, mock_evaluator):
        """Weights that don't sum to 1.0 are normalised so the total becomes 1.0."""
        mock_evaluator.return_value = Mock()
        # Deliberately unbalanced weights
        optimizer = _make_optimizer(
            metric_weights={"quality": 3.0, "speed": 1.0}
        )
        total = sum(optimizer.metric_weights.values())
        assert abs(total - 1.0) < 1e-9

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_weight_normalization_proportions_preserved(self, mock_evaluator):
        """After normalisation, the relative proportions remain correct."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(
            metric_weights={"quality": 3.0, "speed": 1.0}
        )
        # quality was 3x speed, so after normalisation quality should be 0.75
        assert abs(optimizer.metric_weights["quality"] - 0.75) < 1e-9
        assert abs(optimizer.metric_weights["speed"] - 0.25) < 1e-9

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_weight_normalization_already_normalised(self, mock_evaluator):
        """Weights already summing to 1.0 remain unchanged."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(
            metric_weights={"quality": 0.6, "speed": 0.4}
        )
        assert abs(optimizer.metric_weights["quality"] - 0.6) < 1e-9
        assert abs(optimizer.metric_weights["speed"] - 0.4) < 1e-9

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_weight_normalization_three_metrics(self, mock_evaluator):
        """Three-metric weights are also normalised correctly."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer(
            metric_weights={"quality": 4.0, "speed": 3.0, "resource": 3.0}
        )
        total = sum(optimizer.metric_weights.values())
        assert abs(total - 1.0) < 1e-9

    @patch(f"{MODULE}.CompositeBenchmarkEvaluator")
    def test_default_weights_sum_to_one(self, mock_evaluator):
        """Default metric_weights are already normalised."""
        mock_evaluator.return_value = Mock()
        optimizer = _make_optimizer()
        total = sum(optimizer.metric_weights.values())
        assert abs(total - 1.0) < 1e-9