import stat
import uuid
from pathlib import Path
from urllib.parse import urlparse, urlunparse

import pytest

import ray
import ray.train.collective
from ray import train
from ray.train import Checkpoint, CheckpointConfig, RunConfig, ScalingConfig
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
from ray.train.torch import TorchTrainer
from ray.train.v2._internal.constants import CHECKPOINT_MANAGER_SNAPSHOT_FILENAME
from ray.train.v2._internal.execution.storage import StorageContext
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
from ray.train.v2.api.exceptions import WorkerGroupError
from ray.train.v2.api.result import Result


def uri_join(base_uri: str, *paths: str) -> str:
    """
    Join a base URI (local or remote) with one or more subpaths.
    Preserves query parameters and scheme.
    """
    parsed = urlparse(base_uri)
    new_path = "/".join([p.strip("/") for p in [parsed.path, *paths] if p])

    # If it's a local path (no scheme), ensure we preserve the leading /
    if not parsed.scheme and not new_path.startswith("/"):
        new_path = "/" + new_path

    return urlunparse(
        (
            parsed.scheme,
            parsed.netloc,
            new_path,
            parsed.params,
            parsed.query,
            parsed.fragment,
        )
    )


def build_dummy_trainer(
    exp_name: str,
    storage_path: str,
    num_iterations: int,
    num_checkpoints: int,
    train_loop_config: dict,
):
    """Build a dummy TorchTrainer for testing purposes."""

    def worker_loop(_config):
        for i in range(num_iterations):
            # Do some random reports in between checkpoints.
            train.report({"metric_a": -100, "metric_b": -100})

            if ray.train.get_context().get_world_rank() == 0:
                with create_dict_checkpoint({"iter": i}) as checkpoint:
                    train.report(
                        metrics={"metric_a": i, "metric_b": -i},
                        checkpoint=checkpoint,
                    )
            else:
                train.report(metrics={"metric_a": i, "metric_b": -i})

        # Ensure that all checkpoints are saved before the RuntimeError
        ray.train.collective.barrier()
        raise RuntimeError()

    trainer = TorchTrainer(
        train_loop_per_worker=worker_loop,
        train_loop_config=train_loop_config,
        scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
        run_config=RunConfig(
            name=exp_name,
            storage_path=storage_path,
            checkpoint_config=CheckpointConfig(
                num_to_keep=num_checkpoints,
                checkpoint_score_attribute="metric_a",
                checkpoint_score_order="max",
            ),
        ),
    )
    return trainer


def test_result_repr():
    """Test that the Result __repr__ function can return a string."""
    res = Result(
        metrics={"iter": 0, "metric": 1.0},
        checkpoint=Checkpoint("/bucket/path/ckpt0"),
        error=None,
        path="/bucket/path",
    )
    assert isinstance(repr(res), str)
    assert "Checkpoint(filesystem=local, path=/bucket/path/ckpt0)" in repr(res)
    assert "metrics={'iter': 0, 'metric': 1.0}" in repr(res)


def test_get_best_checkpoint():
    res = Result(
        metrics={},
        checkpoint=None,
        error=None,
        path="/bucket/path",
        best_checkpoints=[
            (Checkpoint("/bucket/path/ckpt0"), {"iter": 0, "metric": 1.0}),
            (Checkpoint("/bucket/path/ckpt1"), {"iter": 1, "metric": 2.0}),
            (Checkpoint("/bucket/path/ckpt2"), {"iter": 2, "metric": 3.0}),
            (Checkpoint("/bucket/path/ckpt3"), {"iter": 3, "metric": 4.0}),
        ],
    )
    assert (
        res.get_best_checkpoint(metric="metric", mode="max").path
        == "/bucket/path/ckpt3"
    )
    assert (
        res.get_best_checkpoint(metric="metric", mode="min").path
        == "/bucket/path/ckpt0"
    )


def test_get_best_checkpoint_nested_metrics():
    """Test that get_best_checkpoint works with nested metric dictionaries."""
    # Test with nested metric structure
    res = Result(
        metrics={},
        checkpoint=None,
        error=None,
        path="/bucket/path",
        best_checkpoints=[
            (
                Checkpoint("/bucket/path/ckpt0"),
                {
                    "iter": 0,
                    "env_runners": {"episode_return_mean": 100.0, "num_episodes": 10},
                },
            ),
            (
                Checkpoint("/bucket/path/ckpt1"),
                {
                    "iter": 1,
                    "env_runners": {"episode_return_mean": 200.0, "num_episodes": 10},
                },
            ),
            (
                Checkpoint("/bucket/path/ckpt2"),
                {
                    "iter": 2,
                    "env_runners": {"episode_return_mean": 300.0, "num_episodes": 10},
                },
            ),
            (
                Checkpoint("/bucket/path/ckpt3"),
                {
                    "iter": 3,
                    "env_runners": {"episode_return_mean": 400.0, "num_episodes": 10},
                },
            ),
        ],
    )

    # Test max mode with nested metric
    assert (
        res.get_best_checkpoint(
            metric="env_runners/episode_return_mean", mode="max"
        ).path
        == "/bucket/path/ckpt3"
    )

    # Test min mode with nested metric
    assert (
        res.get_best_checkpoint(
            metric="env_runners/episode_return_mean", mode="min"
        ).path
        == "/bucket/path/ckpt0"
    )

    # Test that flat keys still work (backwards compatibility)
    res_flat = Result(
        metrics={},
        checkpoint=None,
        error=None,
        path="/bucket/path",
        best_checkpoints=[
            (
                Checkpoint("/bucket/path/ckpt0"),
                {"iter": 0, "env_runners/episode_return_mean": 100.0},
            ),
            (
                Checkpoint("/bucket/path/ckpt1"),
                {"iter": 1, "env_runners/episode_return_mean": 200.0},
            ),
        ],
    )
    assert (
        res_flat.get_best_checkpoint(
            metric="env_runners/episode_return_mean", mode="max"
        ).path
        == "/bucket/path/ckpt1"
    )


