import logging
import os
import pickle
import re
import tempfile
import time
import uuid
from contextlib import contextmanager
from pathlib import Path
from typing import List, Optional, Tuple

import pyarrow.fs
import pytest

import ray
from ray import train, tune
from ray._common.test_utils import simulate_s3_bucket
from ray.air._internal.uri_utils import URI
from ray.air.constants import EXPR_RESULT_FILE
from ray.train._checkpoint import Checkpoint
from ray.train._internal.storage import (
    StorageContext,
    _delete_fs_path,
    _download_from_fs_path,
)
from ray.train.base_trainer import TrainingFailedError
from ray.train.data_parallel_trainer import DataParallelTrainer
from ray.tune.trainable.trainable import _DICT_CHECKPOINT_FILE_NAME


class TestConstants:
    NUM_ITERATIONS = 6  # == num_checkpoints == num_artifacts
    NUM_TRIALS = 2
    NUM_WORKERS = 3

    SCORE_KEY = "score"


@contextmanager
def mock_s3_bucket_uri():
    port = 5002
    region = "us-west-2"
    with simulate_s3_bucket(port=port, region=region) as s3_uri:
        import boto3

        s3 = boto3.client(
            "s3", region_name=region, endpoint_url=f"http://localhost:{port}"
        )
        # Bucket name will be autogenerated/unique per test
        bucket_name = URI(s3_uri).name
        s3.create_bucket(
            Bucket=bucket_name,
            CreateBucketConfiguration={"LocationConstraint": region},
        )
        # Disable server HTTP request logging
        logging.getLogger("werkzeug").setLevel(logging.WARNING)
        yield URI(s3_uri)
        logging.getLogger("werkzeug").setLevel(logging.INFO)


@contextmanager
def dummy_context_manager(*args, **kwargs):
    yield "dummy value"


@pytest.fixture(autouse=True, scope="module")
def ray_start_4_cpus():
    ray.init(num_cpus=4)
    yield
    ray.shutdown()


def _create_mock_custom_fs(custom_fs_root_dir: Path) -> pyarrow.fs.FileSystem:
    from fsspec.implementations.dirfs import DirFileSystem
    from fsspec.implementations.local import LocalFileSystem

    custom_fs_root_dir.mkdir(parents=True, exist_ok=True)
    storage_filesystem = pyarrow.fs.PyFileSystem(
        pyarrow.fs.FSSpecHandler(
            DirFileSystem(path=str(custom_fs_root_dir), fs=LocalFileSystem())
        )
    )
    return storage_filesystem


@contextmanager
def _resolve_storage_type(
    storage_path_type: str, tmp_path: Path
) -> Tuple[str, Optional[pyarrow.fs.FileSystem]]:
    storage_path, storage_filesystem = None, None

    context_manager = (
        mock_s3_bucket_uri if storage_path_type == "cloud" else dummy_context_manager
    )

    with context_manager() as cloud_storage_path:
        if storage_path_type == "nfs":
            storage_path = str(tmp_path / "fake_nfs")
        elif storage_path_type == "cloud":
            storage_path = str(cloud_storage_path)
        elif storage_path_type == "custom_fs":
            storage_path = "mock_bucket"
            storage_filesystem = _create_mock_custom_fs(tmp_path / "custom_fs")

        yield storage_path, storage_filesystem


