"""设备解析：cuda / mps / cpu 三后端的选择、AMP 默认值与显存统计。"""

import pytest

torch = pytest.importorskip("torch")

from utils.device import (  # noqa: E402
    BACKEND_PRIORITY,
    DeviceNotAvailableError,
    PeakMemory,
    amp_supported,
    available_backends,
    default_amp,
    describe_device,
    memory_stats,
    pin_memory_supported,
    resolve_amp,
    resolve_device,
    synchronize,
)


def test_available_backends_is_ordered_and_always_has_cpu():
    backends = available_backends()
    assert backends[-1] == "cpu"
    assert backends == [b for b in BACKEND_PRIORITY if b in backends]
    assert len(set(backends)) == len(backends)


def test_auto_picks_the_highest_priority_available_backend():
    device = resolve_device("auto")
    assert device.type == available_backends()[0]
    assert resolve_device("") == device  # 空串等同 auto


def test_cpu_is_always_resolvable():
    assert resolve_device("cpu").type == "cpu"


def test_explicit_backend_matches_hardware_or_raises():
    """显式请求不可用的后端必须报错，不能悄悄退回 CPU——那会让训练慢几十倍还不吭声。"""
    for backend in ("cuda", "mps"):
        if backend in available_backends():
            assert resolve_device(backend).type == backend
        else:
            with pytest.raises(DeviceNotAvailableError, match=backend):
                resolve_device(backend)


def test_unknown_device_raises():
    with pytest.raises(DeviceNotAvailableError, match="未知设备"):
        resolve_device("tpu")


def test_amp_defaults_follow_the_backend():
    # CUDA 有 tensor core，默认开；MPS 上 autocast 实测更慢，默认关；CPU 不支持
    assert default_amp("cuda") is True
    assert default_amp("mps") is False
    assert default_amp("cpu") is False
    assert amp_supported("cuda") and amp_supported("mps")
    assert not amp_supported("cpu")


def test_resolve_amp_handles_auto_and_explicit():
    assert resolve_amp("auto", "cuda") is True
    assert resolve_amp("auto", "mps") is False
    assert resolve_amp(None, "mps") is False
    assert resolve_amp(True, "mps") is True  # 显式打开，MPS 支持 autocast
    assert resolve_amp(True, "cpu") is False  # CPU 不支持，显式打开也无效
    assert resolve_amp(False, "cuda") is False


def test_pin_memory_only_for_cuda():
    assert pin_memory_supported("cuda")
    assert not pin_memory_supported("mps")
    assert not pin_memory_supported("cpu")


def test_describe_device_is_readable():
    assert describe_device("cpu") == "CPU"
    assert "MPS" in describe_device("mps")


def test_memory_stats_shape_per_backend():
    assert memory_stats("cpu") == {}
    stats = memory_stats(resolve_device("auto"))
    device_type = resolve_device("auto").type
    if device_type == "cuda":
        assert {"allocated_mb", "reserved_mb", "peak_allocated_mb"} <= set(stats)
    elif device_type == "mps":
        assert {"allocated_mb", "driver_allocated_mb"} <= set(stats)


def test_peak_memory_tracks_growth_on_accelerators():
    device = resolve_device("auto")
    peak = PeakMemory(device)

    if device.type == "cpu":
        assert peak.peak_mb == 0.0 and peak.reserved_mb == 0.0
        return

    # MPS 没有峰值 API，靠采样，所以峰值是下界（exact=False）
    assert peak.exact == (device.type != "mps")
    buffer = torch.empty(8 * 1024 * 1024 // 4, device=device)  # 约 8 MB
    peak.update()
    synchronize(device)
    assert peak.peak_mb >= 4.0
    del buffer


def test_synchronize_is_a_noop_on_cpu():
    synchronize("cpu")  # 不应抛异常