"""Tests for obliteratus.blend — complementary abliteration blending."""

from __future__ import annotations

import json
from pathlib import Path

import torch
from safetensors.torch import load_file, save_file


def _make_model(tmpdir: Path, value: float, n_tensors: int = 3):
    """Create a minimal model with all weights set to a constant value."""
    tensors = {f"layer.{i}.weight": torch.full((4, 4), value) for i in range(n_tensors)}
    shard = "model-00001-of-00001.safetensors"
    save_file(tensors, str(tmpdir / shard))
    index = {"metadata": {}, "weight_map": {k: shard for k in tensors}}
    (tmpdir / "model.safetensors.index.json").write_text(json.dumps(index))
    (tmpdir / "config.json").write_text(json.dumps({"model_type": "test"}))
    return tensors


class TestBlendModels:
    def test_lerp_blend(self, tmp_path):
        from obliteratus.blend import blend_models

        a_dir = tmp_path / "model_a"
        b_dir = tmp_path / "model_b"
        out_dir = tmp_path / "blended"
        a_dir.mkdir()
        b_dir.mkdir()

        _make_model(a_dir, 0.0)
        _make_model(b_dir, 10.0)

        result = blend_models(str(a_dir), str(b_dir), str(out_dir), alpha=0.6)

        assert result["blended_tensors"] == 3
        assert result["alpha"] == 0.6

        merged = load_file(str(out_dir / "model-00001-of-00001.safetensors"))
        # 0.6 * 10.0 + 0.4 * 0.0 = 6.0
        assert torch.allclose(merged["layer.0.weight"], torch.full((4, 4), 6.0))

    def test_alpha_zero_is_pure_a(self, tmp_path):
        from obliteratus.blend import blend_models

        a_dir = tmp_path / "a"
        b_dir = tmp_path / "b"
        out = tmp_path / "out"
        a_dir.mkdir()
        b_dir.mkdir()

        _make_model(a_dir, 42.0)
        _make_model(b_dir, 99.0)

        blend_models(str(a_dir), str(b_dir), str(out), alpha=0.0)
        merged = load_file(str(out / "model-00001-of-00001.safetensors"))
        assert torch.allclose(merged["layer.0.weight"], torch.full((4, 4), 42.0))

    def test_alpha_one_is_pure_b(self, tmp_path):
        from obliteratus.blend import blend_models

        a_dir = tmp_path / "a"
        b_dir = tmp_path / "b"
        out = tmp_path / "out"
        a_dir.mkdir()
        b_dir.mkdir()

        _make_model(a_dir, 42.0)
        _make_model(b_dir, 99.0)

        blend_models(str(a_dir), str(b_dir), str(out), alpha=1.0)
        merged = load_file(str(out / "model-00001-of-00001.safetensors"))
        assert torch.allclose(merged["layer.0.weight"], torch.full((4, 4), 99.0))

    def test_writes_metadata(self, tmp_path):
        from obliteratus.blend import blend_models

        a_dir = tmp_path / "a"
        b_dir = tmp_path / "b"
        out = tmp_path / "out"
        a_dir.mkdir()
        b_dir.mkdir()

        _make_model(a_dir, 1.0)
        _make_model(b_dir, 2.0)

        blend_models(str(a_dir), str(b_dir), str(out), alpha=0.5)

        meta = json.loads((out / "blend_metadata.json").read_text())
        assert meta["alpha"] == 0.5
        assert meta["blend_method"] == "lerp"

    def test_copies_config(self, tmp_path):
        from obliteratus.blend import blend_models

        a_dir = tmp_path / "a"
        b_dir = tmp_path / "b"
        out = tmp_path / "out"
        a_dir.mkdir()
        b_dir.mkdir()

        _make_model(a_dir, 1.0)
        _make_model(b_dir, 2.0)

        blend_models(str(a_dir), str(b_dir), str(out))
        assert (out / "config.json").exists()


class TestBlendSearch:
    def test_creates_multiple_blends(self, tmp_path):
        from obliteratus.blend import blend_search

        a_dir = tmp_path / "a"
        b_dir = tmp_path / "b"
        out = tmp_path / "search"
        a_dir.mkdir()
        b_dir.mkdir()

        _make_model(a_dir, 0.0)
        _make_model(b_dir, 10.0)

        results = blend_search(str(a_dir), str(b_dir), str(out), alphas=[0.3, 0.7])

        assert len(results) == 2
        assert (out / "blend_30").exists()
        assert (out / "blend_70").exists()


class TestCLIDispatch:
    def test_blend_dispatch(self, tmp_path):
        from obliteratus.cli import main as cli_main

        a_dir = tmp_path / "a"
        b_dir = tmp_path / "b"
        out = tmp_path / "out"
        a_dir.mkdir()
        b_dir.mkdir()

        _make_model(a_dir, 1.0)
        _make_model(b_dir, 2.0)

        cli_main(["blend",
                   "--model-a", str(a_dir),
                   "--model-b", str(b_dir),
                   "--alpha", "0.6",
                   "--output", str(out)])

        assert (out / "blend_metadata.json").exists()