from aitop.core.gpu.base import GPUInfo
from aitop.ui.components.process_metrics import (
    build_accelerator_identities,
    calculate_process_device_breakdown,
    calculate_process_gpu_metrics,
    format_device_breakdown,
    recommended_device_display_limit,
    summarize_process_device_resources,
)


def test_process_metrics_prefers_nvidia_sm_util_over_gpu_fallback() -> None:
    process = {"pid": 100}
    gpus = [
        GPUInfo(
            index=0,
            name="NVIDIA H100",
            utilization=97.0,
            memory_used=40960.0,
            memory_total=81920.0,
            temperature=70.0,
            power_draw=320.0,
            power_limit=350.0,
            processes=[
                {"pid": 100, "name": "trainer", "memory": 20000.0, "sm_util": 43.0},
                {"pid": 101, "name": "trainer", "memory": 20000.0, "sm_util": 40.0},
            ],
        )
    ]

    gpu_percent, vram_percent = calculate_process_gpu_metrics(process, gpus)

    assert gpu_percent == 43.0
    assert round(vram_percent, 1) == 24.4


def test_process_metrics_parses_cu_occupancy_percent_string() -> None:
    process = {"pid": 200}
    gpus = [
        GPUInfo(
            index=1,
            name="AMD MI250",
            utilization=88.0,
            memory_used=10000.0,
            memory_total=20000.0,
            temperature=65.0,
            power_draw=250.0,
            power_limit=300.0,
            processes=[
                {"pid": 200, "name": "worker", "memory": 5000.0, "cu_occupancy": "72%"}
            ],
        )
    ]

    gpu_percent, vram_percent = calculate_process_gpu_metrics(process, gpus)

    assert gpu_percent == 72.0
    assert round(vram_percent, 1) == 25.0


def test_process_metrics_falls_back_to_memory_ratio_without_process_util() -> None:
    process = {"pid": 300}
    gpus = [
        GPUInfo(
            index=0,
            name="NVIDIA A100",
            utilization=90.0,
            memory_used=30000.0,
            memory_total=40000.0,
            temperature=62.0,
            power_draw=280.0,
            power_limit=300.0,
            processes=[
                {"pid": 300, "name": "train-a", "memory": 3000.0},
                {"pid": 301, "name": "train-b", "memory": 7000.0},
            ],
        )
    ]

    gpu_percent, vram_percent = calculate_process_gpu_metrics(process, gpus)

    # 3000/(3000+7000) * 90
    assert round(gpu_percent, 1) == 27.0
    # 3000/40000 * 100
    assert round(vram_percent, 1) == 7.5


def test_process_metrics_normalizes_byte_based_memory_values() -> None:
    process = {"pid": 400}
    gpus = [
        GPUInfo(
            index=0,
            name="AMD GPU",
            utilization=60.0,
            memory_used=0.0,
            memory_total=16000.0,
            temperature=50.0,
            power_draw=100.0,
            power_limit=200.0,
            processes=[
                # 2 GiB in bytes
                {"pid": 400, "name": "job", "memory": 2 * 1024 * 1024 * 1024},
            ],
        )
    ]

    gpu_percent, vram_percent = calculate_process_gpu_metrics(process, gpus)

    # No direct per-process util, single process fallback uses full GPU util
    assert round(gpu_percent, 1) == 60.0
    # 2048MB / 16000MB * 100
    assert round(vram_percent, 1) == 12.8


def test_process_metrics_normalizes_kib_memory_values() -> None:
    process = {"pid": 450}
    gpus = [
        GPUInfo(
            index=0,
            name="NVIDIA GPU",
            utilization=80.0,
            memory_used=0.0,
            memory_total=16000.0,
            temperature=0.0,
            power_draw=0.0,
            power_limit=0.0,
            processes=[
                # 2 GiB in KiB
                {"pid": 450, "name": "job", "memory": 2 * 1024 * 1024},
            ],
        )
    ]

    gpu_percent, vram_percent = calculate_process_gpu_metrics(process, gpus)

    assert round(gpu_percent, 1) == 80.0
    # 2048MB / 16000MB * 100
    assert round(vram_percent, 1) == 12.8


