from spelunkai_training.client import build_class_mapping, parse_dataset_version


def test_parse_dataset_version():
    payload = {
        "id": 1,
        "set_id": 2,
        "name": "v1",
        "frames": [
            {
                "frame": {
                    "id": 10, "session_name": "run01", "frame_index": 0,
                    "image_path": "run01_frames/frame_000000.png", "width": 1280, "height": 720,
                },
                "labels": [
                    {"sub_class_id": 5, "x": 1.0, "y": 2.0, "width": 3.0, "height": 4.0},
                ],
            },
        ],
    }

    result = parse_dataset_version(payload)

    assert result.id == 1
    assert result.set_id == 2
    assert result.name == "v1"
    assert len(result.frames) == 1
    frame = result.frames[0]
    assert frame.image_path == "run01_frames/frame_000000.png"
    assert len(frame.labels) == 1
    assert frame.labels[0].sub_class_id == 5


def test_build_class_mapping_is_stable_and_ordered_by_id():
    payload = {
        "main_classes": [
            {"id": 1, "name": "Enemy", "sub_classes": [{"id": 5, "name": "Snake"}, {"id": 2, "name": "Bat"}]},
            {"id": 2, "name": "Hazard", "sub_classes": [{"id": 9, "name": "Spikes"}]},
        ],
    }

    mapping = build_class_mapping(payload)

    assert mapping == {2: 0, 5: 1, 9: 2}