def _get_local_inspect_dir(
    root_local_path: Path,
    storage_path: str,
    storage_filesystem: Optional[pyarrow.fs.FileSystem],
    storage_local_path: Path = None,
) -> Tuple[Path, str]:
    """Downloads the storage path -> local dir for inspecting contents.

    Args:
        root_local_path: Local directory to use as the inspect root.
        storage_path: The storage path or URI to download from.
        storage_filesystem: Optional custom filesystem to use.
        storage_local_path: Local path that ``storage_path`` mirrors on disk
            when no remote storage is configured.

    Returns:
        Tuple: (local_inspect_dir, storage_fs_path), where storage_fs_path
            is the path to the storage path on the filesystem (e.g., prefix stripped).
            This is used to check the correctness of paths returned from `Result`'s,
            since URIs are hard to do comparisons with.
    """
    local_inspect_dir = root_local_path / "inspect"
    if storage_path:
        if storage_filesystem:
            fs, storage_fs_path = storage_filesystem, storage_path
        else:
            fs, storage_fs_path = pyarrow.fs.FileSystem.from_uri(storage_path)
        _download_from_fs_path(
            fs=fs, fs_path=storage_fs_path, local_path=str(local_inspect_dir)
        )
    else:
        fs, storage_fs_path = pyarrow.fs.LocalFileSystem(), str(storage_local_path)
        local_inspect_dir = storage_local_path

    return local_inspect_dir, storage_fs_path


def _get_checkpoint_index(checkpoint_dir_name: str) -> int:
    """Gets the checkpoint index from the checkpoint directory name."""
    return int(checkpoint_dir_name.split("_")[-1])


def _create_checkpoint_shard_filename(rank_str: str) -> str:
    return f"checkpoint_shard-rank={rank_str}.pkl"


def _get_checkpoint_shard_rank(checkpoint_shard_filename: str) -> int:
    """Get the checkpoint shard rank from the filename."""
    pattern = _create_checkpoint_shard_filename(r"(\d+)")
    match = re.search(pattern, checkpoint_shard_filename)
    assert match
    return int(match.group(1))


def train_fn(config):
    in_trainer = config.get("in_trainer", False)
    if in_trainer:
        from ray.train._internal.session import _TrainSession, get_session

        train_session = get_session()

        assert isinstance(train_session, _TrainSession)
        assert train_session.storage
        assert train_session.storage.checkpoint_fs_path

        # Check that the working dir for each worker is the shared trial dir.
        assert (
            Path.cwd() == Path(train_session.storage.trial_working_directory).resolve()
        )

    start = 0

    checkpoint = train.get_checkpoint()
    if checkpoint:
        custom_restore_fn = config.get("custom_restore_fn")
        if custom_restore_fn:
            state = custom_restore_fn(checkpoint)
        else:
            with checkpoint.as_directory() as checkpoint_dir:
                with open(os.path.join(checkpoint_dir, "checkpoint.pkl"), "rb") as f:
                    state = pickle.load(f)
        print("Loaded back state from checkpoint:", state)
        start = state["iter"] + 1

    for i in range(start, config.get("num_iterations", 5)):
        time.sleep(config.get("time_per_iter", 0.25))

        metrics = {"iter": i, TestConstants.SCORE_KEY: i}

        # Save an artifact in the local trial dir.
        rank = train.get_context().get_world_rank()
        artifact_file_name = (
            f"artifact-rank={rank}-iter={i}.txt"
            if in_trainer
            else f"artifact-iter={i}.txt"
        )
        with open(artifact_file_name, "w") as f:
            f.write(f"{i}")

        if in_trainer and train.get_context().get_world_rank() in config.get(
            "no_checkpoint_ranks", []
        ):
            train.report(metrics)
        else:
            with tempfile.TemporaryDirectory() as temp_dir:
                with open(os.path.join(temp_dir, "checkpoint.pkl"), "wb") as f:
                    pickle.dump({"iter": i}, f)

                if in_trainer:
                    checkpoint_file_name = _create_checkpoint_shard_filename(str(rank))
                    with open(os.path.join(temp_dir, checkpoint_file_name), "wb") as f:
                        pickle.dump({"iter": i}, f)

                with config.get("custom_save_fn", dummy_context_manager)(temp_dir):
                    train.report(
                        metrics, checkpoint=Checkpoint.from_directory(temp_dir)
                    )
                # `train.report` should not have deleted this!
                assert os.path.exists(temp_dir)

        if i in config.get("fail_iters", []):
            raise RuntimeError(f"Failing on iter={i}!!")