@pytest.mark.parametrize("path_type", ["str", "PathLike"])
@pytest.mark.parametrize("pass_storage_filesystem", [True, False])
@pytest.mark.parametrize("trailing_slash", [False, True])
def test_result_restore(
    ray_start_4_cpus,
    tmp_path,
    path_type,
    pass_storage_filesystem,
    trailing_slash,
):
    """Test Result.from_path functionality similar to v1 test_result_restore."""

    num_iterations = 3
    num_checkpoints = 2

    storage_path = str(tmp_path)

    # Add UUID to ensure test isolation when sharing module-scoped S3 mock
    exp_name = f"test_result_restore_v2-{uuid.uuid4().hex[:8]}"

    trainer = build_dummy_trainer(
        exp_name,
        storage_path,
        num_iterations,
        num_checkpoints,
        train_loop_config={"a": 1, "b": 2},
    )
    with pytest.raises(WorkerGroupError):
        trainer.fit()

    if pass_storage_filesystem:
        storage_context = StorageContext(
            storage_path=storage_path,
            experiment_dir_name=exp_name,
        )

        trial_dir = storage_context.experiment_fs_path
        file_system = storage_context.storage_filesystem
    else:
        trial_dir = uri_join(storage_path, exp_name)
        file_system = None

    # Add trailing slash if parameterized to test that case
    if trailing_slash:
        trial_dir = trial_dir + "/"

    # For PathLike test, only use Path() for local paths, not URIs
    if path_type == "PathLike":
        trial_dir_arg = Path(trial_dir)
    else:
        trial_dir_arg = trial_dir

    result = Result.from_path(
        trial_dir_arg,
        storage_filesystem=file_system,
    )

    assert result.checkpoint
    assert len(result.best_checkpoints) == num_checkpoints

    """
    Top-2 checkpoints with metrics:

                        | iter   | metric_a    metric_b
    checkpoint_000002        2            2          -2
    checkpoint_000001        1            1          -1
    """
    # Check if the checkpoints bounded with correct metrics
    best_ckpt_a = result.get_best_checkpoint(metric="metric_a", mode="max")
    assert load_dict_checkpoint(best_ckpt_a)["iter"] == num_iterations - 1

    best_ckpt_b = result.get_best_checkpoint(metric="metric_b", mode="max")
    assert load_dict_checkpoint(best_ckpt_b)["iter"] == num_iterations - num_checkpoints

    with pytest.raises(RuntimeError, match="Invalid metric name.*"):
        result.get_best_checkpoint(metric="invalid_metric", mode="max")


def test_result_from_path_read_only_storage(
    ray_start_4_cpus,
    tmp_path,
):
    """Reproduces https://github.com/ray-project/ray/issues/64305.

    If checkpoints live on write-once / read-only storage, check that
    ``Result.from_path`` / ``get_best_checkpoint`` are read-only operations.
    """

    def train_fn():
        with create_dict_checkpoint({"iter": 1}) as ckpt:
            ray.train.report({"epoch": 1}, ckpt)
        with create_dict_checkpoint({"iter": 2}) as ckpt:
            ray.train.report({"epoch": 2}, ckpt)

    trainer = DataParallelTrainer(
        train_fn,
        run_config=RunConfig(
            name="test_read_only_storage",
            storage_path=str(tmp_path),
        ),
    )
    trainer.fit()

    # Strip write permissions from the experiment directory, so that any new write fails.
    no_write_mask = ~(stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH)
    experiment_dir = Path(tmp_path) / "test_read_only_storage"
    assert experiment_dir.exists()
    original_modes = {
        path: stat.S_IMODE(path.stat().st_mode)
        for path in [experiment_dir, *experiment_dir.rglob("*")]
    }
    for path in original_modes:
        path.chmod(original_modes[path] & no_write_mask)

    try:
        # Reading the best checkpoint must not require write access.
        result = Result.from_path(str(experiment_dir))
        assert result.checkpoint
        assert len(result.best_checkpoints) == 2

        best_ckpt = result.get_best_checkpoint(metric="epoch", mode="max")
        assert load_dict_checkpoint(best_ckpt)["iter"] == 2
    finally:
        for path, mode in original_modes.items():
            path.chmod(mode)


def test_result_from_path_validation(
    ray_start_4_cpus,
    tmp_path,
):
    """Test that Result.from_path raises RuntimeError when folder or snapshot file doesn't exist."""

    nonexistent_folder = str(tmp_path / "nonexistent_experiment")
    existing_folder = str(tmp_path / "existing_experiment")

    # Test 1: Folder doesn't exist
    with pytest.raises(RuntimeError, match="Experiment folder .* doesn't exist."):
        Result.from_path(nonexistent_folder)

    # Test 2: Folder exists but snapshot file doesn't exist
    Path(existing_folder).mkdir(parents=True, exist_ok=True)

    with pytest.raises(
        RuntimeError,
        match=f"Failed to restore the Result object: {CHECKPOINT_MANAGER_SNAPSHOT_FILENAME} doesn't exist in the experiment folder. Make sure that this is an output directory created "
        "by a Ray Train run.",
    ):
        Result.from_path(existing_folder)


if __name__ == "__main__":
    import sys

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