import sys
import time
from random import random

import pytest

import ray
import ray.cluster_utils
from ray._common.test_utils import wait_for_condition
from ray.util.placement_group import placement_group, remove_placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy


def run_mini_integration_test(cluster, pg_removal=True, num_pgs=999):
    # This test checks the race condition between remove / creation.
    # This test shouldn't be flaky. If it fails on the last ray.get
    # that highly likely indicates a real bug.
    # It also runs 3 times to make sure the test consistently passes.
    # When 999 resource quantity is used, it fails about every other time
    # when the test was written.
    resource_quantity = num_pgs
    num_nodes = 5
    custom_resources = {"pg_custom": resource_quantity}

    # TODO(sang): Cluster setup. Remove when running in real clusters.
    nodes = []
    for _ in range(num_nodes):
        nodes.append(cluster.add_node(num_cpus=3, resources=custom_resources))
    cluster.wait_for_nodes()
    num_nodes = len(nodes)

    ray.init(address=cluster.address)
    bundles = [{"pg_custom": 1}] * num_nodes

    @ray.remote(num_cpus=0, resources={"pg_custom": 1}, max_calls=0)
    def mock_task():
        time.sleep(0.1)
        return True

    @ray.remote(num_cpus=0)
    def pg_launcher(num_pgs_to_create):
        print("Creating pgs")
        pgs = []
        for _ in range(num_pgs_to_create):
            pgs.append(placement_group(bundles, strategy="STRICT_SPREAD"))

        pgs_removed = []
        pgs_unremoved = []
        # Randomly choose placement groups to remove.
        if pg_removal:
            print("removing pgs")
        for pg in pgs:
            if random() < 0.5 and pg_removal:
                pgs_removed.append(pg)
            else:
                pgs_unremoved.append(pg)
        print(len(pgs_unremoved))

        tasks = []
        # Randomly schedule tasks or actors on placement groups that
        # are not removed.
        for pg in pgs_unremoved:
            for i in range(num_nodes):
                tasks.append(
                    mock_task.options(
                        scheduling_strategy=PlacementGroupSchedulingStrategy(
                            placement_group=pg, placement_group_bundle_index=i
                        )
                    ).remote()
                )
        # Remove the rest of placement groups.
        if pg_removal:
            for pg in pgs_removed:
                remove_placement_group(pg)
        ray.get(tasks)
        # Since placement groups are scheduled, remove them.
        for pg in pgs_unremoved:
            remove_placement_group(pg)

    pg_launchers = []
    for _ in range(3):
        pg_launchers.append(pg_launcher.remote(num_pgs // 3))

    ray.get(pg_launchers, timeout=240)
    ray.shutdown()

    ray.init(address=cluster.address)

    cluster_resources = ray.cluster_resources()
    cluster_resources.pop("memory")
    cluster_resources.pop("object_store_memory")

    def wait_for_resource_recovered():
        for resource, val in ray.available_resources().items():
            if resource in cluster_resources and cluster_resources[resource] != val:
                return False
            if "_group_" in resource:
                return False
        return True

    wait_for_condition(wait_for_resource_recovered)


@pytest.mark.parametrize("execution_number", range(1))
def test_placement_group_create_only(ray_start_cluster, execution_number):
    """PG mini integration test without remove_placement_group

    When there are failures, this will help identifying if issues are
    from removal or not.
    """
    run_mini_integration_test(ray_start_cluster, pg_removal=False, num_pgs=333)


@pytest.mark.parametrize("execution_number", range(3))
def test_placement_group_remove_stress(ray_start_cluster, execution_number):
    """Full PG mini integration test that runs many
    concurrent remove_placement_group
    """
    run_mini_integration_test(ray_start_cluster, pg_removal=True, num_pgs=999)


if __name__ == "__main__":

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