"""Lightweight expand bridge routing smoke tests."""

from __future__ import annotations

from model_osmosis.cli import _auto_detect_expand_scheme, _expand_route_next_command, _parse_expand_ref, ExpandModelRef


def test_auto_detect_expand_scheme_by_suffix(tmp_path):
    osmoimg = tmp_path / "model.osmoimg"
    osmoimg.write_bytes(b"")
    assert _auto_detect_expand_scheme(str(osmoimg)) == "osmoimg"

    osmo = tmp_path / "model.osmo"
    osmo.write_bytes(b"")
    assert _auto_detect_expand_scheme(str(osmo)) == "osmo"

    onnx = tmp_path / "model.onnx"
    onnx.write_bytes(b"")
    assert _auto_detect_expand_scheme(str(onnx)) == "onnx"

    tf_h5 = tmp_path / "model.h5"
    tf_h5.write_bytes(b"")
    assert _auto_detect_expand_scheme(str(tf_h5)) == "tf"

    jax_npz = tmp_path / "my_jax_model.npz"
    jax_npz.write_bytes(b"")
    assert _auto_detect_expand_scheme(str(jax_npz)) == "jax"

    pt = tmp_path / "model.pt"
    pt.write_bytes(b"")
    assert _auto_detect_expand_scheme(str(pt)) == "pt"


def test_parse_expand_ref_supports_extended_schemes():
    assert _parse_expand_ref("vision:timm/resnet18") == ("vision", "timm/resnet18")
    assert _parse_expand_ref("audio:openai/whisper-tiny") == ("audio", "openai/whisper-tiny")
    assert _parse_expand_ref("diffusion:runwayml/stable-diffusion-v1-5") == (
        "diffusion",
        "runwayml/stable-diffusion-v1-5",
    )
    assert _parse_expand_ref("mlx:mlx-community/phi-2-4bit") == ("mlx", "mlx-community/phi-2-4bit")
    assert _parse_expand_ref("lora:base::adapter.safetensors") == ("lora", "base::adapter.safetensors")


def test_expand_route_next_command_is_deterministic():
    ref = ExpandModelRef(raw="onnx:./model.onnx", scheme="onnx", value="./model.onnx", model=None)
    assert _expand_route_next_command(ref) == "osmosis convert onnx:./model.onnx --to pytorch -o expanded_source.pt"

    ref = ExpandModelRef(raw="osmoimg:./model.osmoimg", scheme="osmoimg", value="./model.osmoimg", model=None)
    assert _expand_route_next_command(ref) == "osmosis export ./model.osmoimg --format pytorch -o expanded_source.pt"

    ref = ExpandModelRef(raw="vision:resnet18", scheme="vision", value="resnet18", model=None)
    assert _expand_route_next_command(ref) == "osmosis expand vision:resnet18 --fallback-mode distill"
