import importlib
import sys

import pytest

import ray.train
from ray.train import FailureConfig, RunConfig, ScalingConfig
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer


@pytest.mark.parametrize(
    "operation, raise_error",
    [
        (lambda: FailureConfig(fail_fast=True), True),
        (lambda: RunConfig(verbose=0), True),
        (lambda: FailureConfig(), False),
        (lambda: RunConfig(), False),
        (lambda: ScalingConfig(trainer_resources={"CPU": 1}), True),
        (lambda: ScalingConfig(), False),
    ],
)
def test_api_configs(operation, raise_error):
    if raise_error:
        with pytest.raises(DeprecationWarning):
            operation()
    else:
        try:
            operation()
        except Exception as e:
            pytest.fail(f"Default Operation raised an exception: {e}")


def test_run_config_default_failure_config():
    """Test that RunConfig creates a default FailureConfig from v2 API, not v1."""
    # Import the v2 FailureConfig and v1 FailureConfig for comparison
    from ray.train.v2.api.config import FailureConfig as FailureConfigV2

    # Create a RunConfig without specifying failure_config
    run_config = RunConfig()

    # Verify that the default failure_config is the v2 version
    assert run_config.failure_config is not None
    assert isinstance(run_config.failure_config, FailureConfigV2)
    assert type(run_config.failure_config) is FailureConfigV2

    # Verify that explicitly passing None also creates v2 FailureConfig
    run_config_explicit_none = RunConfig(failure_config=None)
    assert run_config_explicit_none.failure_config is not None
    assert isinstance(run_config_explicit_none.failure_config, FailureConfigV2)
    assert type(run_config_explicit_none.failure_config) is FailureConfigV2


def test_scaling_config_total_resources():
    """Test the patched scaling config total resources calculation."""
    num_workers = 2
    num_cpus_per_worker = 1
    num_gpus_per_worker = 1

    scaling_config = ScalingConfig(
        num_workers=num_workers,
        use_gpu=True,
        resources_per_worker={"CPU": num_cpus_per_worker, "GPU": num_gpus_per_worker},
    )
    scaling_config.total_resources == {
        "CPU": num_workers * num_cpus_per_worker,
        "GPU": num_workers * num_gpus_per_worker,
    }


def test_trainer_restore():
    with pytest.raises(DeprecationWarning):
        DataParallelTrainer.restore("dummy")

    with pytest.raises(DeprecationWarning):
        DataParallelTrainer.can_restore("dummy")


def test_serialized_imports(ray_start_4_cpus):
    """Check that captured imports are deserialized properly without circular imports."""

    from ray.train.lightgbm import LightGBMTrainer
    from ray.train.torch import TorchTrainer
    from ray.train.xgboost import XGBoostTrainer

    if sys.version_info < (3, 12):
        from ray.train.tensorflow import TensorflowTrainer
    else:
        TensorflowTrainer = None

    @ray.remote
    def dummy_task():
        _ = (TorchTrainer, TensorflowTrainer, XGBoostTrainer, LightGBMTrainer)

    ray.get(dummy_task.remote())


def test_v1_config_validation():
    """Test that V1 configs raise an error when V2 is enabled."""
    import ray.air

    with pytest.raises(ValueError, match="ray.train.ScalingConfig"):
        DataParallelTrainer(lambda: None, scaling_config=ray.air.ScalingConfig())

    with pytest.raises(ValueError, match="ray.train.RunConfig"):
        DataParallelTrainer(lambda: None, run_config=ray.air.RunConfig())

    with pytest.raises(ValueError, match="ray.train.FailureConfig"):
        DataParallelTrainer(
            lambda: None,
            run_config=ray.train.RunConfig(failure_config=ray.air.FailureConfig()),
        )


@pytest.mark.parametrize("env_v2_enabled", [False, True])
def test_train_v2_import(monkeypatch, env_v2_enabled):
    monkeypatch.setenv("RAY_TRAIN_V2_ENABLED", str(int(env_v2_enabled)))

    # Load from the public `ray.train` module
    # isort: off
    importlib.reload(ray.train)
    from ray.train import FailureConfig, Result, RunConfig

    # isort: on

    # Import from the absolute module paths as references
    from ray.train.v2.api.config import (
        FailureConfig as FailureConfigV2,
        RunConfig as RunConfigV2,
    )
    from ray.train.v2.api.result import Result as ResultV2

    if env_v2_enabled:
        assert RunConfig is RunConfigV2
        assert FailureConfig is FailureConfigV2
        assert Result is ResultV2
    else:
        assert RunConfig is not RunConfigV2
        assert FailureConfig is not FailureConfigV2
        assert Result is not ResultV2


@pytest.mark.parametrize("env_v2_enabled", [False, True])
def test_report_callback_v2_only_arguments(monkeypatch, env_v2_enabled):
    monkeypatch.setenv("RAY_TRAIN_V2_ENABLED", str(int(env_v2_enabled)))
    import ray.train.lightning._lightning_utils

    if env_v2_enabled:
        from ray.train.v2.api.report_config import CheckpointUploadMode
        from ray.train.v2.api.validation_config import ValidationConfig

        def validation_fn(checkpoint):
            return {"val_score": 1}

        def train_fn():
            ray.train.lightning._lightning_utils.RayTrainReportCallback(
                checkpoint_upload_mode=CheckpointUploadMode.SYNC, validation=True
            )

        trainer = DataParallelTrainer(
            train_fn,
            validation_config=ValidationConfig(fn=validation_fn),
        )
        trainer.fit()
    else:
        with pytest.raises(
            ValueError,
            match="`checkpoint_upload_mode` is only supported in Ray Train v2",
        ):
            ray.train.lightning._lightning_utils.RayTrainReportCallback(
                checkpoint_upload_mode="anything"
            )
        with pytest.raises(
            ValueError, match="`validation` is only supported in Ray Train v2"
        ):
            ray.train.lightning._lightning_utils.RayTrainReportCallback(
                validation="anything"
            )


if __name__ == "__main__":
    import sys

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