"""Tests for smart merge engine recipe selection."""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np

from model_osmosis.merging.smart_merge_engine import (
    DEFAULT_PACK_VERSION,
    collect_tensor_groups_from_state_dicts,
    resolve_strategy_engine_mode,
    select_recipe_for_tensor_groups,
)


def _make_groups(seed: int = 42):
    rng = np.random.default_rng(seed)
    groups = []
    for _ in range(6):
        a = rng.standard_normal((8, 8)).astype(np.float32)
        b = (a + 0.02 * rng.standard_normal((8, 8))).astype(np.float32)
        c = (a + 0.03 * rng.standard_normal((8, 8))).astype(np.float32)
        groups.append([a, b, c])
    return groups


def test_collect_tensor_groups_from_state_dicts_deterministic():
    states = [
        {"w1": np.ones((2, 2), dtype=np.float32), "w2": np.ones((3, 3), dtype=np.float32)},
        {"w1": np.ones((2, 2), dtype=np.float32) * 2.0, "w2": np.ones((3, 3), dtype=np.float32) * 3.0},
    ]
    groups_a = collect_tensor_groups_from_state_dicts(states, seed=7)
    groups_b = collect_tensor_groups_from_state_dicts(states, seed=7)
    assert len(groups_a) == len(groups_b)
    for ga, gb in zip(groups_a, groups_b):
        for ta, tb in zip(ga, gb):
            np.testing.assert_array_equal(ta, tb)


def test_smart_selection_is_deterministic(monkeypatch, tmp_path):
    monkeypatch.setenv("OSMOSIS_STRATEGY_HOME", str(tmp_path / "smart"))
    groups = _make_groups(11)

    first = select_recipe_for_tensor_groups(
        tensor_groups=groups,
        model_count=3,
        seed=42,
        strategy_engine_mode="auto",
        context_label="merge",
        source_refs=["pt:a.pt", "pt:b.pt", "pt:c.pt"],
    )
    second = select_recipe_for_tensor_groups(
        tensor_groups=groups,
        model_count=3,
        seed=42,
        strategy_engine_mode="auto",
        context_label="merge",
        source_refs=["pt:a.pt", "pt:b.pt", "pt:c.pt"],
    )

    assert first.recipe_id == second.recipe_id
    assert first.strategy_name == second.strategy_name
    assert first.recipe_pack_version == second.recipe_pack_version
    assert first.selector_version == second.selector_version


def test_strategy_engine_off_forces_fallback(monkeypatch):
    monkeypatch.setenv("OSMOSIS_STRATEGY_ENGINE", "off")
    assert resolve_strategy_engine_mode("auto") == "off"

    selection = select_recipe_for_tensor_groups(
        tensor_groups=_make_groups(5),
        model_count=3,
        seed=42,
        strategy_engine_mode="auto",
    )
    assert selection.strategy_name == "weighted_average"
    assert selection.fallback_reason == "strategy_engine_off"


def test_global_optimizer_off_overrides_strategy(monkeypatch):
    monkeypatch.setenv("OSMOSIS_OPTIMIZER", "off")
    assert resolve_strategy_engine_mode("auto") == "off"


def test_invalid_pack_checksum_falls_back_to_builtin(monkeypatch, tmp_path):
    home = tmp_path / "smart"
    home.mkdir(parents=True, exist_ok=True)
    pack_path = home / "recipe_pack.json"
    pack_path.write_text(
        json.dumps(
            {
                "pack_version": "broken-pack",
                "selector_version": "x",
                "recipes": [
                    {
                        "recipe_id": "broken",
                        "strategy_name": "weighted_average",
                        "strategy_kwargs": {},
                    }
                ],
                "checksum_sha256": "deadbeef",
            }
        ),
        encoding="utf-8",
    )
    monkeypatch.setenv("OSMOSIS_STRATEGY_HOME", str(home))

    selection = select_recipe_for_tensor_groups(
        tensor_groups=_make_groups(8),
        model_count=3,
        seed=9,
        strategy_engine_mode="auto",
    )
    assert selection.recipe_pack_version == DEFAULT_PACK_VERSION
    assert "fallback_builtin" in selection.fallback_reason or selection.fallback_reason == "builtin_pack"


def test_selection_updates_local_state(monkeypatch, tmp_path):
    home = tmp_path / "smart"
    monkeypatch.setenv("OSMOSIS_STRATEGY_HOME", str(home))
    _ = select_recipe_for_tensor_groups(
        tensor_groups=_make_groups(14),
        model_count=3,
        seed=42,
        strategy_engine_mode="auto",
        source_refs=["pt:a.pt", "pt:b.pt", "pt:c.pt"],
    )
    state_path = home / "local_state.json"
    assert state_path.exists()
    state = json.loads(state_path.read_text(encoding="utf-8"))
    assert state["schema_version"] == 1
    assert isinstance(state["recipe_priors"], dict)
