"""Tests for deterministic compression optimizer selection."""

from __future__ import annotations

import numpy as np

from model_osmosis.format.osmo_format import ArchitectureSpec, OsmoManifest, OsmoModel
from model_osmosis.optimizer.compression_engine import (
    resolve_compression_engine_mode,
    select_compression_selection,
)


def _make_osmo(seed: int = 0, width: int = 64) -> OsmoModel:
    rng = np.random.default_rng(seed)
    params = {
        "layer0.weight": rng.standard_normal((width, width), dtype=np.float32),
        "layer0.bias": rng.standard_normal((width,), dtype=np.float32),
    }
    return OsmoModel(
        manifest=OsmoManifest(name="compression-test", parameter_count=sum(arr.size for arr in params.values())),
        architecture=ArchitectureSpec(layers=[], hidden_size=width, num_layers=1, architecture_type="mlp"),
        parameters=params,
    )


def test_resolve_compression_engine_global_off(monkeypatch):
    monkeypatch.setenv("OSMOSIS_OPTIMIZER", "off")
    assert resolve_compression_engine_mode("auto") == "off"


def test_select_compression_selection_lossless_is_deterministic(monkeypatch, tmp_path):
    monkeypatch.setenv("OSMOSIS_OPTIMIZER_HOME", str(tmp_path / "optimizer"))
    osmo = _make_osmo(seed=11, width=48)

    first = select_compression_selection(
        osmo_model=osmo,
        compression_profile="lossless-max",
        compression_engine_mode="auto",
    )
    second = select_compression_selection(
        osmo_model=osmo,
        compression_profile="lossless-max",
        compression_engine_mode="auto",
    )

    assert first.profile_id == second.profile_id
    assert first.pack_version == second.pack_version
    assert first.selector_version == second.selector_version
    assert tuple(first.zstd_levels) == tuple(second.zstd_levels)


def test_select_compression_selection_legacy_passthrough(monkeypatch, tmp_path):
    monkeypatch.setenv("OSMOSIS_OPTIMIZER_HOME", str(tmp_path / "optimizer"))
    osmo = _make_osmo(seed=5, width=16)
    selection = select_compression_selection(
        osmo_model=osmo,
        compression_profile="legacy",
        compression_engine_mode="auto",
    )
    assert selection.compression_profile == "legacy"
    assert selection.profile_id == "legacy_passthrough"
    assert tuple(selection.zstd_levels) == ()
