"""
Tests for GateWay GPU driver setup automation.
"""

from unittest.mock import AsyncMock, Mock, patch

import pytest

from verlex_server.gateway.backend.gpu_setup import (
    CUDAVersion,
    DriverConfig,
    GPUDetector,
    GPUInfo,
    GPUSetupManager,
    GPUSetupResult,
    GPUSetupStatus,
    GPUVendor,
    LinuxDistro,
    LinuxDistroDetector,
    NVIDIADriverInstaller,
    detect_gpus,
    get_recommended_cuda_version,
)


class TestGPUInfo:
    """Tests for GPUInfo dataclass."""

    def test_to_dict(self):
        """Test converting GPUInfo to dictionary."""
        gpu = GPUInfo(
            vendor=GPUVendor.NVIDIA,
            model="Tesla V100",
            device_id="0",
            pci_bus_id="00:1e.0",
            memory_mb=16384,
            driver_version="535.54.03",
            cuda_version="12.1",
            compute_capability="7.0",
        )

        data = gpu.to_dict()

        assert data["vendor"] == "nvidia"
        assert data["model"] == "Tesla V100"
        assert data["memory_mb"] == 16384
        assert data["driver_version"] == "535.54.03"
        assert data["compute_capability"] == "7.0"


class TestDriverConfig:
    """Tests for DriverConfig dataclass."""

    def test_default_values(self):
        """Test default configuration values."""
        config = DriverConfig()

        assert config.cuda_version == "12.1"
        assert config.cudnn_version == "8.9"
        assert config.install_container_toolkit is True
        assert config.silent is True

    def test_to_dict(self):
        """Test converting to dictionary."""
        config = DriverConfig(
            cuda_version="11.8",
            cudnn_version="8.6",
        )

        data = config.to_dict()

        assert data["cuda_version"] == "11.8"
        assert data["cudnn_version"] == "8.6"


class TestGPUSetupResult:
    """Tests for GPUSetupResult dataclass."""

    def test_to_dict(self):
        """Test converting result to dictionary."""
        result = GPUSetupResult(
            status=GPUSetupStatus.COMPLETED,
            gpus_detected=[
                GPUInfo(
                    vendor=GPUVendor.NVIDIA,
                    model="A100",
                    device_id="0",
                    pci_bus_id="00:1e.0",
                    memory_mb=40960,
                )
            ],
            driver_installed=True,
            driver_version="535.54.03",
            cuda_installed=True,
            cuda_version="12.1",
        )

        data = result.to_dict()

        assert data["status"] == "completed"
        assert len(data["gpus_detected"]) == 1
        assert data["driver_installed"] is True
        assert data["cuda_version"] == "12.1"


class TestGPUDetector:
    """Tests for GPUDetector."""

    def test_compute_capability_mapping(self):
        """Test compute capability mapping for known GPUs."""
        detector = GPUDetector()

        assert detector.NVIDIA_COMPUTE_CAPABILITIES.get("H100") == "9.0"
        assert detector.NVIDIA_COMPUTE_CAPABILITIES.get("A100") == "8.0"
        assert detector.NVIDIA_COMPUTE_CAPABILITIES.get("V100") == "7.0"
        assert detector.NVIDIA_COMPUTE_CAPABILITIES.get("T4") == "7.5"

    def test_cuda_driver_requirements(self):
        """Test CUDA driver version requirements."""
        detector = GPUDetector()

        assert detector.CUDA_DRIVER_REQUIREMENTS.get("12.1") == "530.30.02"
        assert detector.CUDA_DRIVER_REQUIREMENTS.get("11.8") == "520.61.05"

    def test_get_minimum_driver_for_cuda(self):
        """Test getting minimum driver for CUDA version."""
        detector = GPUDetector()

        min_driver = detector.get_minimum_driver_for_cuda("12.1")
        assert min_driver == "530.30.02"

        min_driver = detector.get_minimum_driver_for_cuda("11.8")
        assert min_driver == "520.61.05"

    def test_is_driver_compatible(self):
        """Test driver compatibility check."""
        detector = GPUDetector()

        # Driver 535 should be compatible with CUDA 12.1 (needs 530+)
        assert detector.is_driver_compatible("535.54.03", "12.1") is True

        # Driver 515 should not be compatible with CUDA 12.1
        assert detector.is_driver_compatible("515.43.04", "12.1") is False

        # Driver 525 should be compatible with CUDA 11.8
        assert detector.is_driver_compatible("525.60.13", "11.8") is True

    def test_version_comparison(self):
        """Test version string comparison."""
        detector = GPUDetector()

        assert detector._compare_versions("535.54.03", "530.30.02") > 0
        assert detector._compare_versions("530.30.02", "535.54.03") < 0
        assert detector._compare_versions("530.30.02", "530.30.02") == 0
        assert detector._compare_versions("530.30", "530.30.02") < 0

    @pytest.mark.asyncio
    async def test_detect_without_gpu(self):
        """Test detection when no GPU is present."""
        detector = GPUDetector()

        # Mock nvidia-smi not found
        with patch('shutil.which', return_value=None):
            gpus = await detector._detect_nvidia_smi()
            assert gpus == []

    def test_get_compute_capability(self):
        """Test getting compute capability from model name."""
        detector = GPUDetector()

        assert detector._get_compute_capability("NVIDIA A100-SXM4-40GB") == "8.0"
        assert detector._get_compute_capability("Tesla V100-PCIE-32GB") == "7.0"
        assert detector._get_compute_capability("NVIDIA H100 PCIe") == "9.0"
        assert detector._get_compute_capability("Unknown GPU") is None


