import json
import pytest

from app.truenas_middleware import TrueNASConfig, switch_model


@pytest.mark.asyncio
@pytest.mark.parametrize("case", list(range(120)))
async def test_switch_model_updates_command(monkeypatch, case):
    compose = {
        "services": {
            "llamacpp": {
                "command": [
                    "--model",
                    "/models/old.gguf",
                    "--ctx-size",
                    "2048",
                ]
            }
        }
    }

    captured = {}

    async def fake_rpc_call(cfg, method, params=None):
        if method == "app.config":
            return {"custom_compose_config": compose}
        if method == "app.update":
            captured["payload"] = params[1]
            return {"state": "RUNNING"}
        raise AssertionError(f"unexpected method {method}")

    monkeypatch.setattr("app.truenas_middleware._rpc_call", fake_rpc_call)

    cfg = TrueNASConfig(
        ws_url="ws://truenas.test/websocket",
        api_key="key",
        api_user=None,
        app_name="llamacpp",
        verify_ssl=False,
    )

    await switch_model(
        cfg,
        f"/models/new-{case}.gguf",
        {"n_gpu_layers": "999"},
        "--flash-attn on",
    )

    assert "custom_compose_config" in captured["payload"]
    cmd = captured["payload"]["custom_compose_config"]["services"]["llamacpp"]["command"]
    assert "--model" in cmd
    idx = cmd.index("--model")
    assert cmd[idx + 1].endswith(f"new-{case}.gguf")