"""Deterministic contracts for accelerator detection and fallback behavior."""

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import torch

from obliteratus import device


@pytest.mark.parametrize(
    ("cuda", "mps", "expected"),
    [(True, False, "cuda"), (False, True, "mps"), (False, False, "cpu")],
)
def test_auto_device_preference(monkeypatch, cuda, mps, expected):
    monkeypatch.setattr(device, "is_cuda", lambda: cuda)
    monkeypatch.setattr(device, "is_mps", lambda: mps)
    assert device.get_device() == expected
    assert device.is_gpu_available() is (cuda or mps)


def test_explicit_device_validation(monkeypatch):
    monkeypatch.setattr(device, "is_cuda", lambda: False)
    monkeypatch.setattr(device, "is_mps", lambda: False)

    assert device.get_device("cpu") == "cpu"
    with pytest.raises(RuntimeError, match="CUDA was requested.*device='auto'"):
        device.get_device("cuda:0")
    with pytest.raises(RuntimeError, match="MPS was requested.*device='auto'"):
        device.get_device("mps")
    with pytest.raises(ValueError, match="Unknown device 'tpu'"):
        device.get_device("tpu")
    with pytest.raises(ValueError, match="Unknown device 'cuda:gpu'"):
        device.get_device("cuda:gpu")

    monkeypatch.setattr(device, "is_cuda", lambda: True)
    assert device.get_device("cuda") == "cuda"
    assert device.get_device("cuda:3") == "cuda:3"
    monkeypatch.setattr(device, "is_mps", lambda: True)
    assert device.get_device("mps") == "mps"


def test_names_and_device_counts(monkeypatch):
    monkeypatch.setattr(device, "is_cuda", lambda: True)
    monkeypatch.setattr(device.torch.cuda, "get_device_name", lambda _index: "Test GPU")
    monkeypatch.setattr(device.torch.cuda, "device_count", lambda: 4)
    assert device.get_device_name() == "Test GPU"
    assert device.device_count() == 4

    monkeypatch.setattr(device, "is_cuda", lambda: False)
    monkeypatch.setattr(device, "is_mps", lambda: True)
    monkeypatch.setattr(device.platform, "processor", lambda: "M3")
    assert device.get_device_name() == "Apple M3 (MPS)"
    assert device.device_count() == 1

    monkeypatch.setattr(device, "is_mps", lambda: False)
    assert device.get_device_name() == "CPU"
    assert device.device_count() == 0


def test_system_memory_sources_and_fallback(monkeypatch):
    gib = 1024**3
    fake_psutil = SimpleNamespace(
        virtual_memory=lambda: SimpleNamespace(total=32 * gib, available=12 * gib),
    )
    monkeypatch.setitem(__import__("sys").modules, "psutil", fake_psutil)
    assert device._system_memory_gb() == (32.0, 12.0)

    monkeypatch.delitem(__import__("sys").modules, "psutil", raising=False)
    real_import = __import__("builtins").__import__

    def reject_psutil(name, *args, **kwargs):
        if name == "psutil":
            raise ImportError
        return real_import(name, *args, **kwargs)

    monkeypatch.setattr("builtins.__import__", reject_psutil)
    monkeypatch.setattr(device.os, "sysconf", lambda name: {"SC_PHYS_PAGES": 4, "SC_PAGE_SIZE": gib}[name])
    assert device._system_memory_gb() == (4.0, 2.4)

    monkeypatch.setattr(device.os, "sysconf", Mock(side_effect=ValueError))
    assert device._system_memory_gb() == (16.0, 8.0)


def test_memory_info_for_cuda_and_cuda_fallback(monkeypatch):
    gib = 1024**3
    monkeypatch.setattr(device, "is_cuda", lambda: True)
    monkeypatch.setattr(device, "get_device_name", lambda: "GPU")
    monkeypatch.setattr(device.torch.cuda, "mem_get_info", lambda _index: (6 * gib, 8 * gib))
    monkeypatch.setattr(device.torch.cuda, "memory_allocated", lambda _index: 1 * gib)
    monkeypatch.setattr(device.torch.cuda, "memory_reserved", lambda _index: 2 * gib)
    assert device.get_memory_info(2) == device.MemoryInfo(1, 2, 8, 6, "GPU")

    monkeypatch.setattr(device.torch.cuda, "mem_get_info", Mock(side_effect=RuntimeError("unsupported")))
    monkeypatch.setattr(
        device.torch.cuda,
        "get_device_properties",
        lambda _index: SimpleNamespace(total_memory=10 * gib),
    )
    assert device.get_memory_info(2) == device.MemoryInfo(total_gb=10, free_gb=10, device_name="GPU")


def test_memory_info_for_mps_cpu_and_total_free(monkeypatch):
    monkeypatch.setattr(device, "is_cuda", lambda: False)
    monkeypatch.setattr(device, "is_mps", lambda: True)
    monkeypatch.setattr(device, "get_device_name", lambda: "MPS")
    monkeypatch.setattr(device, "_system_memory_gb", lambda: (20.0, 8.0))
    assert device.get_memory_info() == device.MemoryInfo(6, 0, 14, 8, "MPS")
    assert device.get_total_free_gb() == pytest.approx(5.6)

    monkeypatch.setattr(device, "is_mps", lambda: False)
    monkeypatch.setattr(device, "get_device_name", lambda: "CPU")
    assert device.get_memory_info() == device.MemoryInfo(total_gb=20, free_gb=8, device_name="CPU")
    assert device.get_total_free_gb() == 0