class TestLinuxDistroDetector:
    """Tests for LinuxDistroDetector."""

    def test_parse_ubuntu_os_release(self):
        """Test parsing Ubuntu os-release."""
        content = '''
PRETTY_NAME="Ubuntu 22.04.3 LTS"
NAME="Ubuntu"
VERSION_ID="22.04"
VERSION="22.04.3 LTS (Jammy Jellyfish)"
VERSION_CODENAME=jammy
ID=ubuntu
ID_LIKE=debian
'''
        distro = LinuxDistroDetector._parse_os_release(content)
        assert distro == LinuxDistro.UBUNTU_22_04

    def test_parse_debian_os_release(self):
        """Test parsing Debian os-release."""
        content = '''
PRETTY_NAME="Debian GNU/Linux 12 (bookworm)"
NAME="Debian GNU/Linux"
VERSION_ID="12"
VERSION="12 (bookworm)"
ID=debian
'''
        distro = LinuxDistroDetector._parse_os_release(content)
        assert distro == LinuxDistro.DEBIAN_12

    def test_parse_centos_os_release(self):
        """Test parsing CentOS os-release."""
        content = '''
NAME="CentOS Stream"
VERSION="8"
ID="centos"
ID_LIKE="rhel fedora"
VERSION_ID="8"
'''
        distro = LinuxDistroDetector._parse_os_release(content)
        assert distro == LinuxDistro.CENTOS_8

    def test_parse_amazon_linux_os_release(self):
        """Test parsing Amazon Linux os-release."""
        content = '''
NAME="Amazon Linux"
VERSION="2"
ID="amzn"
ID_LIKE="centos rhel fedora"
VERSION_ID="2"
'''
        distro = LinuxDistroDetector._parse_os_release(content)
        assert distro == LinuxDistro.AMAZON_LINUX_2

    def test_get_package_manager_apt(self):
        """Test getting apt package manager for Debian-based distros."""
        assert LinuxDistroDetector.get_package_manager(LinuxDistro.UBUNTU_22_04) == "apt"
        assert LinuxDistroDetector.get_package_manager(LinuxDistro.DEBIAN_12) == "apt"

    def test_get_package_manager_yum(self):
        """Test getting yum package manager for RHEL-based distros."""
        assert LinuxDistroDetector.get_package_manager(LinuxDistro.CENTOS_8) == "yum"
        assert LinuxDistroDetector.get_package_manager(LinuxDistro.RHEL_8) == "yum"
        assert LinuxDistroDetector.get_package_manager(LinuxDistro.AMAZON_LINUX_2) == "yum"


