from dataclasses import dataclass
from typing import Any

import aitop.data_collector as data_collector_module
from aitop.data_collector import DataCollector


class _FakeAIMonitor:
    def get_ai_processes(
        self, gpu_processes: list[dict[str, Any]], force_refresh: bool
    ) -> list[dict[str, Any]]:
        del gpu_processes, force_refresh
        return []


class _FakeMemoryMonitor:
    def get_memory_stats(self) -> dict[str, Any]:
        return {"percent": 0.0}

    def get_memory_by_type(self) -> dict[str, float]:
        return {"cached": 0.0}


@dataclass
class _FakeGPU:
    index: int
    processes: list[dict[str, Any]]


@dataclass
class _QuickMetricsGPU:
    index: int
    utilization: float
    memory_used: float
    memory_total: float


class _HealthyMonitor:
    def __init__(self, gpus: list[_FakeGPU]) -> None:
        self._gpus = gpus

    def get_gpu_info(self) -> list[_FakeGPU]:
        return self._gpus


class _FailingMonitor:
    def get_gpu_info(self) -> list[_FakeGPU]:
        raise RuntimeError("vendor monitor failed")


def _make_collector(monkeypatch: Any) -> DataCollector:
    monkeypatch.setattr(
        data_collector_module.GPUMonitorFactory,
        "create_monitors",
        lambda: [],
    )
    monkeypatch.setattr(
        data_collector_module.GPUMonitorFactory,
        "detect_vendors",
        lambda: [],
    )
    monkeypatch.setattr(
        data_collector_module.GPUMonitorFactory,
        "detect_dependency_warnings",
        lambda: [],
    )
    monkeypatch.setattr(data_collector_module, "AIProcessMonitor", _FakeAIMonitor)
    monkeypatch.setattr(
        data_collector_module,
        "SystemMemoryMonitor",
        _FakeMemoryMonitor,
    )
    return DataCollector(max_workers=1)


def test_collect_full_gpu_info_continues_when_one_vendor_fails(
    monkeypatch: Any,
) -> None:
    collector = _make_collector(monkeypatch)
    gpu = _FakeGPU(
        index=2,
        processes=[{"pid": 99, "name": "python train.py", "memory": 1024.0}],
    )

    collector.gpu_monitors = [_HealthyMonitor([gpu]), _FailingMonitor()]
    collector.vendors = ["nvidia", "amd"]

    gpu_info, gpu_processes = collector._collect_full_gpu_info()

    assert gpu_info == [(gpu, "nvidia")]
    assert gpu_processes == [
        {
            "pid": 99,
            "name": "python train.py",
            "memory": 1024.0,
            "gpu_index": 2,
            "gpu_vendor": "nvidia",
        }
    ]
    assert "gpu_index" not in gpu.processes[0]
    assert any(
        error["component"] == "gpu_amd" and "vendor monitor failed" in error["error"]
        for error in collector._cache["errors"]
    )


def test_collect_full_gpu_info_returns_empty_on_critical_iteration_error(
    monkeypatch: Any,
) -> None:
    collector = _make_collector(monkeypatch)
    collector.gpu_monitors = None  # type: ignore[assignment]
    collector.vendors = ["nvidia"]

    gpu_info, gpu_processes = collector._collect_full_gpu_info()

    assert gpu_info == []
    assert gpu_processes == []
    assert any(
        error["component"] == "gpu_collection" for error in collector._cache["errors"]
    )


def test_update_gpu_quick_data_updates_only_present_metrics(
    monkeypatch: Any,
) -> None:
    collector = _make_collector(monkeypatch)
    gpu = _QuickMetricsGPU(
        index=0,
        utilization=12.0,
        memory_used=256.0,
        memory_total=1024.0,
    )
    collector._cache["gpu_info"] = [(gpu, "nvidia")]

    collector._update_gpu_quick_data({"nvidia:0": {"utilization": 77.0}})

    assert gpu.utilization == 77.0
    assert gpu.memory_used == 256.0
    assert gpu.memory_total == 1024.0


def test_update_gpu_quick_data_ignores_unknown_gpu_keys(
    monkeypatch: Any,
) -> None:
    collector = _make_collector(monkeypatch)
    gpu = _QuickMetricsGPU(
        index=1,
        utilization=10.0,
        memory_used=128.0,
        memory_total=2048.0,
    )
    collector._cache["gpu_info"] = [(gpu, "amd")]

    collector._update_gpu_quick_data({"nvidia:0": {"utilization": 44.0}})

    assert gpu.utilization == 10.0
    assert gpu.memory_used == 128.0
    assert gpu.memory_total == 2048.0