class ClassTrainable(tune.Trainable):
    """Implement (almost) the same thing as `train_fn` but as a class."""

    def setup(self, config):
        # Save some markers in the trial dir.
        tmp_path = config.get("tmp_path")
        self.fail_markers = {
            i: tmp_path / f"fail_marker_{self.trial_id}_iter={i}"
            for i in config.get("fail_iters", [])
        }
        setup_marker = tmp_path / f"setup_marker_{self.trial_id}"
        if not setup_marker.exists():
            for marker in self.fail_markers.values():
                marker.touch()
            setup_marker.touch()

        self.save_as_dict = config.get("save_checkpoint_as_dict", False)

    def step(self) -> dict:
        if self.iteration in self.fail_markers:
            marker = self.fail_markers[self.iteration]
            if marker.exists():
                marker.unlink()
                raise RuntimeError(f"Failing on iter={self.iteration}")

        # Save an artifact in the local trial dir.
        artifact_file_name = f"artifact-iter={self.iteration}.txt"
        with open(artifact_file_name, "w") as f:
            f.write(f"{self.iteration}")

        return {
            "score": 1,
            "done": self.iteration >= self.config.get("num_iterations") - 1,
            "should_checkpoint": True,
        }

    def save_checkpoint(self, temp_checkpoint_dir) -> str:
        if self.save_as_dict:
            return {"dummy": "data"}
        (Path(temp_checkpoint_dir) / "checkpoint.pkl").write_text("dummy")
        return temp_checkpoint_dir

    def load_checkpoint(self, checkpoint_dict_or_path):
        print("Loading state from:", checkpoint_dict_or_path)
        print("At iteration =", self.iteration)
        if self.save_as_dict:
            assert checkpoint_dict_or_path == {"dummy": "data"}
        else:
            assert (
                Path(checkpoint_dict_or_path) / "checkpoint.pkl"
            ).read_text() == "dummy"


def _resume_from_checkpoint(
    checkpoint: Checkpoint,
    expected_state: dict,
    storage_path: Optional[str] = None,
    storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
):
    print(f"\nStarting run with `resume_from_checkpoint`: {checkpoint}\n")

    def assert_fn(config):
        checkpoint_to_check = train.get_checkpoint()
        with checkpoint_to_check.as_directory() as checkpoint_dir:
            with open(os.path.join(checkpoint_dir, "checkpoint.pkl"), "rb") as f:
                state = pickle.load(f)

        print("Loaded state from `resume_from_checkpoint`:", state)
        print("Expected state:", expected_state)
        assert state == expected_state, (state, expected_state)

        dummy_ckpt = tempfile.mkdtemp()
        with open(os.path.join(dummy_ckpt, "dummy.txt"), "w") as f:
            f.write("data")
        train.report({"dummy": 1}, checkpoint=Checkpoint.from_directory(dummy_ckpt))

    trainer = DataParallelTrainer(
        assert_fn,
        scaling_config=train.ScalingConfig(num_workers=2),
        run_config=train.RunConfig(
            name="test_resume_from_checkpoint",
            storage_path=storage_path,
            storage_filesystem=storage_filesystem,
        ),
        resume_from_checkpoint=checkpoint,
    )
    result = trainer.fit()

    # Make sure that the checkpoint indexing starts from scratch.
    assert Path(
        result.checkpoint.path
    ).name == StorageContext._make_checkpoint_dir_name(0)

    # Clean up this run's experiment directory immediately after.
    _delete_fs_path(result.filesystem, Path(result.path).parent.as_posix())


