"""Tests for the EngineProvisioner.

Validates:
- VRAM-based model selection
- Provisioning flow (select + download)
- GPU driver recommendations
"""

from __future__ import annotations

from unittest.mock import patch

import pytest

from llmhosts.discovery.models import GPUInfo, HardwareProfile
from llmhosts.discovery.provisioner import EngineProvisioner, _select_model_for_hardware

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _hw_with_vram(vram_mb: int, ram_gb: float = 32.0) -> HardwareProfile:
    """Build a HardwareProfile with the specified VRAM (via GPUInfo list)."""
    gpus = [GPUInfo(name="Test GPU", vram_total_mb=vram_mb, vram_free_mb=vram_mb)] if vram_mb > 0 else []
    return HardwareProfile(ram_total_gb=ram_gb, gpus=gpus)


# ---------------------------------------------------------------------------
# VRAM-based model selection
# ---------------------------------------------------------------------------


@pytest.mark.unit
@pytest.mark.parametrize(
    "vram_mb,expected_model",
    [
        (0, "qwen3:1.7b"),  # No GPU -- RAM/2 = 16GB -> gpt-oss:20b
        (2048, "qwen3:1.7b"),  # 2 GB -- ultra-light
        (4096, "qwen3:4b"),  # 4 GB -- small discrete / iGPU
        (6144, "qwen3:4b"),  # 6 GB -- upper small tier
        (7168, "qwen3:8b"),  # 7 GB -- RTX 3060 8 GB
        (8192, "qwen3:8b"),  # 8 GB -- RTX 4060
        (10240, "qwen3:14b"),  # 10 GB -- RTX 3080 10 GB
        (12288, "qwen3:14b"),  # 12 GB -- RTX 3060 12 GB
        (16384, "gpt-oss:20b"),  # 16 GB -- RTX 4080
        (18432, "gpt-oss:20b"),  # 18 GB -- RTX 4090 Mobile
        (24576, "qwen3:32b"),  # 24 GB -- RTX 4090
        (32768, "qwen3:32b"),  # 32 GB -- RTX 5090
        (49152, "llama3.3:70b"),  # 48 GB -- dual GPU / A100
    ],
)
def test_model_selection_by_vram(vram_mb: int, expected_model: str) -> None:
    """Model selection should scale with available VRAM."""
    if vram_mb == 0:
        # No GPU: effective = ram/2 = 32/2 = 16 -> should be gpt-oss:20b
        expected_model = "gpt-oss:20b"
    hw = _hw_with_vram(vram_mb)
    model_name, _ = _select_model_for_hardware(hw)
    assert model_name == expected_model


@pytest.mark.unit
def test_cpu_only_uses_half_ram() -> None:
    """Without GPU, effective VRAM = RAM/2."""
    # 16GB RAM -> 8GB effective -> should pick qwen3:8b
    hw = _hw_with_vram(0, ram_gb=16.0)
    model_name, _ = _select_model_for_hardware(hw)
    assert model_name == "qwen3:8b"


@pytest.mark.unit
def test_minimal_ram_gets_tiny_model() -> None:
    """With very little RAM and no GPU, should get the smallest model."""
    hw = _hw_with_vram(0, ram_gb=4.0)
    model_name, _ = _select_model_for_hardware(hw)
    # 4GB RAM -> 2GB effective -> should pick qwen3:1.7b
    assert model_name == "qwen3:1.7b"


# ---------------------------------------------------------------------------
# Provisioning flow
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_provision_success_flow() -> None:
    """Full provision flow: select + download should succeed."""
    status_messages: list[str] = []
    provisioner = EngineProvisioner(on_status=status_messages.append)

    hw = _hw_with_vram(8192, ram_gb=16.0)

    with patch.object(provisioner, "_download_model", return_value=True):
        result = await provisioner.provision(hw)

    assert result.success is True
    assert result.engine_installed is True
    assert result.model_pulled == "qwen3:8b"
    assert result.install_method == "core"
    assert result.duration_seconds >= 0


@pytest.mark.unit
async def test_provision_download_failure() -> None:
    """When model download fails, result should indicate failure."""
    provisioner = EngineProvisioner()

    hw = _hw_with_vram(8192, ram_gb=16.0)

    with patch.object(provisioner, "_download_model", return_value=False):
        result = await provisioner.provision(hw)

    assert result.success is False
    assert result.engine_installed is True  # Core is always available


# ---------------------------------------------------------------------------
# GPU driver recommendations
# ---------------------------------------------------------------------------


@pytest.mark.unit
def test_gpu_driver_recs_linux_no_gpu() -> None:
    """On Linux without GPU, should recommend NVIDIA and AMD drivers."""
    hw = _hw_with_vram(0, ram_gb=16.0)

    with patch("platform.system", return_value="Linux"):
        recs = EngineProvisioner.gpu_driver_recommendations(hw)

    assert len(recs) >= 1
    combined = " ".join(recs).lower()
    assert "nvidia" in combined or "amd" in combined


@pytest.mark.unit
def test_gpu_driver_recs_with_gpu() -> None:
    """With GPU already detected, should return no recommendations."""
    hw = _hw_with_vram(8192, ram_gb=16.0)
    recs = EngineProvisioner.gpu_driver_recommendations(hw)
    assert len(recs) == 0


@pytest.mark.unit
def test_suggest_and_pull() -> None:
    """suggest_and_pull should suggest based on hardware."""
    hw = _hw_with_vram(24576, ram_gb=32.0)

    model_name, _ = _select_model_for_hardware(hw)
    assert model_name == "qwen3:32b"
