import json
import os
from typing import Dict, Union

import pytest

import ray
from ray import tune
from ray.train._internal.storage import StorageContext
from ray.train.tests.util import create_dict_checkpoint
from ray.tune.trainable import wrap_function


@pytest.fixture
def ray_start_2_cpus():
    address_info = ray.init(num_cpus=2)
    yield address_info
    # The code after the yield will run as teardown code.
    ray.shutdown()


class SavingTrainable(tune.Trainable):
    def __init__(self, return_type: str, *args, **kwargs):
        self.return_type = return_type
        super(SavingTrainable, self).__init__(*args, **kwargs)

    def step(self):
        return {"iter": self.training_iteration}

    def save_checkpoint(self, tmp_checkpoint_dir: str):
        checkpoint_data = {"data": 1}

        if self.return_type == "object":
            return checkpoint_data

        subdir = os.path.join(tmp_checkpoint_dir, "subdir")
        os.makedirs(subdir, exist_ok=True)
        checkpoint_file = os.path.join(subdir, "checkpoint.pkl")
        with open(checkpoint_file, "w") as f:
            f.write(json.dumps(checkpoint_data))

        if self.return_type == "root":
            return tmp_checkpoint_dir
        elif self.return_type == "subdir":
            return subdir
        elif self.return_type == "checkpoint":
            return checkpoint_file

    def load_checkpoint(self, checkpoint: Union[Dict, str]):
        if self.return_type == "object":
            assert isinstance(checkpoint, dict)
            checkpoint_data = checkpoint
            checkpoint_file = None
        elif self.return_type == "root":
            assert "subdir" not in checkpoint
            checkpoint_file = os.path.join(checkpoint, "subdir", "checkpoint.pkl")
        elif self.return_type == "subdir":
            assert "subdir" in checkpoint
            assert "checkpoint.pkl" not in checkpoint
            checkpoint_file = os.path.join(checkpoint, "checkpoint.pkl")
        else:  # self.return_type == "checkpoint"
            assert checkpoint.endswith("subdir/checkpoint.pkl")
            checkpoint_file = checkpoint

        if checkpoint_file:
            with open(checkpoint_file, "rb") as f:
                checkpoint_data = json.load(f)

        checkpoint_data = {
            key: value
            for key, value in checkpoint_data.items()
            if not key.startswith("_")
        }
        assert checkpoint_data == {"data": 1}, checkpoint_data


def function_trainable(config):
    with create_dict_checkpoint({"checkpoint_data": 5}) as checkpoint:
        tune.report({"metric": 4}, checkpoint=checkpoint)


@pytest.mark.parametrize("return_type", ["object", "root"])
def test_save_load_checkpoint_path_class(ray_start_2_cpus, return_type, tmpdir):
    """Assert that restoring from a Trainable.save() future works with
    class trainables.

    Needs Ray cluster so we get actual futures.
    """
    trainable = ray.remote(SavingTrainable).remote(return_type=return_type)

    # Train one step
    ray.get(trainable.train.remote())

    # Save checkpoint
    saving_future = trainable.save.remote()

    # Check for errors
    ray.get(saving_future)

    restoring_future = trainable.restore.remote(saving_future)

    ray.get(restoring_future)


def test_save_load_checkpoint_path_fn(ray_start_2_cpus, tmp_path):
    """Assert that restoring from a Trainable.save() future works with
    function trainables.

    Needs Ray cluster so we get actual futures.
    """
    trainable_cls = wrap_function(function_trainable)
    trainable = ray.remote(trainable_cls).remote(
        storage=StorageContext(
            storage_path=str(tmp_path),
            experiment_dir_name="exp",
            trial_dir_name="trial",
        )
    )
    ray.get(trainable.train.remote())

    saving_future = trainable.save.remote()

    # Check for errors
    ray.get(saving_future)

    restoring_future = trainable.restore.remote(saving_future)

    ray.get(restoring_future)


if __name__ == "__main__":
    import sys

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