import builtins
import io
from types import SimpleNamespace

import pytest
import torch

from comfy.cli_args import args as cli_args

if not torch.cuda.is_available():
    cli_args.cpu = True

import comfy.model_management as model_management  # noqa: E402


MIXED_TOPOLOGY = {
    "card1": {
        "realpath": "/sys/devices/pci0000:00/0000:00:01.1/0000:01:00.0/0000:02:00.0/0000:03:00.0",
        "vram": 17163091968,
        "gtt": 16359088128,
    },
    "card2": {
        "realpath": "/sys/devices/pci0000:00/0000:00:01.2/0000:04:00.0/0000:05:00.0/0000:06:00.0",
        "vram": 17163091968,
        "gtt": 16359088128,
    },
    "card3": {
        "realpath": "/sys/devices/pci0000:00/0000:00:08.1/0000:1d:00.0",
        "vram": 536870912,
        "gtt": 16359088128,
    },
}


def install_fake_drm(monkeypatch, cards):
    files = {}
    realpaths = {}
    for card, values in cards.items():
        device_dir = f"/sys/class/drm/{card}/device"
        files[f"{device_dir}/vendor"] = values.get("vendor", "0x1002")
        files[f"{device_dir}/mem_info_vram_total"] = str(values["vram"])
        files[f"{device_dir}/mem_info_gtt_total"] = str(values["gtt"])
        realpaths[device_dir] = values["realpath"]

    monkeypatch.setattr(
        model_management.os,
        "listdir",
        lambda path: [*cards, "card1-DP-1"] if path == "/sys/class/drm" else [],
    )
    monkeypatch.setattr(model_management.os.path, "exists", lambda path: path in files)
    monkeypatch.setattr(model_management.os.path, "realpath", lambda path: realpaths.get(path, path))
    monkeypatch.setattr(
        builtins,
        "open",
        lambda path, *args, **kwargs: io.StringIO(files[str(path)]),
    )


def properties(bus, *, integrated, domain=0, device=0):
    return SimpleNamespace(
        pci_domain_id=domain,
        pci_bus_id=bus,
        pci_device_id=device,
        is_integrated=int(integrated),
    )


@pytest.fixture(autouse=True)
def gpu_mode(monkeypatch):
    monkeypatch.setattr(model_management, "cpu_state", model_management.CPUState.GPU)
    monkeypatch.setattr(model_management, "is_amd", lambda: True)
    monkeypatch.setattr(model_management, "is_nvidia", lambda: False)


def test_mixed_topology_maps_each_device_by_canonical_bdf(monkeypatch):
    install_fake_drm(monkeypatch, MIXED_TOPOLOGY)
    device_properties = {
        0: properties(3, integrated=False),
        1: properties(6, integrated=False),
        2: properties(29, integrated=True),
    }
    monkeypatch.setattr(
        model_management.torch.cuda,
        "get_device_properties",
        lambda device: device_properties[device.index],
    )

    assert model_management._amd_vram_gtt_totals(torch.device("cuda", 0)) == (
        17163091968,
        16359088128,
    )
    assert model_management._amd_vram_gtt_totals(torch.device("cuda", 1)) == (
        17163091968,
        16359088128,
    )
    assert model_management._amd_vram_gtt_totals(torch.device("cuda", 2)) == (
        536870912,
        16359088128,
    )
    assert model_management.is_integrated_gpu(torch.device("cuda", 0)) is False
    assert model_management.is_integrated_gpu(torch.device("cuda", 1)) is False
    assert model_management.is_integrated_gpu(torch.device("cuda", 2)) is True
    assert model_management.integrated_gpu_is_shared_heavy(torch.device("cuda", 0)) is False
    assert model_management.integrated_gpu_is_shared_heavy(torch.device("cuda", 1)) is False
    assert model_management.integrated_gpu_is_shared_heavy(torch.device("cuda", 2)) is True


def test_dedicated_heavy_integrated_gpu_stays_non_shared(monkeypatch):
    install_fake_drm(monkeypatch, {"card1": MIXED_TOPOLOGY["card1"]})
    monkeypatch.setattr(
        model_management.torch.cuda,
        "get_device_properties",
        lambda device: properties(3, integrated=True),
    )

    device = torch.device("cuda", 0)
    assert model_management.is_integrated_gpu(device) is True
    assert model_management.integrated_gpu_is_shared_heavy(device) is False


