"""Tests for adaptive cross-merge recipe selector."""

from __future__ import annotations

import numpy as np

from model_osmosis.merging.cross_merge_engine import select_cross_merge_recipe
from model_osmosis.optimizer.self_optimize import get_optimizer_status


def _classic_models():
    return [object(), object()]


def test_cross_selector_is_deterministic(monkeypatch, tmp_path):
    monkeypatch.setenv("OSMOSIS_OPTIMIZER_HOME", str(tmp_path / "optimizer"))
    sample = np.random.default_rng(0).standard_normal((64, 16)).astype(np.float32)

    first = select_cross_merge_recipe(
        models=_classic_models(),
        sample_data=sample,
        weights=[0.5, 0.5],
        seed=42,
        strategy_engine_mode="auto",
        cross_opt_stage="shadow",
        source_refs=["pt:a.pt", "pt:b.pt"],
        epochs=8,
    )
    second = select_cross_merge_recipe(
        models=_classic_models(),
        sample_data=sample,
        weights=[0.5, 0.5],
        seed=42,
        strategy_engine_mode="auto",
        cross_opt_stage="shadow",
        source_refs=["pt:a.pt", "pt:b.pt"],
        epochs=8,
    )

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


def test_cross_selector_can_promote_after_canary_passes(monkeypatch, tmp_path):
    monkeypatch.setenv("OSMOSIS_OPTIMIZER_HOME", str(tmp_path / "optimizer"))
    sample = np.random.default_rng(1).standard_normal((32, 8)).astype(np.float32)

    last = None
    for _ in range(6):
        last = select_cross_merge_recipe(
            models=_classic_models(),
            sample_data=sample,
            seed=7,
            strategy_engine_mode="auto",
            cross_opt_stage="active",
            source_refs=["pt:a.pt", "pt:b.pt"],
            epochs=6,
        )

    assert last is not None
    status = get_optimizer_status()
    classic = status["subsystems"]["cross_merge_classic"]
    assert classic["canary_attempts"] >= 1
    assert classic["canary_passes"] >= 1
    assert classic["canary_status"] in {"canary", "promoted"}


def test_cross_selector_rolls_back_on_quality_regression(monkeypatch, tmp_path):
    monkeypatch.setenv("OSMOSIS_OPTIMIZER_HOME", str(tmp_path / "optimizer"))
    sample = np.random.default_rng(2).standard_normal((16, 4)).astype(np.float32)

    import model_osmosis.merging.cross_merge_engine as eng

    def _forced_scores(*, recipes, context, subsystem_metadata):
        out = {}
        for rec in recipes:
            rid = str(rec["recipe_id"])
            out[rid] = 10.0 if "fast" in rid else 1.0
        return out

    def _forced_quality(recipe, context):
        rid = str(recipe.get("recipe_id", ""))
        return 0.0 if "fast" in rid else 1.0

    monkeypatch.setattr(eng, "_score_recipes", _forced_scores)
    monkeypatch.setattr(eng, "_proxy_quality", _forced_quality)

    result = select_cross_merge_recipe(
        models=_classic_models(),
        sample_data=sample,
        seed=11,
        strategy_engine_mode="auto",
        cross_opt_stage="active",
        source_refs=["pt:a.pt", "pt:b.pt"],
        epochs=4,
    )
    assert result.canary_status == "rollback"
    assert result.rollback_reason == "quality_regression_blocked"