def test_total_cuda_memory_sums_query_and_fallback(monkeypatch):
    gib = 1024**3
    monkeypatch.setattr(device, "is_cuda", lambda: True)
    monkeypatch.setattr(device.torch.cuda, "device_count", lambda: 2)
    monkeypatch.setattr(
        device.torch.cuda,
        "mem_get_info",
        Mock(side_effect=[(3 * gib, 4 * gib), RuntimeError("no query")]),
    )
    monkeypatch.setattr(
        device.torch.cuda,
        "get_device_properties",
        lambda _index: SimpleNamespace(total_memory=5 * gib),
    )
    assert device.get_total_free_gb() == 8.0


def test_cache_cleanup_paths_are_best_effort(monkeypatch):
    cuda_empty = Mock()
    monkeypatch.setattr(device, "is_cuda", lambda: True)
    monkeypatch.setattr(device.torch.cuda, "empty_cache", cuda_empty)
    device.empty_cache()
    cuda_empty.assert_called_once_with()

    sync = Mock()
    reset = Mock()
    monkeypatch.setattr(device.torch.cuda, "empty_cache", Mock(side_effect=RuntimeError("busy")))
    monkeypatch.setattr(device.torch.cuda, "synchronize", sync)
    monkeypatch.setattr(device.torch.cuda, "reset_peak_memory_stats", reset)
    monkeypatch.setattr(device.gc, "collect", Mock())
    device.free_gpu_memory()
    sync.assert_called_once_with()
    reset.assert_called_once_with()

    mps_empty = Mock()
    mps_sync = Mock(side_effect=RuntimeError("busy"))
    monkeypatch.setattr(device, "is_cuda", lambda: False)
    monkeypatch.setattr(device, "is_mps", lambda: True)
    monkeypatch.setattr(device.torch, "mps", SimpleNamespace(empty_cache=mps_empty, synchronize=mps_sync))
    device.empty_cache()
    mps_empty.side_effect = RuntimeError("busy")
    device.free_gpu_memory()
    assert mps_empty.call_count == 2
    mps_sync.assert_called_once_with()


def test_cuda_cleanup_tolerates_every_recovery_failure(monkeypatch):
    monkeypatch.setattr(device, "is_cuda", lambda: True)
    monkeypatch.setattr(device.torch.cuda, "empty_cache", Mock(side_effect=RuntimeError("busy")))
    monkeypatch.setattr(device.torch.cuda, "synchronize", Mock(side_effect=RuntimeError("lost")))
    monkeypatch.setattr(
        device.torch.cuda,
        "reset_peak_memory_stats",
        Mock(side_effect=RuntimeError("lost")),
    )
    monkeypatch.setattr(device.gc, "collect", Mock())
    device.free_gpu_memory()


@pytest.mark.parametrize(
    ("version", "expected"),
    [("2.2.9", False), ("2.3.0", True), ("2.10.1", True)],
)
def test_mps_bfloat16_version_boundary(monkeypatch, version, expected):
    monkeypatch.setattr(device, "is_cuda", lambda: False)
    monkeypatch.setattr(device.torch, "__version__", version)
    assert device.supports_bfloat16("mps") is expected


def test_seed_dtype_and_capability_contracts(monkeypatch):
    manual_seed = Mock()
    cuda_seed = Mock()
    monkeypatch.setattr(device.torch, "manual_seed", manual_seed)
    monkeypatch.setattr(device.torch.cuda, "manual_seed_all", cuda_seed)
    monkeypatch.setattr(device, "is_cuda", lambda: True)
    device.set_seed_all(7)
    manual_seed.assert_called_once_with(7)
    cuda_seed.assert_called_once_with(7)

    assert device.default_dtype("cpu") is torch.float32
    assert device.default_dtype("cuda") is torch.float16
    monkeypatch.setattr(device.torch.cuda, "get_device_capability", lambda _index: (8, 0))
    assert device.supports_bfloat16("cuda") is True
    monkeypatch.setattr(device.torch.cuda, "get_device_capability", lambda _index: (7, 5))
    assert device.supports_bfloat16("cuda") is False
    monkeypatch.setattr(device, "is_cuda", lambda: False)
    assert device.supports_bfloat16("cuda") is False
    assert device.supports_bfloat16("cpu") is True
    assert device.supports_float64("mps") is False
    assert device.supports_float64("cpu") is True
    assert device.supports_bitsandbytes("cuda:1") is True
    assert device.supports_bitsandbytes("mps") is False
    assert device.supports_device_map_auto("cuda") is True
    assert device.supports_device_map_auto("cpu") is False


def test_svd_dtype_and_oom_matching():
    assert device.safe_svd_dtype(torch.ones(1, dtype=torch.float64)) is torch.float64
    assert device.safe_svd_dtype(torch.ones(1, dtype=torch.float16)) is torch.float32
    assert device.is_oom_error(torch.cuda.OutOfMemoryError("oom")) is True
    assert device.is_oom_error(RuntimeError("MPS backend out of memory")) is True
    assert device.is_oom_error(RuntimeError("other")) is False


def test_mps_svd_dtype_uses_float32_without_requiring_mps_hardware():
    tensor = SimpleNamespace(device=SimpleNamespace(type="mps"), dtype=torch.float64)
    assert device.safe_svd_dtype(tensor) is torch.float32


def test_configure_cuda_allocator(monkeypatch):
    monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False)
    monkeypatch.setattr(device, "is_cuda", lambda: True)
    device.configure_cuda_alloc()
    assert device.os.environ["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True"

    monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "existing")
    device.configure_cuda_alloc()
    assert device.os.environ["PYTORCH_CUDA_ALLOC_CONF"] == "existing"