import ray
from ray._common.test_utils import wait_for_condition
from ray.autoscaler.v2.sdk import get_cluster_status
import time
from logger import logger
from typing import Dict

ray.init("auto")

# Sync with the compute config.
HEAD_NODE_CPU = 0
WORKER_NODE_CPU = 4
IDLE_TERMINATION_S = 60 * 5  # 5 min
DEFAULT_RETRY_INTERVAL_MS = 15 * 1000  # 15 sec


def check_cluster(target_num_nodes: int, target_resources: Dict[str, float]):
    gcs_address = ray.get_runtime_context().gcs_address
    cluster_status = get_cluster_status(gcs_address)

    assert (
        len(cluster_status.active_nodes) + len(cluster_status.idle_nodes)
    ) == target_num_nodes

    for k, v in target_resources.items():
        assert cluster_status.total_resources().get(k, 0) == v

    return True


ctx = {
    "num_cpus": 0,
    "num_nodes": 1,
}
logger.info(f"Starting cluster with {ctx['num_nodes']} nodes, {ctx['num_cpus']} cpus")
check_cluster(
    target_num_nodes=ctx["num_nodes"], target_resources={"CPU": ctx["num_cpus"]}
)


# Request for cluster resources
def test_request_cluster_resources(ctx: dict):
    from ray.autoscaler._private.commands import request_resources

    request_resources(num_cpus=8)

    ctx["num_cpus"] += 8
    ctx["num_nodes"] += 8 // WORKER_NODE_CPU

    # Assert on number of worker nodes.
    logger.info(
        f"Requesting cluster constraints: {ctx['num_nodes']} nodes, "
        f"{ctx['num_cpus']} cpus"
    )
    wait_for_condition(
        check_cluster,
        timeout=60 * 5,  # 5min
        retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
        target_num_nodes=ctx["num_nodes"],
        target_resources={"CPU": ctx["num_cpus"]},
    )

    # Reset the cluster constraints.
    request_resources(num_cpus=0)

    ctx["num_cpus"] -= 8
    ctx["num_nodes"] -= 8 // WORKER_NODE_CPU
    logger.info(
        f"Waiting for cluster go idle after constraint cleared: {ctx['num_nodes']} "
        f"nodes, {ctx['num_cpus']} cpus"
    )
    wait_for_condition(
        check_cluster,
        timeout=60 + IDLE_TERMINATION_S,  # 1min + idle timeout
        retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
        target_num_nodes=ctx["num_nodes"],
        target_resources={"CPU": ctx["num_cpus"]},
    )


# Run actors/tasks that exceed the cluster resources should upscale the cluster
def test_run_tasks_concurrent(ctx: dict):
    num_tasks = 2
    num_actors = 2

    @ray.remote(num_cpus=WORKER_NODE_CPU)
    def f():
        while True:
            time.sleep(1)

    @ray.remote(num_cpus=WORKER_NODE_CPU)
    class Actor:
        def __init__(self):
            pass

    tasks = [f.remote() for _ in range(num_tasks)]
    actors = [Actor.remote() for _ in range(num_actors)]

    ctx["num_cpus"] += (num_tasks + num_actors) * WORKER_NODE_CPU
    ctx["num_nodes"] += num_tasks + num_actors

    logger.info(f"Waiting for {ctx['num_nodes']} nodes, {ctx['num_cpus']} cpus")
    wait_for_condition(
        check_cluster,
        timeout=60 * 5,  # 5min
        retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
        target_num_nodes=ctx["num_nodes"],
        target_resources={"CPU": ctx["num_cpus"]},
    )

    [ray.cancel(task) for task in tasks]
    [ray.kill(actor) for actor in actors]

    ctx["num_cpus"] -= (num_actors + num_tasks) * WORKER_NODE_CPU
    ctx["num_nodes"] -= num_actors + num_tasks

    logger.info(
        f"Waiting for cluster to scale down to {ctx['num_nodes']} nodes, "
        f"{ctx['num_cpus']} cpus"
    )
    wait_for_condition(
        check_cluster,
        timeout=60 + IDLE_TERMINATION_S,
        retry_interval_ms=DEFAULT_RETRY_INTERVAL_MS,
        target_num_nodes=ctx["num_nodes"],
        target_resources={"CPU": ctx["num_cpus"]},
    )


test_request_cluster_resources(ctx)
test_run_tasks_concurrent(ctx)