class TestNVIDIADriverInstaller:
    """Tests for NVIDIADriverInstaller."""

    def test_cuda_urls_available(self):
        """Test that CUDA URLs are defined for common versions."""
        assert "12.1" in NVIDIADriverInstaller.CUDA_URLS
        assert "11.8" in NVIDIADriverInstaller.CUDA_URLS

        # Check that URLs exist for Ubuntu
        assert LinuxDistro.UBUNTU_22_04 in NVIDIADriverInstaller.CUDA_URLS["12.1"]

    @pytest.mark.asyncio
    async def test_installer_initialization(self):
        """Test installer initialization."""
        config = DriverConfig(cuda_version="12.1")
        callback = Mock()

        installer = NVIDIADriverInstaller(
            config=config,
            distro=LinuxDistro.UBUNTU_22_04,
            output_callback=callback,
        )

        assert installer.config.cuda_version == "12.1"
        assert installer.distro == LinuxDistro.UBUNTU_22_04


class TestGPUSetupManager:
    """Tests for GPUSetupManager."""

    @pytest.mark.asyncio
    async def test_manager_initialization(self):
        """Test manager initialization."""
        manager = GPUSetupManager(
            cuda_version="12.1",
            cudnn_version="8.9",
        )

        assert manager.cuda_version == "12.1"
        assert manager.cudnn_version == "8.9"

    @pytest.mark.asyncio
    async def test_is_gpu_ready_no_gpu(self):
        """Test is_gpu_ready when no GPU is present."""
        manager = GPUSetupManager()

        with patch.object(manager.detector, 'detect', new_callable=AsyncMock) as mock_detect:
            mock_detect.return_value = []

            ready = await manager.is_gpu_ready()
            assert ready is False

    @pytest.mark.asyncio
    async def test_get_gpu_info(self):
        """Test getting GPU info."""
        manager = GPUSetupManager()

        mock_gpu = GPUInfo(
            vendor=GPUVendor.NVIDIA,
            model="A100",
            device_id="0",
            pci_bus_id="00:1e.0",
            memory_mb=40960,
        )

        with patch.object(manager.detector, 'detect', new_callable=AsyncMock) as mock_detect:
            with patch.object(manager.detector, 'get_driver_version', new_callable=AsyncMock) as mock_driver:
                with patch.object(manager.detector, 'get_cuda_version', new_callable=AsyncMock) as mock_cuda:
                    mock_detect.return_value = [mock_gpu]
                    mock_driver.return_value = "535.54.03"
                    mock_cuda.return_value = "12.1"

                    info = await manager.get_gpu_info()

                    assert info["gpu_count"] == 1
                    assert info["driver_version"] == "535.54.03"
                    assert info["cuda_version"] == "12.1"
                    assert info["total_memory_mb"] == 40960

    @pytest.mark.asyncio
    async def test_run_diagnostic(self):
        """Test running GPU diagnostic."""
        manager = GPUSetupManager()

        with patch.object(manager.detector, 'detect', new_callable=AsyncMock) as mock_detect:
            mock_detect.return_value = []

            results = await manager.run_diagnostic()

            assert results["gpu_detected"] is False
            assert "No GPU detected" in results["issues"]


class TestConvenienceFunctions:
    """Tests for module-level convenience functions."""

    def test_get_recommended_cuda_version_modern_gpu(self):
        """Test recommended CUDA version for modern GPUs."""
        assert get_recommended_cuda_version("NVIDIA A100") == "12.1"
        assert get_recommended_cuda_version("NVIDIA H100 PCIe") == "12.1"
        assert get_recommended_cuda_version("GeForce RTX 4090") == "12.1"
        assert get_recommended_cuda_version("NVIDIA L4") == "12.1"

    def test_get_recommended_cuda_version_older_gpu(self):
        """Test recommended CUDA version for older GPUs."""
        # GPUs not in the modern list should get CUDA 11.8
        assert get_recommended_cuda_version("Tesla K80") == "11.8"
        assert get_recommended_cuda_version("Unknown GPU") == "11.8"

    @pytest.mark.asyncio
    async def test_detect_gpus_function(self):
        """Test detect_gpus convenience function."""
        with patch('verlex_server.gateway.backend.gpu_setup.GPUDetector') as MockDetector:
            mock_instance = MockDetector.return_value
            mock_instance.detect = AsyncMock(return_value=[])

            gpus = await detect_gpus()
            assert gpus == []


