"""CPU-safe contract tests for Bayesian optimization helpers."""

from __future__ import annotations

import builtins
import sys
import types

import pytest
import torch
import torch.nn as nn

from obliteratus import bayesian_optimizer as bo


class _Pipeline:
    def __init__(self):
        self.refusal_directions = {}
        self.handle = None
        self._strong_layers = []
        self.harmful_prompts = ["harmful one", "harmful two"]
        self.use_chat_template = False
        self.freed = 0
        self.logs = []

    def _get_model_device(self, _model):
        return torch.device("cpu")

    def _maybe_apply_chat_template(self, prompts):
        return [f"<chat>{prompt}</chat>" for prompt in prompts]

    def _free_gpu_memory(self):
        self.freed += 1

    def log(self, message):
        self.logs.append(message)


def _install_fake_optuna(monkeypatch, study=None):
    optuna = types.ModuleType("optuna")
    optuna.logging = types.SimpleNamespace(WARNING=30, set_verbosity=lambda _level: None)
    optuna.create_study = lambda **_kwargs: study

    samplers = types.ModuleType("optuna.samplers")

    class TPESampler:
        def __init__(self, **_kwargs):
            pass

    samplers.TPESampler = TPESampler
    monkeypatch.setitem(sys.modules, "optuna", optuna)
    monkeypatch.setitem(sys.modules, "optuna.samplers", samplers)


class _FakeTrial:
    number = 0

    params = {
        "attn_max_weight": 0.8,
        "attn_peak_position": 0.0,
        "attn_min_weight": 0.1,
        "attn_spread": 0.6,
        "mlp_max_weight": 0.6,
        "mlp_peak_position": 1.0,
        "mlp_min_weight": 0.2,
        "mlp_spread": 0.6,
        "dir_idx": 1.0,
    }

    values = (0.2, 0.1)

    def __init__(self, number=0, params=None, values=None):
        self.number = number
        if params is not None:
            self.params = params
        if values is not None:
            self.values = values

    def suggest_float(self, name, _low, _high):
        return self.params[name]


class _FakeStudy:
    def __init__(self, best_trials):
        self.best_trials = best_trials
        self.enqueued = []
        self.objective_values = []

    def enqueue_trial(self, params):
        self.enqueued.append(params)

    def optimize(self, objective, n_trials, show_progress_bar):
        assert show_progress_bar is False
        for number in range(n_trials):
            self.objective_values.append(objective(_FakeTrial(number=number)))


class _TokenBatch(dict):
    def __init__(self):
        super().__init__(input_ids=torch.tensor([[1, 2]]))


class _ReferenceTokenizer:
    def __call__(self, *_args, **_kwargs):
        return _TokenBatch()


class _Layer(nn.Module):
    def __init__(self):
        super().__init__()
        self.self_attn = nn.Module()
        self.self_attn.o_proj = nn.Linear(2, 2)
        self.mlp = nn.Module()
        self.mlp.down_proj = nn.Linear(2, 2)


class _ReferenceModel(nn.Module):
    def __init__(self, layers):
        super().__init__()
        self.model = nn.Module()
        self.model.layers = nn.ModuleList(layers)

    def forward(self, **_kwargs):
        return types.SimpleNamespace(logits=torch.tensor([[[0.0, 1.0]]]))


def _optimization_pipeline(layers):
    pipeline = _Pipeline()
    pipeline.handle = types.SimpleNamespace(
        model=_ReferenceModel(layers),
        tokenizer=_ReferenceTokenizer(),
        architecture="llama",
    )
    pipeline._strong_layers = list(range(len(layers)))
    pipeline.refusal_directions = {
        idx: torch.tensor([float(idx + 1), 1.0])
        for idx in pipeline._strong_layers
    }
    pipeline.norm_preserve = True
    pipeline.projections = []
    pipeline.moe_calls = []

    def project_out(module, direction, names, norm_preserve, regularization):
        pipeline.projections.append(
            {
                "module": module,
                "direction": direction.detach().clone(),
                "names": tuple(names),
                "norm_preserve": norm_preserve,
                "regularization": regularization,
            }
        )
        for name in names:
            proj = getattr(module, name, None)
            if proj is not None and hasattr(proj, "weight"):
                proj.weight.data.add_(10.0)
                return 1
        return 0

    def project_moe(module, direction, **kwargs):
        pipeline.moe_calls.append((module, direction.detach().clone(), kwargs))

    pipeline._project_out_advanced = project_out
    pipeline._project_moe_experts = project_moe
    return pipeline