def _assert_storage_contents(
    local_inspect_dir: Path,
    exp_name: str,
    checkpoint_config: train.CheckpointConfig,
    trainable_name: str,
    test_trainer: bool,
    no_checkpoint_ranks: List[int] = None,
    constants: type = TestConstants,
):
    no_checkpoint_ranks = no_checkpoint_ranks or []

    # Second, inspect the contents of the storage path
    storage_path_ls = list(local_inspect_dir.glob("*"))
    assert len(storage_path_ls) == 1  # Only expect 1 experiment dir
    exp_dir = storage_path_ls[0]
    assert exp_dir.name == exp_name

    # Files synced by the driver
    assert len(list(exp_dir.glob("tuner.pkl"))) == 1
    if test_trainer:
        assert len(list(exp_dir.glob("trainer.pkl"))) == 1
    # 2 copies of these files:
    # 1 for the initial run, and 1 for the manually restored run.
    assert len(list(exp_dir.glob("basic-variant-state-*"))) == 2
    assert len(list(exp_dir.glob("experiment_state-*"))) == 2

    # Files synced by the worker
    assert (
        len(list(exp_dir.glob(f"{trainable_name}*"))) == 1
        if test_trainer
        else constants.NUM_TRIALS
    )
    for trial_dir in exp_dir.glob(f"{trainable_name}*"):
        # If set, expect num_to_keep. Otherwise, expect to see all of them.
        expected_num_checkpoints = (
            checkpoint_config.num_to_keep or constants.NUM_ITERATIONS
        )

        assert len(list(trial_dir.glob("checkpoint_*"))) == expected_num_checkpoints
        checkpoint_idxs = sorted(
            [
                _get_checkpoint_index(checkpoint_dir.name)
                for checkpoint_dir in trial_dir.glob("checkpoint_*")
            ]
        )
        # Ex: If num_to_keep=2 out of 6 total checkpoints,
        # expect checkpoint_004 and checkpoint_005.
        assert checkpoint_idxs == list(
            range(
                constants.NUM_ITERATIONS - expected_num_checkpoints,
                constants.NUM_ITERATIONS,
            )
        )

        for checkpoint_dir in trial_dir.glob("checkpoint_*"):
            # 1 shared checkpoint.pkl file, written by the trainable / all workers.
            assert (
                len(list(checkpoint_dir.glob("checkpoint.pkl"))) == 1
                # NOTE: Dict checkpoint is only for the ClassTrainable.
                or len(list(checkpoint_dir.glob(_DICT_CHECKPOINT_FILE_NAME))) == 1
            )
            if test_trainer:
                # 1 checkpoint shard per worker.
                # Unless the worker did not report a checkpoint (no_checkpoint_ranks).
                assert {
                    _get_checkpoint_shard_rank(checkpoint_shard.name)
                    for checkpoint_shard in checkpoint_dir.glob(
                        "checkpoint_shard-*.pkl"
                    )
                } == {
                    i
                    for i in range(constants.NUM_WORKERS)
                    if i not in no_checkpoint_ranks
                }

        if test_trainer:
            expected_num_artifacts = constants.NUM_ITERATIONS * constants.NUM_WORKERS
        else:
            expected_num_artifacts = constants.NUM_ITERATIONS
        assert len(list(trial_dir.glob("artifact-*"))) == expected_num_artifacts

        # NOTE: This result file is synced by the driver.
        assert len(list(trial_dir.glob(EXPR_RESULT_FILE))) == 1