class TestGPUVendor:
    """Tests for GPUVendor enum."""

    def test_vendor_values(self):
        """Test vendor enum values."""
        assert GPUVendor.NVIDIA.value == "nvidia"
        assert GPUVendor.AMD.value == "amd"
        assert GPUVendor.INTEL.value == "intel"
        assert GPUVendor.NONE.value == "none"


class TestCUDAVersion:
    """Tests for CUDAVersion enum."""

    def test_cuda_version_values(self):
        """Test CUDA version enum values."""
        assert CUDAVersion.CUDA_12_1.value == "12.1"
        assert CUDAVersion.CUDA_11_8.value == "11.8"
        assert CUDAVersion.CUDA_12_4.value == "12.4"


class TestGPUSetupStatus:
    """Tests for GPUSetupStatus enum."""

    def test_status_values(self):
        """Test status enum values."""
        assert GPUSetupStatus.NOT_STARTED.value == "not_started"
        assert GPUSetupStatus.COMPLETED.value == "completed"
        assert GPUSetupStatus.FAILED.value == "failed"
        assert GPUSetupStatus.NO_GPU.value == "no_gpu"


class TestIntegrationScenarios:
    """Integration-style tests for common scenarios."""

    @pytest.mark.asyncio
    async def test_full_setup_no_gpu(self):
        """Test full setup flow when no GPU is present."""
        manager = GPUSetupManager()

        with patch.object(manager.detector, 'detect', new_callable=AsyncMock) as mock_detect:
            mock_detect.return_value = []

            result = await manager.setup()

            assert result.status == GPUSetupStatus.NO_GPU
            assert len(result.gpus_detected) == 0

    @pytest.mark.asyncio
    async def test_already_configured(self):
        """Test when GPU is already configured."""
        manager = GPUSetupManager(cuda_version="12.1")

        mock_gpu = GPUInfo(
            vendor=GPUVendor.NVIDIA,
            model="A100",
            device_id="0",
            pci_bus_id="00:1e.0",
            memory_mb=40960,
        )

        with patch.object(manager, 'is_gpu_ready', new_callable=AsyncMock) as mock_ready:
            with patch.object(manager.detector, 'detect', new_callable=AsyncMock) as mock_detect:
                with patch.object(manager.detector, 'get_driver_version', new_callable=AsyncMock) as mock_driver:
                    with patch.object(manager.detector, 'get_cuda_version', new_callable=AsyncMock) as mock_cuda:
                        mock_ready.return_value = True
                        mock_detect.return_value = [mock_gpu]
                        mock_driver.return_value = "535.54.03"
                        mock_cuda.return_value = "12.1"

                        result = await manager.setup()

                        assert result.status == GPUSetupStatus.COMPLETED
                        assert result.driver_installed is True


class TestEdgeCases:
    """Tests for edge cases and error handling."""

    def test_version_comparison_edge_cases(self):
        """Test version comparison edge cases."""
        detector = GPUDetector()

        # Different length versions
        assert detector._compare_versions("535", "530.30.02") > 0
        assert detector._compare_versions("530.30.02", "535") < 0

        # Same major.minor, different patch
        assert detector._compare_versions("535.54.03", "535.54.02") > 0
        assert detector._compare_versions("535.54.02", "535.54.03") < 0

    def test_unknown_distro_package_manager(self):
        """Test package manager for unknown distro defaults to yum."""
        # Create a mock distro that's not in debian_based set
        distro = LinuxDistro.CENTOS_7
        pkg_manager = LinuxDistroDetector.get_package_manager(distro)
        assert pkg_manager == "yum"

    def test_gpu_info_with_minimal_data(self):
        """Test GPUInfo with minimal required data."""
        gpu = GPUInfo(
            vendor=GPUVendor.NVIDIA,
            model="Unknown",
            device_id="0",
            pci_bus_id="00:00.0",
            memory_mb=0,
        )

        data = gpu.to_dict()
        assert data["vendor"] == "nvidia"
        assert data["model"] == "Unknown"
        assert data["driver_version"] is None
        assert data["cuda_version"] is None
