import pyarrow.fs
import pytest

from ray.train import RunConfig, ScalingConfig


def test_scaling_config_validation():
    assert ScalingConfig(
        num_workers=2, use_gpu=True, resources_per_worker={"CPU": 1}
    ).total_resources == {"CPU": 2, "GPU": 2}

    with pytest.raises(ValueError, match="`use_gpu` is False but `GPU` was found in"):
        ScalingConfig(num_workers=2, use_gpu=False, resources_per_worker={"GPU": 1})

    with pytest.raises(ValueError, match="Cannot specify both"):
        ScalingConfig(num_workers=2, use_gpu=True, use_tpu=True)

    with pytest.raises(
        ValueError,
        match=(
            "If `label_selector` is a list, it must be the same length as "
            "`max_workers`"
        ),
    ):
        ScalingConfig(num_workers=2, label_selector=[{"subcluster": "my_subcluster"}])
    with pytest.raises(
        ValueError,
        match=(
            "If `label_selector` is a list, it must be the same length as "
            "`max_workers`"
        ),
    ):
        ScalingConfig(
            num_workers=(2, 3),
            label_selector=[{"subcluster": "a"}, {"subcluster": "b"}],
        )


def test_label_selector_per_worker():
    # None -> None (no constraint; downstream consumers handle this directly).
    assert ScalingConfig(num_workers=3)._label_selector_per_worker(3) is None

    # Dict -> replicated per worker, decoupled from the original.
    cfg = ScalingConfig(num_workers=2, label_selector={"zone": "a"})
    result = cfg._label_selector_per_worker(2)
    assert result == [{"zone": "a"}, {"zone": "a"}]
    result[0]["zone"] = "b"
    assert cfg.label_selector == {"zone": "a"}

    # List -> sliced to num_workers, decoupled from the original.
    cfg = ScalingConfig(
        num_workers=(1, 3),
        label_selector=[{"a": "1"}, {"a": "2"}, {"a": "3"}],
    )
    assert cfg._label_selector_per_worker(2) == [{"a": "1"}, {"a": "2"}]


def test_scaling_config_accelerator_type():
    scaling_config = ScalingConfig(num_workers=2, use_gpu=True, accelerator_type="A100")
    assert scaling_config.accelerator_type == "A100"
    assert scaling_config._resources_per_worker_not_none == {
        "GPU": 1,
        "accelerator_type:A100": 0.001,
    }
    assert scaling_config.total_resources == {
        "GPU": 2,
        "accelerator_type:A100": 0.002,
    }
    assert scaling_config.additional_resources_per_worker == {
        "accelerator_type:A100": 0.001
    }


def test_scaling_config_tpu_min_workers_multiple():
    with pytest.raises(ValueError, match="min_workers"):
        ScalingConfig(
            num_workers=(1, 2),
            use_tpu=True,
            topology="2x2x2",
            accelerator_type="TPU-V4",
            resources_per_worker={"TPU": 4},
        )


def test_storage_filesystem_repr():
    """Test for https://github.com/ray-project/ray/pull/40851"""
    config = RunConfig(storage_filesystem=pyarrow.fs.S3FileSystem())
    repr(config)


def test_scaling_config_default_workers():
    """Test that num_workers defaults to 1 for non-TPU workloads."""
    config = ScalingConfig()
    assert config.num_workers == 1
    assert config.total_resources == {"CPU": 1}

    config_gpu = ScalingConfig(use_gpu=True)
    assert config_gpu.num_workers == 1
    assert config_gpu.total_resources == {"GPU": 1}


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", "-x", __file__]))