from gpu_rent.inventory import (
    gpu_boot_image_score,
    gpu_quota_from_compute,
    match_label,
    pick_boot_image,
    rank_flavors,
    resolve_flavor,
)


class FakeFlavor:
    def __init__(self, flavor_id: str, name: str, extra=None, disabled: bool = False):
        self.id = flavor_id
        self.name = name
        self.extra_specs = extra or {}
        self.is_disabled = disabled
        self.vcpus = 8
        self.ram = 32768


def test_match_4090_24_not_48():
    a = FakeFlavor("1", "GPU 1x RTX 4090 24GB")
    b = FakeFlavor("2", "GPU 1x RTX 4090 48GB")
    assert match_label("4090-24", a)
    assert not match_label("4090-24", b)
    assert match_label("4090-48", b)


def test_rank_skips_disabled():
    flavors = [
        FakeFlavor("d", "RTX 4090 24GB", disabled=True),
        FakeFlavor("ok", "RTX A5000 24GB"),
    ]
    ranked = rank_flavors(flavors, ("4090-24", "a5000"))
    assert [x.id for x in ranked] == ["ok"]
    assert ranked[0].label == "a5000"


def test_gpu_quota_from_compute():
    assert gpu_quota_from_compute({"cores": 10}) is None
    assert gpu_quota_from_compute({"gpu": 0}) == 0
    assert gpu_quota_from_compute({"GPU_limit": 2}) == 2


def test_resolve_flavor_explicit_and_preference():
    flavors = [
        FakeFlavor("d", "RTX 4090 24GB", disabled=True),
        FakeFlavor("a5", "RTX A5000 24GB"),
        FakeFlavor("ok", "GPU 1x RTX 4090 24GB"),
    ]
    picked = resolve_flavor(flavors, ("4090-24", "a5000"))
    assert picked.id == "ok"
    explicit = resolve_flavor(flavors, ("a5000",), explicit="a5")
    assert explicit.id == "a5"


def test_resolve_flavor_prefers_list_over_default_id():
    flavors = [
        FakeFlavor("a5", "RTX A5000 24GB"),
        FakeFlavor("ok", "GPU 1x RTX 4090 24GB"),
    ]
    picked = resolve_flavor(flavors, ("4090-24", "a5000"), default_id="a5", fallback=True)
    assert picked.id == "ok"


def test_resolve_flavor_no_fallback_requires_default():
    flavors = [FakeFlavor("ok", "GPU 1x RTX 4090 24GB")]
    try:
        resolve_flavor(flavors, ("4090-24",), fallback=False)
        raise AssertionError("expected ValueError")
    except ValueError:
        pass
    picked = resolve_flavor(flavors, ("a5000",), default_id="ok", fallback=False)
    assert picked.id == "ok"


def test_pick_boot_image_prefers_24_580_without_docker():
    class Img:
        def __init__(self, name):
            self.name = name

    images = [
        Img("Ubuntu 24.04 LTS 64-bit GPU Driver 580 Docker"),
        Img("Ubuntu 24.04 LTS 64-bit GPU Driver 535"),
        Img("Ubuntu 24.04 LTS 64-bit GPU Driver 580"),
        Img("Data Science VM (Ubuntu 22.04 LTS 64-bit)"),
    ]
    picked = pick_boot_image(images)
    assert picked.name == "Ubuntu 24.04 LTS 64-bit GPU Driver 580"
    assert gpu_boot_image_score(picked.name) > gpu_boot_image_score(images[0].name)