@pytest.mark.parametrize(
    ("vram", "gtt", "expected"),
    [
        (1024, 1023, False),
        (1024, 1024, True),
        (1024, 1025, True),
        (0, 1024, True),
    ],
)
def test_shared_heavy_threshold(monkeypatch, vram, gtt, expected):
    card = {**MIXED_TOPOLOGY["card3"], "vram": vram, "gtt": gtt}
    install_fake_drm(monkeypatch, {"card3": card})
    monkeypatch.setattr(
        model_management.torch.cuda,
        "get_device_properties",
        lambda device: properties(29, integrated=True),
    )

    assert model_management.integrated_gpu_is_shared_heavy(torch.device("cuda", 0)) is expected


def test_full_domain_bus_device_bdf_is_matched(monkeypatch):
    card = {
        **MIXED_TOPOLOGY["card3"],
        "realpath": "/sys/devices/pci1234:00/1234:ab:1f.0",
    }
    install_fake_drm(monkeypatch, {"card3": card})
    monkeypatch.setattr(
        model_management.torch.cuda,
        "get_device_properties",
        lambda device: properties(0xAB, integrated=True, domain=0x1234, device=0x1F),
    )

    assert model_management._amd_vram_gtt_totals(torch.device("cuda", 0)) == (
        536870912,
        16359088128,
    )


def test_multiple_candidates_without_bdf_match_do_not_fall_back(monkeypatch):
    install_fake_drm(monkeypatch, MIXED_TOPOLOGY)
    monkeypatch.setattr(
        model_management.torch.cuda,
        "get_device_properties",
        lambda device: properties(127, integrated=True),
    )

    assert model_management._amd_vram_gtt_totals(torch.device("cuda", 0)) is None


def test_non_amd_drm_candidate_is_ignored(monkeypatch):
    cards = {
        "card1": MIXED_TOPOLOGY["card1"],
        "card2": {**MIXED_TOPOLOGY["card2"], "vendor": "0x10de"},
    }
    install_fake_drm(monkeypatch, cards)
    monkeypatch.setattr(
        model_management.torch.cuda,
        "get_device_properties",
        lambda device: properties(6, integrated=False),
    )

    assert model_management._amd_vram_gtt_totals(torch.device("cuda", 0)) == (
        17163091968,
        16359088128,
    )


def test_single_candidate_can_fall_back_when_pci_properties_are_unavailable(monkeypatch):
    install_fake_drm(monkeypatch, {"card3": MIXED_TOPOLOGY["card3"]})
    monkeypatch.setattr(
        model_management.torch.cuda,
        "get_device_properties",
        lambda device: (_ for _ in ()).throw(RuntimeError("properties unavailable")),
    )

    assert model_management._amd_vram_gtt_totals(torch.device("cuda", 0)) == (
        536870912,
        16359088128,
    )


def test_unavailable_totals_default_integrated_gpu_to_shared(monkeypatch):
    install_fake_drm(monkeypatch, {})

    assert model_management._amd_vram_gtt_totals(torch.device("cuda", 0)) is None
    assert model_management.integrated_gpu_is_shared_heavy(torch.device("cuda", 0)) is True


def test_non_integer_totals_default_integrated_gpu_to_shared(monkeypatch):
    card = {**MIXED_TOPOLOGY["card3"], "vram": "invalid", "gtt": 16359088128}
    install_fake_drm(monkeypatch, {"card3": card})

    assert model_management._amd_vram_gtt_totals(torch.device("cuda", 0)) is None
    assert model_management.integrated_gpu_is_shared_heavy(torch.device("cuda", 0)) is True


def test_cpu_state_never_reports_integrated_gpu(monkeypatch):
    monkeypatch.setattr(model_management, "cpu_state", model_management.CPUState.CPU)
    monkeypatch.setattr(
        model_management.torch.cuda,
        "get_device_properties",
        lambda device: properties(29, integrated=True),
    )

    assert model_management.is_integrated_gpu(torch.device("cuda", 0)) is False


def test_non_amd_backend_has_no_amd_totals(monkeypatch):
    install_fake_drm(monkeypatch, MIXED_TOPOLOGY)
    monkeypatch.setattr(model_management, "is_amd", lambda: False)

    assert model_management._amd_vram_gtt_totals(torch.device("cuda", 0)) is None