def test_parametric_layer_weight_boundaries():
    assert bo._parametric_layer_weight(0, 1, 0.8, 0.5, 0.1, 0.2) == pytest.approx(0.8)

    # At the peak, the kernel returns the maximum weight.
    assert bo._parametric_layer_weight(2, 5, 0.9, 0.5, 0.1, 0.25) == pytest.approx(0.9)

    # At the tent edge, it reaches the minimum weight.
    assert bo._parametric_layer_weight(1, 5, 0.9, 0.5, 0.1, 0.25) == pytest.approx(0.1)

    # Outside the spread cutoff, the layer is skipped.
    assert bo._parametric_layer_weight(0, 5, 0.9, 0.5, 0.1, 0.24) == pytest.approx(0.0)

    # Tiny or negative spread is clamped to 0.01.
    assert bo._parametric_layer_weight(0, 101, 0.7, 0.0, 0.2, -1.0) == pytest.approx(0.7)


def test_interpolate_direction_handles_empty_clamps_exact_and_normalized_interpolation():
    pipeline = _Pipeline()
    assert torch.equal(bo._interpolate_direction(pipeline, layer_idx=3, float_dir_idx=1.0), torch.zeros(1))

    pipeline.refusal_directions = {
        2: torch.tensor([3.0, 0.0]),
        5: torch.tensor([0.0, 4.0]),
        9: torch.tensor([1.0, 1.0]),
    }

    low = bo._interpolate_direction(pipeline, layer_idx=5, float_dir_idx=-10.0)
    assert torch.allclose(low, torch.tensor([1.0, 0.0]))

    high = bo._interpolate_direction(pipeline, layer_idx=5, float_dir_idx=99.0)
    assert torch.allclose(high, torch.tensor([2**-0.5, 2**-0.5]))

    exact = bo._interpolate_direction(pipeline, layer_idx=5, float_dir_idx=1.0)
    assert torch.allclose(exact, torch.tensor([0.0, 1.0]))

    interpolated = bo._interpolate_direction(pipeline, layer_idx=5, float_dir_idx=0.5)
    expected = torch.tensor([1.5, 2.0])
    expected = expected / expected.norm()
    assert torch.allclose(interpolated, expected)
    assert interpolated.norm().item() == pytest.approx(1.0)


def test_run_bayesian_optimization_returns_empty_when_optuna_missing(monkeypatch):
    real_import = builtins.__import__

    def fake_import(name, *args, **kwargs):
        if name == "optuna" or name.startswith("optuna."):
            raise ImportError("no optuna in this test")
        return real_import(name, *args, **kwargs)

    monkeypatch.setattr(builtins, "__import__", fake_import)

    pipeline = _Pipeline()
    pipeline.handle = object()
    pipeline._strong_layers = [0]

    assert bo.run_bayesian_optimization(pipeline) == {}


def test_run_bayesian_optimization_returns_empty_without_handle_or_layers(monkeypatch):
    _install_fake_optuna(monkeypatch)

    pipeline = _Pipeline()
    pipeline.handle = None
    pipeline._strong_layers = [0]
    assert bo.run_bayesian_optimization(pipeline) == {}

    pipeline.handle = types.SimpleNamespace(model=object(), tokenizer=object())
    pipeline._strong_layers = []
    assert bo.run_bayesian_optimization(pipeline) == {}


