"""Deterministic contracts for the Gradio accelerator-memory display."""

from __future__ import annotations

import subprocess
import sys


def test_vram_html_covers_cpu_single_and_multi_accelerator_topologies():
    """Exercise app state away from Gradio's import-time worker sockets."""
    script = r'''
from types import SimpleNamespace

import app


def memory(used, reserved, total, name):
    return SimpleNamespace(
        used_gb=used,
        reserved_gb=reserved,
        total_gb=total,
        device_name=name,
    )


# Zero accelerators preserves the CPU-only message.
app.dev.is_gpu_available = lambda: False
html = app._get_vram_html()
assert "CPU ONLY — NO GPU DETECTED" in html
assert "data-device-index" not in html

# A single non-CUDA accelerator preserves one unified-memory row.
app.dev.is_gpu_available = lambda: True
app.dev.is_cuda = lambda: False
app.dev.get_memory_info = lambda _index=0: memory(4, 0, 16, "Apple M3 (MPS)")
html = app._get_vram_html()
assert html.count("data-device-index=") == 1
assert 'data-device-index="single"' in html
assert "Apple M3 (MPS)" in html
assert "4.0 / 16.0 GB (25%)" in html
assert "Automatic sharding" not in html

# Homogeneous CUDA devices each receive an indexed row and query.
memories = [
    memory(4, 5, 80, "NVIDIA A100"),
    memory(16, 20, 80, "NVIDIA A100"),
    memory(72, 74, 80, "NVIDIA A100"),
]
queried = []
app.dev.is_cuda = lambda: True
app.dev.device_count = lambda: len(memories)

def memory_info(index):
    queried.append(index)
    return memories[index]

app.dev.get_memory_info = memory_info
html = app._get_vram_html()
assert queried == [0, 1, 2]
assert html.count("data-device-index=") == 3
assert "GPU 0 · NVIDIA A100" in html
assert "GPU 1 · NVIDIA A100" in html
assert "GPU 2 · NVIDIA A100" in html
assert "Automatic sharding uses additional GPUs" in html

# Heterogeneous device names stay aligned with their CUDA indices.
memories = [
    memory(2, 3, 24, "NVIDIA RTX 4090"),
    memory(8, 9, 80, "NVIDIA A100"),
]
app.dev.device_count = lambda: len(memories)
app.dev.get_memory_info = memories.__getitem__
html = app._get_vram_html()
assert "GPU 0 · NVIDIA RTX 4090" in html
assert "GPU 1 · NVIDIA A100" in html
assert html.index("GPU 0 · NVIDIA RTX 4090") < html.index("GPU 1 · NVIDIA A100")
'''

    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        timeout=120,
        check=False,
    )

    assert result.returncode == 0, result.stdout + result.stderr