from __future__ import annotations

from hytop.core.sorting import next_sort_field_index, sort_with_missing_last
from hytop.gpu.sort import sort_gpu_keys_grouped
from hytop.net.sort import sort_net_keys_grouped


class TestGpuSortingHelpers:
    def test_grouped_default_host_then_gpu(self):
        keys = {("node02", 1), ("node01", 2), ("node01", 0)}
        ordered = sort_gpu_keys_grouped(keys, hosts=["node01", "node02"])
        assert ordered == [("node01", 0), ("node01", 2), ("node02", 1)]

    def test_metric_sort_keeps_missing_last(self):
        base = [("node01", 0), ("node01", 1), ("node02", 0)]
        values = {("node01", 0): 50.0, ("node02", 0): 10.0}
        ordered = sort_with_missing_last(base, lambda key: values.get(key), desc=True)
        assert ordered == [("node01", 0), ("node02", 0), ("node01", 1)]


class TestNetSortingHelpers:
    def test_grouped_default_host_mode_nic(self):
        keys = {
            ("node02", "ib:mlx5_1/p1"),
            ("node01", "ib:mlx5_0/p1"),
            ("node01", "eth:p6p1"),
            ("node01", "eth:p14p1"),
        }
        ordered = sort_net_keys_grouped(keys, hosts=["node01", "node02"])
        assert ordered == [
            ("node01", "eth:p14p1"),
            ("node01", "eth:p6p1"),
            ("node01", "ib:mlx5_0/p1"),
            ("node02", "ib:mlx5_1/p1"),
        ]

    def test_metric_sort_keeps_missing_last(self):
        base = [("node01", "eth:p6p1"), ("node01", "ib:mlx5_0/p1"), ("node02", "eth:p14p1")]
        values = {("node01", "ib:mlx5_0/p1"): 30.0, ("node02", "eth:p14p1"): 10.0}
        ordered = sort_with_missing_last(base, lambda key: values.get(key), desc=False)
        assert ordered == [
            ("node02", "eth:p14p1"),
            ("node01", "ib:mlx5_0/p1"),
            ("node01", "eth:p6p1"),
        ]


class TestSortStateHelpers:
    def test_next_sort_field_index_cycles(self):
        assert next_sort_field_index(0, 3) == 1
        assert next_sort_field_index(2, 3) == 0
        assert next_sort_field_index(0, 0) == 0