def test_run_bayesian_optimization_returns_empty_when_reference_logits_fail(monkeypatch):
    _install_fake_optuna(monkeypatch)

    class RaisingTokenizer:
        def __call__(self, *_args, **_kwargs):
            raise RuntimeError("tokenization failed")

    pipeline = _Pipeline()
    pipeline.handle = types.SimpleNamespace(
        model=object(),
        tokenizer=RaisingTokenizer(),
    )
    pipeline._strong_layers = [0]

    assert bo.run_bayesian_optimization(pipeline, n_kl_prompts=2) == {}
    assert pipeline.freed == 1
    assert "Failed to collect reference logits" in pipeline.logs[-1]


def test_measure_refusal_rate_restores_chat_template_and_counts_generation_failures():
    class TokenBatch(dict):
        def __init__(self, token_count):
            super().__init__(input_ids=torch.arange(token_count).unsqueeze(0))

    class Tokenizer:
        def __init__(self):
            self.calls = 0

        def __call__(self, prompt, **_kwargs):
            self.calls += 1
            return TokenBatch(3 if "one" in prompt else 4)

        def decode(self, token_ids, **_kwargs):
            if token_ids.numel() == 0:
                return ""
            return "I cannot assist with that request"

    class Model:
        def __init__(self):
            self.calls = 0

        def generate(self, **kwargs):
            self.calls += 1
            if self.calls == 2:
                raise RuntimeError("generation failed")
            input_ids = kwargs["input_ids"]
            continuation = torch.tensor([[99, 100]])
            return torch.cat([input_ids, continuation], dim=1)

    pipeline = _Pipeline()
    pipeline.use_chat_template = False
    pipeline.handle = types.SimpleNamespace(model=Model(), tokenizer=Tokenizer())

    rate = bo._measure_refusal_rate(pipeline, n_prompts=2, max_new_tokens=4)

    assert rate == pytest.approx(0.5)
    assert pipeline.use_chat_template is False
    assert pipeline.freed == 1


def test_measure_kl_divergence_skips_failures_and_missing_reference_logits():
    class TokenBatch(dict):
        def __init__(self):
            super().__init__(input_ids=torch.tensor([[1, 2]]))

    class Tokenizer:
        def __call__(self, prompt, **_kwargs):
            if "bad" in prompt:
                raise RuntimeError("tokenization failed")
            return TokenBatch()

    class Output:
        logits = torch.tensor([[[0.0, 1.0, 2.0]]])

    class Model:
        def __call__(self, **_kwargs):
            return Output()

    pipeline = _Pipeline()
    pipeline.handle = types.SimpleNamespace(model=Model(), tokenizer=Tokenizer())
    reference = [torch.tensor([0.0, 1.0, 2.0])]

    kl = bo._measure_kl_divergence(
        pipeline,
        reference_logits=reference,
        prompts=["good", "bad", "ignored because no reference"],
    )

    assert kl == pytest.approx(0.0)
    assert pipeline.freed == 1


