import sys
import time
from pathlib import Path
from typing import List

import pytest

import ray
import ray.train
from ray.cluster_utils import Cluster
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
from ray.train.v2._internal.constants import HEALTH_CHECK_INTERVAL_S_ENV_VAR
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer


@pytest.fixture
def cluster():
    cluster = Cluster(initialize_head=True, head_node_args=dict(num_cpus=0))
    cluster.wait_for_nodes()
    ray.init(
        address=cluster.address,
        runtime_env={"working_dir": str(Path(__file__).parent)},
    )
    yield cluster
    ray.shutdown()
    cluster.shutdown()


def train_fn(config: dict):
    train_context = ray.train.get_context()
    rank = train_context.get_world_rank()

    start_epoch = 1
    checkpoint = ray.train.get_checkpoint()
    min_world_size = None
    max_world_size = None
    if checkpoint:
        checkpoint_data = load_dict_checkpoint(checkpoint)
        start_epoch = checkpoint_data["epoch"] + 1
        min_world_size = checkpoint_data.get("min_world_size")
        max_world_size = checkpoint_data.get("max_world_size")
        if rank == 0:
            print("Restoring from epoch: ", start_epoch)

    for epoch in range(start_epoch, config.get("num_epochs", 60) + 1):
        world_size = train_context.get_world_size()
        if min_world_size is None:
            min_world_size = world_size
        if max_world_size is None:
            max_world_size = world_size
        min_world_size = min(min_world_size, world_size)
        max_world_size = max(max_world_size, world_size)
        # TODO: This test injects errors by "killing nodes," which ungracefully
        # kills processes. This means that any backlog in the checkpoint queue
        # will not be flushed to the controller.
        # This means that the checkpoint populated on restore may not be
        # the most recent one.
        # Set the poll interval < health check interval to reduce the
        # backlog size to mitigate the issue.
        time.sleep(2 * config.get("health_check_interval_s", 1))

        with create_dict_checkpoint(
            {
                "epoch": epoch,
                "min_world_size": min_world_size,
                "max_world_size": max_world_size,
            }
        ) as checkpoint:
            ray.train.report(
                {
                    "epoch": epoch,
                    "world_size": world_size,
                    "min_world_size": min_world_size,
                    "max_world_size": max_world_size,
                },
                checkpoint=checkpoint if rank == 0 else None,
                checkpoint_dir_name=f"checkpoint-epoch={epoch}",
            )
        if rank == 0:
            print("Finished epoch: ", epoch)


def test_elastic_training(monkeypatch, tmp_path, cluster):
    """End to end test for elastic training.

    This test covers:
    * Elastic startup (0 resources -> min resources)
    * Elastic scale up while running (min resources -> max resources)
    * Elastic scale down due to failure while running
    * Checkpointing + restoration
    * Preemption failure handling
    """
    unit_time_s = 0.1
    health_check_interval_s = unit_time_s
    elastic_resize_monitor_interval_s = unit_time_s * 10
    num_epochs = 30

    monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, str(health_check_interval_s))

    @ray.remote(num_cpus=0)
    def run_training():
        trainer = DataParallelTrainer(
            train_fn,
            train_loop_config={
                "num_epochs": num_epochs,
                "health_check_interval_s": health_check_interval_s,
            },
            scaling_config=ray.train.ScalingConfig(
                num_workers=(4, 32),
                use_gpu=True,
                elastic_resize_monitor_interval_s=elastic_resize_monitor_interval_s,
            ),
            run_config=ray.train.RunConfig(
                storage_path=str(tmp_path),
                checkpoint_config=ray.train.CheckpointConfig(num_to_keep=2),
                # NOTE: The outer test script will inject 2 node failures.
                failure_config=ray.train.FailureConfig(max_failures=2),
            ),
        )
        return trainer.fit()

    run_training_future = run_training.remote()

    start = time.time()
    ALL_NODES = []

    def print_status(message):
        elapsed = time.time() - start
        print()
        print("-" * 80)
        cluster_resources = {
            resource: value
            for resource, value in ray.cluster_resources().items()
            if resource in ("CPU", "GPU")
        }
        print(f"[elapsed={elapsed:.1f}s] {cluster_resources=}")
        print(message)
        print("-" * 80)
        print()

    def sleep(num_units):
        time.sleep(unit_time_s * num_units)

    def add_nodes(gpus: List[int]) -> List:
        added_nodes = []
        for num_gpus in gpus:
            node = cluster.add_node(num_gpus=num_gpus, wait=False)
            added_nodes.append(node)

        print_status(f"Added {len(gpus)} node(s) with num_gpus: {gpus}")
        cluster.wait_for_nodes()
        return added_nodes

    def remove_nodes(nodes: List):
        for node in nodes:
            cluster.remove_node(node)

        cluster.wait_for_nodes()
        print_status(f"Removed nodes: {nodes}")

    # Wait a bit before adding resources.
    print_status("Waiting for training to start...")
    sleep(8)

    # Add a node with 4 GPUs
    ALL_NODES.extend(add_nodes([4]))

    # Wait a bit before adding more resources.
    sleep(8)
    print("Adding 4 GPU node.")
    ALL_NODES.extend(add_nodes([4]))
    sleep(1)
    ALL_NODES.extend(add_nodes([4]))
    # Should not upscale here due to the elastic resize monitor interval.

    # Should upscale to 12 during this sleep.
    sleep(20)

    # Kill a node.
    remove_nodes([ALL_NODES.pop(0)])
    sleep(12)

    # Kill all worker nodes.
    remove_nodes(ALL_NODES)
    ALL_NODES = []

    sleep(8)
    ALL_NODES.extend(add_nodes(gpus=[1] * 16))

    sleep(12)
    # 4 extra GPUs shouldn't be used.
    ALL_NODES.extend(add_nodes(gpus=[4] * 4 + [1] * 4))

    result: ray.train.Result = ray.get(run_training_future)

    print_status(f"Training finished with result: {result}")
    assert not result.error
    assert result.metrics["min_world_size"] >= 4
    assert result.metrics["max_world_size"] <= 32
    assert result.metrics["max_world_size"] >= result.metrics["min_world_size"]
    assert result.checkpoint
    assert Path(result.checkpoint.path).name == f"checkpoint-epoch={num_epochs}"


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", "-x", __file__]))