def test_process_metrics_clamps_overflowing_percentages() -> None:
    process = {"pid": 475}
    gpus = [
        GPUInfo(
            index=0,
            name="NVIDIA GPU",
            utilization=95.0,
            memory_used=0.0,
            memory_total=16000.0,
            temperature=0.0,
            power_draw=0.0,
            power_limit=0.0,
            processes=[
                {"pid": 475, "name": "job", "memory": 25000.0, "sm_util": 2500.0},
            ],
        )
    ]

    gpu_percent, vram_percent = calculate_process_gpu_metrics(process, gpus)

    assert gpu_percent == 100.0
    assert vram_percent == 100.0


def test_process_metrics_handles_npu_unknown_total_memory() -> None:
    process = {"pid": 500}
    gpus = [
        GPUInfo(
            index=0,
            name="AMD NPU",
            utilization=80.0,
            memory_used=0.0,
            memory_total=0.0,
            temperature=0.0,
            power_draw=0.0,
            power_limit=0.0,
            processes=[
                {"pid": 500, "name": "rank0", "memory": 3.0},
                {"pid": 501, "name": "rank1", "memory": 1.0},
            ],
        )
    ]

    gpu_percent, vram_percent = calculate_process_gpu_metrics(process, gpus)

    # 3/(3+1) * 80
    assert round(gpu_percent, 1) == 60.0
    # Unknown total memory should remain 0 instead of inventing a denominator
    assert vram_percent == 0.0


def test_process_device_breakdown_sorts_and_limits_top_three() -> None:
    process = {"pid": 600}
    accelerators = [
        (
            GPUInfo(
                index=0,
                name="NVIDIA A100",
                utilization=95.0,
                memory_used=0.0,
                memory_total=80000.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[{"pid": 600, "memory": 8000.0, "sm_util": 41.0}],
            ),
            "nvidia",
        ),
        (
            GPUInfo(
                index=1,
                name="NVIDIA H100",
                utilization=90.0,
                memory_used=0.0,
                memory_total=80000.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[{"pid": 600, "memory": 12000.0, "sm_util": 70.0}],
            ),
            "nvidia",
        ),
        (
            GPUInfo(
                index=0,
                name="Intel AI Boost",
                utilization=85.0,
                memory_used=0.0,
                memory_total=0.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[{"pid": 600, "memory": 1.0, "gpu_util": 80.0}],
            ),
            "intel_npu",
        ),
        (
            GPUInfo(
                index=2,
                name="NVIDIA L40S",
                utilization=60.0,
                memory_used=0.0,
                memory_total=80000.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[{"pid": 600, "memory": 2000.0, "sm_util": 20.0}],
            ),
            "nvidia",
        ),
    ]

    breakdown = calculate_process_device_breakdown(process, accelerators, max_devices=3)

    assert [entry["device"] for entry in breakdown] == ["NPU0", "GPU1", "GPU0"]
    assert [entry["utilization"] for entry in breakdown] == [80.0, 70.0, 41.0]


def test_build_accelerator_identities_scopes_global_indexes_per_class() -> None:
    accelerators = [
        (
            GPUInfo(
                index=0,
                name="NVIDIA A100",
                utilization=0.0,
                memory_used=0.0,
                memory_total=0.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[],
            ),
            "nvidia",
        ),
        (
            GPUInfo(
                index=0,
                name="AMD MI300",
                utilization=0.0,
                memory_used=0.0,
                memory_total=0.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[],
            ),
            "amd",
        ),
        (
            GPUInfo(
                index=0,
                name="Intel AI Boost",
                utilization=0.0,
                memory_used=0.0,
                memory_total=0.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[],
            ),
            "intel_npu",
        ),
        (
            GPUInfo(
                index=0,
                name="AMD NPU",
                utilization=0.0,
                memory_used=0.0,
                memory_total=0.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[],
            ),
            "amd_npu",
        ),
    ]

    identities = build_accelerator_identities(accelerators)

    assert [item["global_label"] for item in identities] == [
        "GPU0",
        "GPU1",
        "NPU0",
        "NPU1",
    ]
    assert [item["local_label"] for item in identities] == [
        "GPU0",
        "GPU0",
        "NPU0",
        "NPU0",
    ]