def test_run_bayesian_optimization_pareto_path_enqueues_warm_start_and_restores(monkeypatch):
    layers = [_Layer(), _Layer()]
    originals = [
        layer.self_attn.o_proj.weight.detach().clone()
        for layer in layers
    ] + [
        layer.mlp.down_proj.weight.detach().clone()
        for layer in layers
    ]
    best_params = {
        "attn_max_weight": 0.8,
        "attn_peak_position": 0.0,
        "attn_min_weight": 0.1,
        "attn_spread": 0.6,
        "mlp_max_weight": 0.6,
        "mlp_peak_position": 1.0,
        "mlp_min_weight": 0.2,
        "mlp_spread": 0.6,
        "dir_idx": 1.0,
    }
    worse_params = {**best_params, "dir_idx": 0.0}
    study = _FakeStudy(
        best_trials=[
            _FakeTrial(params=worse_params, values=(0.5, 0.1)),
            _FakeTrial(params=best_params, values=(0.2, 0.3)),
        ]
    )
    _install_fake_optuna(monkeypatch, study)
    monkeypatch.setattr(bo, "_measure_refusal_rate", lambda *_args, **_kwargs: 0.25)
    monkeypatch.setattr(bo, "_measure_kl_divergence", lambda *_args, **_kwargs: 0.1)

    pipeline = _optimization_pipeline(layers)
    pipeline._informed_warm_start = {
        "max_weight": 0.8,
        "peak_position": 0.25,
        "min_weight": 0.02,
        "spread": 0.2,
        "mlp_scale": 0.5,
        "dir_idx": 0.75,
    }

    result = bo.run_bayesian_optimization(
        pipeline,
        n_trials=2,
        n_refusal_prompts=3,
        n_kl_prompts=1,
    )

    assert study.enqueued == [{
        "attn_max_weight": 0.8,
        "attn_peak_position": 0.25,
        "attn_min_weight": 0.02,
        "attn_spread": 0.2,
        "mlp_max_weight": 0.4,
        "mlp_peak_position": 0.25,
        "mlp_min_weight": 0.02,
        "mlp_spread": 0.2,
        "dir_idx": 0.75,
    }]
    assert study.objective_values == [(0.25, 0.1), (0.25, 0.1)]
    assert result == {0: pytest.approx(0.6), 1: pytest.approx(0.7)}
    assert pipeline._bayesian_attn_scale == pytest.approx(0.8)
    assert pipeline._bayesian_mlp_scale == pytest.approx(0.6)
    assert any("Applying interpolated direction" in msg for msg in pipeline.logs)

    restored = [
        layer.self_attn.o_proj.weight
        for layer in layers
    ] + [
        layer.mlp.down_proj.weight
        for layer in layers
    ]
    for live, original in zip(restored, originals, strict=True):
        assert torch.allclose(live, original)

    assert len(pipeline.projections) == 8
    regularizations = sorted({round(call["regularization"], 6) for call in pipeline.projections})
    assert regularizations == [0.2, 0.4, 1.0]
    assert all(call["norm_preserve"] is True for call in pipeline.projections)
    assert all(call["direction"].shape == (2, 1) for call in pipeline.projections)

    expected_direction = torch.tensor([2.0, 1.0])
    expected_direction = expected_direction / expected_direction.norm()
    assert torch.allclose(pipeline.refusal_directions[0], expected_direction)
    assert torch.allclose(pipeline.refusal_directions[1], expected_direction)


def test_run_bayesian_optimization_no_pareto_uses_objective_best_and_restores_after_projection_errors(
    monkeypatch,
):
    layers = [_Layer()]
    original_attn = layers[0].self_attn.o_proj.weight.detach().clone()
    original_mlp = layers[0].mlp.down_proj.weight.detach().clone()
    study = _FakeStudy(best_trials=[])
    _install_fake_optuna(monkeypatch, study)
    monkeypatch.setattr(bo, "_measure_refusal_rate", lambda *_args, **_kwargs: 0.4)
    monkeypatch.setattr(bo, "_measure_kl_divergence", lambda *_args, **_kwargs: 0.2)

    pipeline = _optimization_pipeline(layers)

    def raising_project(*_args, **_kwargs):
        layers[0].self_attn.o_proj.weight.data.add_(5.0)
        layers[0].mlp.down_proj.weight.data.add_(7.0)
        raise RuntimeError("projection failed")

    pipeline._project_out_advanced = raising_project

    result = bo.run_bayesian_optimization(
        pipeline,
        n_trials=1,
        n_refusal_prompts=1,
        n_kl_prompts=1,
    )

    assert study.enqueued == [{
        "attn_max_weight": 0.9,
        "attn_peak_position": 0.0,
        "attn_min_weight": 0.05,
        "attn_spread": 0.3,
        "mlp_max_weight": 0.6,
        "mlp_peak_position": 0.0,
        "mlp_min_weight": 0.05,
        "mlp_spread": 0.3,
        "dir_idx": 0.0,
    }]
    assert study.objective_values == [(0.4, 0.2)]
    assert result == {0: pytest.approx(0.3)}
    assert any("Using best combined score: 0.5000" in msg for msg in pipeline.logs)
    assert pipeline.moe_calls == []
    assert torch.allclose(layers[0].self_attn.o_proj.weight, original_attn)
    assert torch.allclose(layers[0].mlp.down_proj.weight, original_mlp)