@pytest.mark.parametrize("trainable", [train_fn, ClassTrainable])
@pytest.mark.parametrize("storage_path_type", ["nfs", "cloud", "custom_fs"])
@pytest.mark.parametrize(
    "checkpoint_config",
    [train.CheckpointConfig(), train.CheckpointConfig(num_to_keep=2)],
)
def test_tuner(
    tmp_path,
    trainable,
    storage_path_type,
    checkpoint_config: train.CheckpointConfig,
):
    """End-to-end test that the new persistence mode works with the Tuner API.
    This test covers many `storage_path_type` options:
    - storage_path=None --> save locally to the default local path (e.g., ~/ray_results)
    - storage_path="nfs" --> save locally to a fake NFS path
    - storage_path="cloud" --> save to a mock S3 bucket
    - storage_path="custom_fs" --> save to a custom pyarrow filesystem
        - The custom fs is a local filesystem that appends a path prefix to every path.

    This is the expected output at the storage path:

    {storage_path}/{exp_name}
    ├── tuner.pkl                   <- Driver artifacts (global experiment state)
    ├── basic-variant-state.json
    ├── experiment_state.json
    ├── train_fn_a2b9e_00000_0_...
    │   ├── artifact-iter=0.txt     <- Trial artifacts
    │   ├── ...
    │   ├── checkpoint_000000       <- Trial checkpoints
    │   │   └── checkpoint.pkl
    │   ├── ...
    │   ├── events.out.tfevents...  <- Driver artifacts (trial results)
    │   ├── params.json
    │   ├── params.pkl
    │   ├── progress.csv
    │   └── result.json
    └── train_fn_a2b9e_00001_1_...
        └── ...                     <- Same as above
    """
    exp_name = f"tuner_persistence_test-{uuid.uuid4().hex}"

    with _resolve_storage_type(storage_path_type, tmp_path) as (
        storage_path,
        storage_filesystem,
    ):
        run_config = train.RunConfig(
            storage_path=storage_path,
            storage_filesystem=storage_filesystem,
            name=exp_name,
            verbose=0,
            failure_config=train.FailureConfig(max_failures=1),
            checkpoint_config=checkpoint_config,
            sync_config=train.SyncConfig(sync_artifacts=True),
        )
        tuner = tune.Tuner(
            trainable,
            param_space={
                "num_iterations": TestConstants.NUM_ITERATIONS,
                "fail_iters": [2, 4],
                # NOTE: This param is only used in the ClassTrainable.
                "save_checkpoint_as_dict": tune.grid_search([True, False]),
                "tmp_path": tmp_path,
            },
            run_config=run_config,
            # 2 samples (from the grid search). Run 1 at at time to test actor reuse
            tune_config=tune.TuneConfig(num_samples=1, max_concurrent_trials=1),
        )
        result_grid = tuner.fit()
        assert result_grid.errors

        restored_tuner = tune.Tuner.restore(
            path=str(URI(run_config.storage_path) / exp_name),
            trainable=trainable,
            storage_filesystem=storage_filesystem,
            resume_errored=True,
        )
        result_grid = restored_tuner.fit()
        assert not result_grid.errors

        local_inspect_dir, storage_fs_path = _get_local_inspect_dir(
            root_local_path=tmp_path,
            storage_path=run_config.storage_path,
            storage_filesystem=storage_filesystem,
        )

    # First, check that the ResultGrid returns the correct paths.
    print(result_grid)
    experiment_fs_path = result_grid.experiment_path
    assert isinstance(result_grid.filesystem, pyarrow.fs.FileSystem), result_grid
    assert experiment_fs_path == os.path.join(storage_fs_path, exp_name)
    assert len(result_grid) == TestConstants.NUM_TRIALS
    for result in result_grid:
        trial_fs_path = result.path
        assert isinstance(result.filesystem, pyarrow.fs.FileSystem), result
        assert trial_fs_path.startswith(experiment_fs_path)
        for checkpoint, _ in result.best_checkpoints:
            assert checkpoint.path.startswith(trial_fs_path)

    # Next, inspect the storage path contents.
    _assert_storage_contents(
        local_inspect_dir,
        exp_name,
        checkpoint_config,
        trainable_name=trainable.__name__,
        test_trainer=False,
    )