def test_recommended_device_display_limit_scales_with_hardware_count() -> None:
    assert recommended_device_display_limit(0) == 1
    assert recommended_device_display_limit(2) == 2
    assert recommended_device_display_limit(6) == 3
    assert recommended_device_display_limit(16) == 2


def test_format_device_breakdown_compacts_with_overflow_suffix() -> None:
    formatted = format_device_breakdown(
        [
            {"device": "NPU0", "utilization": 61.0, "memory_mb": 512.0},
            {"device": "GPU0", "utilization": 43.0, "memory_mb": 2048.0},
            {"device": "GPU1", "utilization": 22.0, "memory_mb": 1024.0},
        ],
        max_devices=2,
        include_utilization=True,
        include_memory=False,
    )

    assert formatted == "NPU0:61%,GPU0:43%,+1"


def test_format_device_breakdown_clamps_percent_values() -> None:
    formatted = format_device_breakdown(
        [
            {"device": "GPU0", "utilization": 1450.0, "memory_mb": 512.0},
            {"device": "GPU1", "utilization": -5.0, "memory_mb": 256.0},
        ],
        max_devices=2,
        include_utilization=True,
        include_memory=False,
    )

    assert formatted == "GPU0:100%,GPU1:0%"


def test_format_device_breakdown_formats_memory_only() -> None:
    formatted = format_device_breakdown(
        [
            {"device": "GPU0", "utilization": 12.0, "memory_mb": 512.0},
            {"device": "GPU1", "utilization": 9.0, "memory_mb": 2048.0},
        ],
        max_devices=2,
        include_utilization=False,
        include_memory=True,
    )

    assert formatted == "GPU0:512M,GPU1:2.0G"


def test_summarize_process_device_resources_groups_usage() -> None:
    processes = [{"pid": 100}, {"pid": 200}]
    accelerators = [
        (
            GPUInfo(
                index=0,
                name="NVIDIA A100",
                utilization=90.0,
                memory_used=0.0,
                memory_total=80000.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[
                    {"pid": 100, "memory": 16000.0, "sm_util": 60.0},
                    {"pid": 200, "memory": 8000.0, "sm_util": 30.0},
                ],
            ),
            "nvidia",
        ),
        (
            GPUInfo(
                index=0,
                name="Intel AI Boost",
                utilization=70.0,
                memory_used=0.0,
                memory_total=0.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[{"pid": 200, "memory": 1.0, "gpu_util": 40.0}],
            ),
            "intel_npu",
        ),
    ]

    summary = summarize_process_device_resources(processes, accelerators)

    assert [entry["device"] for entry in summary] == ["GPU0", "NPU0"]
    assert summary[0]["process_count"] == 2
    assert round(summary[0]["memory_used_mb"], 1) == 24000.0


def test_summarize_process_device_resources_separates_gpu_vendor_index_collisions() -> None:
    processes = [{"pid": 111}, {"pid": 222}]
    accelerators = [
        (
            GPUInfo(
                index=0,
                name="NVIDIA A100",
                utilization=90.0,
                memory_used=0.0,
                memory_total=80000.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[
                    {"pid": 111, "memory": 16000.0, "sm_util": 60.0},
                ],
            ),
            "nvidia",
        ),
        (
            GPUInfo(
                index=0,
                name="AMD MI300",
                utilization=80.0,
                memory_used=0.0,
                memory_total=64000.0,
                temperature=0.0,
                power_draw=0.0,
                power_limit=0.0,
                processes=[
                    {"pid": 222, "memory": 8000.0, "gpu_util": 40.0},
                ],
            ),
            "amd",
        ),
    ]

    summary = summarize_process_device_resources(processes, accelerators)

    assert len(summary) == 2
    assert {entry["device"] for entry in summary} == {"GPU0", "GPU1"}
    assert all(entry["process_count"] == 1 for entry in summary)
