from __future__ import annotations

import pytest


def test_build_llm_config_dict_maps_tp_and_bundles():
    from argus.service.serve_llm_config import build_llm_config_dict
    from argus.service.serving_spec import ResolvedServingSpec

    resolved = ResolvedServingSpec(
        user_id="alice",
        model_id_suffix="qwen-0.5b",
        model_id_prefix="alice-202601061235",
        model_id="alice-202601061235-qwen-0.5b",
        model_source="/private/hf/x",
        num_replicas=2,
        gpus_per_replica=4,
        engine_kwargs={"gpu_memory_utilization": 0.9},
    )

    cfg = build_llm_config_dict(
        resolved,
        accelerator_type="H20",
        runtime_env_env_vars={"HF_ENDPOINT": "https://hf-mirror.com"},
        cpu_per_gpu=2.0,
    )
    assert cfg["model_loading_config"]["model_id"] == "alice-202601061235-qwen-0.5b"
    assert cfg["model_loading_config"]["model_source"] == "/private/hf/x"
    assert cfg["accelerator_type"] == "H20"
    assert cfg["deployment_config"]["num_replicas"] == 2

    # gpus_per_replica -> tensor_parallel_size
    assert cfg["engine_kwargs"]["tensor_parallel_size"] == 4
    assert cfg["engine_kwargs"]["gpu_memory_utilization"] == 0.9

    # resources_per_bundle reserves the full TP GPU set for each replica.
    bundle = cfg["resources_per_bundle"]
    assert bundle["GPU"] == 4.0
    assert bundle["CPU"] == 8.0


def test_build_llm_config_dict_injects_hf_offline_defaults():
    from argus.service.serve_llm_config import build_llm_config_dict
    from argus.service.serving_spec import ResolvedServingSpec

    resolved = ResolvedServingSpec(
        user_id="alice",
        model_id_suffix="x",
        model_id_prefix="alice-202601061235",
        model_id="alice-202601061235-x",
        model_source="/private/users/alice/models/x",
        num_replicas=1,
        gpus_per_replica=1,
        engine_kwargs=None,
    )
    cfg = build_llm_config_dict(resolved, accelerator_type="H20", runtime_env_env_vars={})
    env = cfg["runtime_env"]["env_vars"]
    assert env["HF_HUB_OFFLINE"] == "1"
    assert env["HF_HOME"] == "/private/hf"
    assert env["HUGGINGFACE_HUB_CACHE"].startswith("/private/hf/")


def test_build_llm_config_dict_requires_accelerator_type():
    from argus.service.serve_llm_config import build_llm_config_dict
    from argus.service.serving_spec import ResolvedServingSpec

    resolved = ResolvedServingSpec(
        user_id="alice",
        model_id_suffix="x",
        model_id_prefix="alice-202601061235",
        model_id="alice-202601061235-x",
        model_source="/private/hf/x",
        num_replicas=1,
        gpus_per_replica=1,
        engine_kwargs=None,
    )
    with pytest.raises(ValueError, match="accelerator_type is required"):
        build_llm_config_dict(resolved, accelerator_type="", runtime_env_env_vars={})