@pytest.mark.parametrize("storage_path_type", ["nfs", "cloud", "custom_fs"])
@pytest.mark.parametrize(
    "checkpoint_config",
    [
        train.CheckpointConfig(),
        train.CheckpointConfig(
            num_to_keep=1,
            checkpoint_score_attribute=TestConstants.SCORE_KEY,
            checkpoint_score_order="max",
        ),
    ],
)
def test_trainer(
    tmp_path, storage_path_type, checkpoint_config: train.CheckpointConfig
):
    """Same end-to-end test as `test_tuner`, but also includes a
    `DataParallelTrainer(resume_from_checkpoint)` test at the end.

    {storage_path}/{exp_name}
    ├── experiment_state-2023-07-28_10-00-38.json       <- Initial exp state
    ├── basic-variant-state-2023-07-28_10-00-38.json
    ├── experiment_state-2023-07-28_10-01-38.json       <- Restored exp state
    ├── basic-variant-state-2023-07-28_10-01-38.json
    ├── trainer.pkl
    ├── tuner.pkl
    └── DataParallelTrainer_46367_00000_0_...
        ├── events.out.tfevents...
        ├── params.json
        ├── params.pkl
        ├── progress.csv
        ├── result.json
        ├── checkpoint_000000
        │   ├── checkpoint.pkl                    <- Shared checkpoint file
        │   ├── checkpoint_shard-rank=0.pkl       <- Worker checkpoint shards
        │   └── checkpoint_shard-rank=1.pkl
        ├── ...
        ├── artifact-rank=0-iter=0.txt            <- Worker artifacts
        ├── artifact-rank=1-iter=0.txt
        ├── ...
        ├── artifact-rank=0-iter=1.txt
        ├── artifact-rank=1-iter=1.txt
        └── ...
    """
    exp_name = f"trainer_persistence_test-{uuid.uuid4().hex}"
    no_checkpoint_ranks = [0]

    with _resolve_storage_type(storage_path_type, tmp_path) as (
        storage_path,
        storage_filesystem,
    ):
        run_config = train.RunConfig(
            storage_path=storage_path,
            storage_filesystem=storage_filesystem,
            name=exp_name,
            verbose=0,
            checkpoint_config=checkpoint_config,
            failure_config=train.FailureConfig(max_failures=1),
            sync_config=train.SyncConfig(sync_artifacts=True),
        )
        trainer = DataParallelTrainer(
            train_fn,
            train_loop_config={
                "in_trainer": True,
                "num_iterations": TestConstants.NUM_ITERATIONS,
                "fail_iters": [2, 4],
                # Test that global rank 0 is not required to checkpoint.
                "no_checkpoint_ranks": no_checkpoint_ranks,
            },
            scaling_config=train.ScalingConfig(num_workers=TestConstants.NUM_WORKERS),
            run_config=run_config,
        )
        print("\nStarting initial run.\n")
        with pytest.raises(TrainingFailedError):
            result = trainer.fit()

        print("\nStarting manually restored run.\n")
        restored_trainer = DataParallelTrainer.restore(
            path=str(URI(run_config.storage_path) / exp_name),
            storage_filesystem=storage_filesystem,
        )
        result = restored_trainer.fit()

        _resume_from_checkpoint(
            result.checkpoint,
            expected_state={"iter": TestConstants.NUM_ITERATIONS - 1},
        )

        local_inspect_dir, storage_fs_path = _get_local_inspect_dir(
            root_local_path=tmp_path,
            storage_path=run_config.storage_path,
            storage_filesystem=storage_filesystem,
        )

    # First, inspect that the result object returns the correct paths.
    print(result)
    trial_fs_path = result.path
    assert trial_fs_path.startswith(storage_fs_path)
    for checkpoint, _ in result.best_checkpoints:
        assert checkpoint.path.startswith(trial_fs_path)

    _assert_storage_contents(
        local_inspect_dir,
        exp_name,
        checkpoint_config,
        trainable_name="DataParallelTrainer",
        test_trainer=True,
        no_checkpoint_ranks=no_checkpoint_ranks,
    )


if __name__ == "__main__":
    import sys

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