"""Tests for GPU passthrough module."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
from typer.testing import CliRunner

from gaol.gpu import (
    AmdGPU,
    DeviceMapping,
    GPUConfig,
    GPUDetector,
    GPUInfo,
    GPUSupport,
    GPUVendor,
    IntelGPU,
    NvidiaGPU,
)


class TestGPUModels:
    """Tests for GPU configuration models."""

    def test_gpu_vendor_enum(self) -> None:
        """Test GPUVendor enum values."""
        assert GPUVendor.NVIDIA.value == "nvidia"
        assert GPUVendor.AMD.value == "amd"
        assert GPUVendor.INTEL.value == "intel"
        assert GPUVendor.AUTO.value == "auto"

    def test_gpu_info_creation(self) -> None:
        """Test creating GPUInfo model."""
        info = GPUInfo(
            index=0,
            name="NVIDIA GeForce RTX 3080",
            vendor=GPUVendor.NVIDIA,
            uuid="GPU-abc123",
            memory_mb=10240,
            pci_bus_id="0000:05:00.0",
        )
        assert info.index == 0
        assert info.name == "NVIDIA GeForce RTX 3080"
        assert info.vendor == GPUVendor.NVIDIA
        assert info.uuid == "GPU-abc123"
        assert info.memory_mb == 10240
        assert info.pci_bus_id == "0000:05:00.0"

    def test_gpu_info_minimal(self) -> None:
        """Test creating GPUInfo with minimal fields."""
        info = GPUInfo(
            index=0,
            name="Generic GPU",
            vendor=GPUVendor.AUTO,
        )
        assert info.index == 0
        assert info.name == "Generic GPU"
        assert info.uuid is None
        assert info.memory_mb is None
        assert info.pci_bus_id is None

    def test_gpu_config_defaults(self) -> None:
        """Test GPUConfig default values."""
        config = GPUConfig()
        assert config.enabled is False
        assert config.vendor == GPUVendor.AUTO
        assert config.device_ids == []
        assert config.capabilities == ["compute", "utility"]
        assert config.memory_limit is None

    def test_gpu_config_enabled(self) -> None:
        """Test GPUConfig with enabled GPU."""
        config = GPUConfig(
            enabled=True,
            vendor=GPUVendor.NVIDIA,
            device_ids=[0, 1],
            capabilities=["compute", "utility", "graphics"],
            memory_limit="8g",
        )
        assert config.enabled is True
        assert config.vendor == GPUVendor.NVIDIA
        assert config.device_ids == [0, 1]
        assert "graphics" in config.capabilities
        assert config.memory_limit == "8g"

    def test_device_mapping_creation(self) -> None:
        """Test creating DeviceMapping model."""
        mapping = DeviceMapping(
            host_path="/dev/nvidia0",
            container_path="/dev/nvidia0",
            permissions="rwm",
        )
        assert mapping.host_path == "/dev/nvidia0"
        assert mapping.container_path == "/dev/nvidia0"
        assert mapping.permissions == "rwm"

    def test_device_mapping_defaults(self) -> None:
        """Test DeviceMapping default values."""
        mapping = DeviceMapping(host_path="/dev/dri/card0")
        assert mapping.host_path == "/dev/dri/card0"
        assert mapping.container_path is None
        assert mapping.permissions == "rwm"

    def test_gpu_support_creation(self) -> None:
        """Test creating GPUSupport model."""
        support = GPUSupport(
            nvidia_available=True,
            nvidia_driver_version="535.154.05",
            nvidia_runtime_available=True,
            rocm_available=False,
            intel_available=True,
        )
        assert support.nvidia_available is True
        assert support.nvidia_driver_version == "535.154.05"
        assert support.nvidia_runtime_available is True
        assert support.rocm_available is False
        assert support.intel_available is True

    def test_gpu_support_defaults(self) -> None:
        """Test GPUSupport default values."""
        support = GPUSupport()
        assert support.nvidia_available is False
        assert support.nvidia_driver_version is None
        assert support.nvidia_runtime_available is False
        assert support.rocm_available is False
        assert support.intel_available is False


class TestNvidiaGPU:
    """Tests for NVIDIA GPU handler."""

    def test_get_docker_config_all_gpus(self) -> None:
        """Test Docker config for all GPUs."""
        config = GPUConfig(enabled=True)
        result = NvidiaGPU.get_docker_config(config)

        assert "device_requests" in result
        assert "environment" in result
        assert result["device_requests"][0]["Driver"] == "nvidia"
        assert result["device_requests"][0]["DeviceIDs"] == ["all"]
        assert result["environment"]["NVIDIA_VISIBLE_DEVICES"] == "all"

    def test_get_docker_config_specific_gpus(self) -> None:
        """Test Docker config for specific GPUs."""
        config = GPUConfig(enabled=True, device_ids=[0, 2])
        result = NvidiaGPU.get_docker_config(config)

        assert result["device_requests"][0]["DeviceIDs"] == ["0", "2"]
        assert result["environment"]["NVIDIA_VISIBLE_DEVICES"] == "0,2"

    def test_get_docker_config_capabilities(self) -> None:
        """Test Docker config with custom capabilities."""
        config = GPUConfig(
            enabled=True,
            capabilities=["compute", "graphics", "video"],
        )
        result = NvidiaGPU.get_docker_config(config)

        assert result["device_requests"][0]["Capabilities"] == [
            ["compute", "graphics", "video"]
        ]
        assert "compute,graphics,video" in result["environment"]["NVIDIA_DRIVER_CAPABILITIES"]

    def test_get_docker_config_memory_limit(self) -> None:
        """Test Docker config with memory limit."""
        config = GPUConfig(enabled=True, memory_limit="4g")
        result = NvidiaGPU.get_docker_config(config)

        assert result["environment"]["NVIDIA_MEM_PER_GPU"] == "4g"

    @patch("pathlib.Path.exists")
    def test_get_nspawn_config(self, mock_exists: MagicMock) -> None:
        """Test nspawn config generation."""
        mock_exists.return_value = True
        config = GPUConfig(enabled=True)
        result = NvidiaGPU.get_nspawn_config(config)

        assert "binds" in result
        assert "environment" in result
        assert result["environment"]["NVIDIA_VISIBLE_DEVICES"] == "all"

    def test_get_libvirt_config(self) -> None:
        """Test libvirt config generation."""
        config = GPUConfig(enabled=True)
        pci_addresses = ["0000:05:00.0", "0000:06:00.0"]
        result = NvidiaGPU.get_libvirt_config(config, pci_addresses)

        assert "hostdev_xml" in result
        assert "requires_iommu" in result
        assert result["requires_iommu"] is True
        assert "0x05" in result["hostdev_xml"]
        assert "0x06" in result["hostdev_xml"]

    @patch("pathlib.Path.exists")
    def test_get_required_devices(self, mock_exists: MagicMock) -> None:
        """Test getting required NVIDIA device paths."""
        mock_exists.return_value = True
        devices = NvidiaGPU.get_required_devices([0])

        assert "/dev/nvidia0" in devices

    def test_get_required_capabilities(self) -> None:
        """Test getting required Linux capabilities."""
        caps = NvidiaGPU.get_required_capabilities()
        assert "SYS_ADMIN" in caps


class TestAmdGPU:
    """Tests for AMD GPU handler."""

    @patch("pathlib.Path.exists")
    @patch("pathlib.Path.iterdir")
    def test_get_docker_config(
        self, mock_iterdir: MagicMock, mock_exists: MagicMock
    ) -> None:
        """Test Docker config for AMD GPUs."""
        mock_exists.return_value = True
        mock_iterdir.return_value = []
        config = GPUConfig(enabled=True)
        result = AmdGPU.get_docker_config(config)

        assert "devices" in result
        assert "environment" in result

    @patch("pathlib.Path.exists")
    @patch("pathlib.Path.iterdir")
    def test_get_docker_config_specific_gpus(
        self, mock_iterdir: MagicMock, mock_exists: MagicMock
    ) -> None:
        """Test Docker config for specific AMD GPUs."""
        mock_exists.return_value = True
        mock_iterdir.return_value = []
        config = GPUConfig(enabled=True, device_ids=[0, 1])
        result = AmdGPU.get_docker_config(config)

        assert "HIP_VISIBLE_DEVICES" in result["environment"]
        assert result["environment"]["HIP_VISIBLE_DEVICES"] == "0,1"

    def test_get_libvirt_config(self) -> None:
        """Test libvirt config for AMD GPUs."""
        config = GPUConfig(enabled=True)
        pci_addresses = ["0000:07:00.0"]
        result = AmdGPU.get_libvirt_config(config, pci_addresses)

        assert "hostdev_xml" in result
        assert "0x07" in result["hostdev_xml"]

    def test_get_render_node_for_gpu(self) -> None:
        """Test getting render node path for GPU index."""
        with patch("pathlib.Path.exists") as mock_exists:
            mock_exists.return_value = True
            result = AmdGPU.get_render_node_for_gpu(0)
            assert result == "/dev/dri/renderD128"

            result = AmdGPU.get_render_node_for_gpu(1)
            assert result == "/dev/dri/renderD129"

    def test_get_required_capabilities(self) -> None:
        """Test getting required Linux capabilities for AMD."""
        caps = AmdGPU.get_required_capabilities()
        assert "SYS_ADMIN" in caps


class TestIntelGPU:
    """Tests for Intel GPU handler."""

    @patch("pathlib.Path.exists")
    @patch("pathlib.Path.iterdir")
    def test_get_docker_config(
        self, mock_iterdir: MagicMock, mock_exists: MagicMock
    ) -> None:
        """Test Docker config for Intel GPUs."""
        mock_exists.return_value = True
        mock_iterdir.return_value = []
        config = GPUConfig(enabled=True)
        result = IntelGPU.get_docker_config(config)

        assert "devices" in result
        assert "environment" in result

    def test_get_libvirt_config(self) -> None:
        """Test libvirt config for Intel GPUs."""
        config = GPUConfig(enabled=True)
        pci_addresses = ["0000:03:00.0"]
        result = IntelGPU.get_libvirt_config(config, pci_addresses)

        assert "hostdev_xml" in result
        assert "0x03" in result["hostdev_xml"]

    def test_is_integrated_gpu(self) -> None:
        """Test detecting integrated GPUs."""
        # Bus 00 = integrated
        assert IntelGPU.is_integrated_gpu("0000:00:02.0") is True

        # Other buses = discrete
        assert IntelGPU.is_integrated_gpu("0000:03:00.0") is False

        # No address = assume integrated
        assert IntelGPU.is_integrated_gpu(None) is True

    def test_get_required_capabilities(self) -> None:
        """Test Intel GPU capabilities (typically none required)."""
        caps = IntelGPU.get_required_capabilities()
        assert isinstance(caps, list)


class TestGPUDetector:
    """Tests for GPU detection utilities."""

    @patch("gaol.gpu.detection.subprocess.run")
    @patch("gaol.gpu.detection.shutil.which")
    def test_detect_nvidia_gpus(
        self, mock_which: MagicMock, mock_run: MagicMock
    ) -> None:
        """Test detecting NVIDIA GPUs via nvidia-smi."""
        mock_which.return_value = "/usr/bin/nvidia-smi"
        mock_run.return_value = MagicMock(
            returncode=0,
            stdout="0, NVIDIA GeForce RTX 3080, GPU-abc123, 10240, 0000:05:00.0\n"
                   "1, NVIDIA GeForce RTX 3080, GPU-def456, 10240, 0000:06:00.0\n",
        )

        gpus = GPUDetector.detect_nvidia()

        assert len(gpus) == 2
        assert gpus[0].index == 0
        assert gpus[0].name == "NVIDIA GeForce RTX 3080"
        assert gpus[0].vendor == GPUVendor.NVIDIA
        assert gpus[0].uuid == "GPU-abc123"
        assert gpus[0].memory_mb == 10240
        assert gpus[0].pci_bus_id == "0000:05:00.0"

    @patch("subprocess.run")
    def test_detect_nvidia_not_available(self, mock_run: MagicMock) -> None:
        """Test when nvidia-smi is not available."""
        mock_run.side_effect = FileNotFoundError()
        gpus = GPUDetector.detect_nvidia()
        assert gpus == []

    @patch("gaol.gpu.detection.subprocess.run")
    @patch("gaol.gpu.detection.shutil.which")
    def test_detect_amd_gpus(
        self, mock_which: MagicMock, mock_run: MagicMock
    ) -> None:
        """Test detecting AMD GPUs via rocm-smi."""
        mock_which.return_value = "/usr/bin/rocm-smi"
        mock_run.return_value = MagicMock(
            returncode=0,
            stdout="GPU[0]: AMD Radeon RX 7900 XTX\nGPU[1]: AMD Radeon RX 6900 XT\n",
        )

        gpus = GPUDetector.detect_amd()

        assert len(gpus) == 2
        assert gpus[0].index == 0
        assert "Radeon" in gpus[0].name
        assert gpus[0].vendor == GPUVendor.AMD

    @patch("subprocess.run")
    def test_detect_amd_not_available(self, mock_run: MagicMock) -> None:
        """Test when rocm-smi is not available."""
        mock_run.side_effect = FileNotFoundError()
        gpus = GPUDetector.detect_amd()
        assert gpus == []

    @patch("pathlib.Path.exists")
    @patch("pathlib.Path.iterdir")
    @patch("subprocess.run")
    def test_detect_intel_gpus(
        self,
        mock_run: MagicMock,
        mock_iterdir: MagicMock,
        mock_exists: MagicMock,
    ) -> None:
        """Test detecting Intel GPUs."""
        mock_exists.return_value = True

        # Create mock Path objects for renderD128 and renderD129
        mock_path1 = MagicMock()
        mock_path1.name = "renderD128"
        mock_path2 = MagicMock()
        mock_path2.name = "renderD129"
        mock_iterdir.return_value = [mock_path1, mock_path2]

        mock_run.return_value = MagicMock(
            returncode=0,
            stdout="00:02.0 VGA compatible controller: Intel Corporation UHD Graphics",
        )

        gpus = GPUDetector.detect_intel()

        # Should detect at least one GPU
        assert len(gpus) >= 0  # May be 0 if detection logic filters it

    @patch.object(GPUDetector, "detect_nvidia")
    @patch.object(GPUDetector, "detect_amd")
    @patch.object(GPUDetector, "detect_intel")
    def test_detect_all(
        self,
        mock_intel: MagicMock,
        mock_amd: MagicMock,
        mock_nvidia: MagicMock,
    ) -> None:
        """Test detecting all GPUs."""
        mock_nvidia.return_value = [
            GPUInfo(index=0, name="NVIDIA GPU", vendor=GPUVendor.NVIDIA)
        ]
        mock_amd.return_value = [
            GPUInfo(index=0, name="AMD GPU", vendor=GPUVendor.AMD)
        ]
        mock_intel.return_value = [
            GPUInfo(index=0, name="Intel GPU", vendor=GPUVendor.INTEL)
        ]

        gpus = GPUDetector.detect_all()

        assert len(gpus) == 3
        vendors = {g.vendor for g in gpus}
        assert GPUVendor.NVIDIA in vendors
        assert GPUVendor.AMD in vendors
        assert GPUVendor.INTEL in vendors

    @patch("subprocess.run")
    @patch("shutil.which")
    @patch("pathlib.Path.exists")
    def test_get_support_info(
        self,
        mock_exists: MagicMock,
        mock_which: MagicMock,
        mock_run: MagicMock,
    ) -> None:
        """Test getting GPU support information."""
        mock_which.side_effect = lambda x: x if x == "nvidia-smi" else None
        mock_exists.return_value = False
        mock_run.return_value = MagicMock(
            returncode=0,
            stdout="535.154.05",
        )

        support = GPUDetector.get_support_info()

        assert isinstance(support, GPUSupport)

    @patch("gaol.gpu.detection.shutil.which")
    def test_is_nvidia_runtime_available(self, mock_which: MagicMock) -> None:
        """Test checking nvidia-container-runtime availability."""
        mock_which.return_value = "/usr/bin/nvidia-container-runtime"
        assert GPUDetector.is_nvidia_runtime_available() is True

        mock_which.return_value = None
        assert GPUDetector.is_nvidia_runtime_available() is False

    @patch("shutil.which")
    def test_is_rocm_available(self, mock_which: MagicMock) -> None:
        """Test checking ROCm availability."""
        mock_which.return_value = "/usr/bin/rocm-smi"
        assert GPUDetector.is_rocm_available() is True

        mock_which.return_value = None
        assert GPUDetector.is_rocm_available() is False


class TestGPUCLI:
    """Tests for GPU CLI commands."""

    @pytest.fixture
    def runner(self) -> CliRunner:
        """Create CLI test runner."""
        return CliRunner()

    @patch.object(GPUDetector, "detect_all")
    def test_gpu_list_command(
        self, mock_detect: MagicMock, runner: CliRunner
    ) -> None:
        """Test 'gpu list' command."""
        from gaol.cli import app

        mock_detect.return_value = [
            GPUInfo(
                index=0,
                name="NVIDIA GeForce RTX 3080",
                vendor=GPUVendor.NVIDIA,
                memory_mb=10240,
                pci_bus_id="0000:05:00.0",
            )
        ]

        result = runner.invoke(app, ["gpu", "list"])

        assert result.exit_code == 0
        assert "NVIDIA" in result.stdout or "nvidia" in result.stdout
        assert "RTX 3080" in result.stdout

    @patch.object(GPUDetector, "detect_all")
    def test_gpu_list_empty(
        self, mock_detect: MagicMock, runner: CliRunner
    ) -> None:
        """Test 'gpu list' with no GPUs."""
        from gaol.cli import app

        mock_detect.return_value = []

        result = runner.invoke(app, ["gpu", "list"])

        assert result.exit_code == 0
        assert "No GPUs detected" in result.stdout

    @patch.object(GPUDetector, "detect_all")
    def test_gpu_list_json(
        self, mock_detect: MagicMock, runner: CliRunner
    ) -> None:
        """Test 'gpu list --json' command."""
        from gaol.cli import app

        mock_detect.return_value = [
            GPUInfo(
                index=0,
                name="Test GPU",
                vendor=GPUVendor.NVIDIA,
            )
        ]

        result = runner.invoke(app, ["gpu", "list", "--json"])

        assert result.exit_code == 0
        assert '"index": 0' in result.stdout
        assert '"name": "Test GPU"' in result.stdout

    @patch.object(GPUDetector, "detect_all")
    def test_gpu_info_command(
        self, mock_detect: MagicMock, runner: CliRunner
    ) -> None:
        """Test 'gpu info' command."""
        from gaol.cli import app

        mock_detect.return_value = [
            GPUInfo(
                index=0,
                name="NVIDIA GeForce RTX 3080",
                vendor=GPUVendor.NVIDIA,
                uuid="GPU-abc123",
                memory_mb=10240,
                pci_bus_id="0000:05:00.0",
            )
        ]

        result = runner.invoke(app, ["gpu", "info", "0"])

        assert result.exit_code == 0
        assert "GPU 0" in result.stdout
        assert "NVIDIA GeForce RTX 3080" in result.stdout

    @patch.object(GPUDetector, "detect_all")
    def test_gpu_info_not_found(
        self, mock_detect: MagicMock, runner: CliRunner
    ) -> None:
        """Test 'gpu info' with non-existent GPU."""
        from gaol.cli import app

        mock_detect.return_value = []

        result = runner.invoke(app, ["gpu", "info", "99"])

        assert result.exit_code == 1
        # Error is written to stderr, but CliRunner captures both
        assert "GPU not found" in result.output

    @patch.object(GPUDetector, "get_support_info")
    @patch.object(GPUDetector, "detect_all")
    def test_gpu_check_command(
        self,
        mock_detect: MagicMock,
        mock_support: MagicMock,
        runner: CliRunner,
    ) -> None:
        """Test 'gpu check' command."""
        from gaol.cli import app

        mock_support.return_value = GPUSupport(
            nvidia_available=True,
            nvidia_driver_version="535.154.05",
            nvidia_runtime_available=True,
            rocm_available=False,
            intel_available=True,
        )
        mock_detect.return_value = [
            GPUInfo(index=0, name="Test GPU", vendor=GPUVendor.NVIDIA)
        ]

        result = runner.invoke(app, ["gpu", "check"])

        assert result.exit_code == 0
        assert "NVIDIA" in result.stdout
        assert "AMD" in result.stdout
        assert "Intel" in result.stdout


class TestContainerConfigWithGPU:
    """Tests for ContainerConfig with GPU settings."""

    @pytest.fixture(autouse=True)
    def rebuild_container_config(self) -> None:
        """Rebuild ContainerConfig to resolve forward references."""
        from gaol.models.container import ContainerConfig
        from gaol.gpu.models import DeviceMapping, GPUConfig

        ContainerConfig.model_rebuild()

    def test_container_config_with_gpu(self) -> None:
        """Test creating a container config with GPU enabled."""
        from gaol.models.container import ContainerConfig

        config = ContainerConfig(
            name="ml-container",
            gpu=GPUConfig(
                enabled=True,
                vendor=GPUVendor.NVIDIA,
                device_ids=[0],
            ),
        )

        assert config.gpu is not None
        assert config.gpu.enabled is True
        assert config.gpu.vendor == GPUVendor.NVIDIA
        assert config.gpu.device_ids == [0]

    def test_container_config_with_devices(self) -> None:
        """Test creating a container config with device mappings."""
        from gaol.models.container import ContainerConfig

        config = ContainerConfig(
            name="device-container",
            devices=[
                DeviceMapping(host_path="/dev/dri/card0"),
                DeviceMapping(host_path="/dev/dri/renderD128"),
            ],
        )

        assert len(config.devices) == 2
        assert config.devices[0].host_path == "/dev/dri/card0"

    def test_container_config_defaults(self) -> None:
        """Test container config GPU defaults."""
        from gaol.models.container import ContainerConfig

        config = ContainerConfig(name="test-container")

        assert config.gpu is None
        assert config.devices == []
