import os
import sys
from unittest.mock import patch

import pytest

import ray
from ray._private.accelerators import HPUAcceleratorManager, hpu
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy


def test_user_configured_more_than_visible(monkeypatch, call_ray_stop_only):
    # Test more hpus are configured than visible.
    monkeypatch.setenv("HABANA_VISIBLE_MODULES", "0,1,2")
    with pytest.raises(ValueError):
        ray.init(resources={"HPU": 4})


@patch(
    "ray._private.accelerators.HPUAcceleratorManager.get_current_node_num_accelerators",  # noqa: E501
    return_value=4,
)
def test_auto_detected_more_than_visible(
    mock_get_num_accelerators, monkeypatch, shutdown_only
):
    # Test more hpus are detected than visible.
    monkeypatch.setenv("HABANA_VISIBLE_MODULES", "0,1,2")
    ray.init()
    _ = mock_get_num_accelerators.called
    assert ray.available_resources()["HPU"] == 3


@patch(
    "ray._private.accelerators.HPUAcceleratorManager.get_current_node_num_accelerators",  # noqa: E501
    return_value=2,
)
def test_auto_detect_resources(mock_get_num_accelerators, shutdown_only):
    # Test that ray node resources are filled with auto detected count.
    ray.init()
    _ = mock_get_num_accelerators.called
    assert ray.available_resources()["HPU"] == 2


def test_get_current_process_visible_accelerator_ids():
    os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] = "0,1,2"
    assert HPUAcceleratorManager.get_current_process_visible_accelerator_ids() == [
        "0",
        "1",
        "2",
    ]  # noqa: E501

    del os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR]
    assert HPUAcceleratorManager.get_current_process_visible_accelerator_ids() is None

    os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] = ""
    assert HPUAcceleratorManager.get_current_process_visible_accelerator_ids() == []

    del os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR]


def test_set_current_process_visible_accelerator_ids():
    HPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0"])
    assert os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] == "0"

    HPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1"])
    assert os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] == "0,1"

    HPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1", "2"])
    assert os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] == "0,1,2"

    del os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR]


@pytest.mark.parametrize(
    "test_config",
    [
        (1, False),
        (0.5, True),
        (3, False),
    ],
)
def test_validate_resource_request_quantity(test_config):
    num_hpus, expect_error = test_config

    if expect_error:
        assert (
            HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[0]
            is False
        )
        assert (
            HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[1]
            is not None
        )
    else:
        assert (
            HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[0]
            is True
        )
        assert (
            HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[1]
            is None
        )


def test_check_accelerator_info():

    if HPUAcceleratorManager.is_initialized():
        assert (
            "Intel-GAUDI" in HPUAcceleratorManager.get_current_node_accelerator_type()
        )
    else:
        assert HPUAcceleratorManager.get_current_node_accelerator_type() is None

    assert HPUAcceleratorManager.get_resource_name() == "HPU"


def test_decorator_args():

    # This is a valid way of using the decorator.
    @ray.remote(resources={"HPU": 1})  # noqa: F811
    class Actor:  # noqa: F811
        def __init__(self):
            pass

    # This is a valid way of using the decorator.
    @ray.remote(num_cpus=1, resources={"HPU": 1})  # noqa: F811
    class Actor:  # noqa: F811
        def __init__(self):
            pass


def test_actor_deletion_with_hpus(shutdown_only):
    ray.init(num_cpus=1, resources={"HPU": 1})

    # When an actor that uses an HPU exits, make sure that the HPU resources
    # are released.

    @ray.remote(resources={"HPU": 1})
    class Actor:
        def getpid(self):
            return os.getpid()

    for _ in range(5):
        # If we can successfully create an actor, that means that enough
        # HPU resources are available.
        a = Actor.remote()
        ray.get(a.getpid.remote())


@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
def test_actor_hpus(ray_start_cluster):
    cluster = ray_start_cluster
    num_nodes = 2
    num_hpus_per_raylet = 2
    for i in range(num_nodes):
        cluster.add_node(num_cpus=10 * 2, resources={"HPU": num_hpus_per_raylet})
    ray.init(address=cluster.address)

    @ray.remote(resources={"HPU": 1})
    class Actor1:
        def __init__(self):
            resource_ids = ray.get_runtime_context().get_accelerator_ids()
            self.hpu_ids = resource_ids.get("HPU")

        def get_location_and_ids(self):
            return (
                ray.get_runtime_context().get_node_id(),
                tuple(self.hpu_ids),
            )

    # Create one actor per HPU.
    actors = [Actor1.remote() for _ in range(num_nodes * num_hpus_per_raylet)]
    # Make sure that no two actors are assigned to the same HPU.
    locations_and_ids = ray.get(
        [actor.get_location_and_ids.remote() for actor in actors]
    )
    node_names = {location for location, hpu_id in locations_and_ids}
    assert len(node_names) == num_nodes
    location_actor_combinations = []
    for node_name in node_names:
        for hpu_id in range(num_hpus_per_raylet):
            location_actor_combinations.append((node_name, (f"{hpu_id}",)))

    assert set(locations_and_ids) == set(location_actor_combinations)

    # Creating a new actor should fail because all of the HPUs are being
    # used.
    a = Actor1.remote()
    ready_ids, _ = ray.wait([a.get_location_and_ids.remote()], timeout=0.01)
    assert ready_ids == []


def test_actor_habana_visible_devices(shutdown_only):
    """Test user can overwrite HABANA_VISIBLE_MODULES
    after the actor is created."""
    ray.init(resources={"HPU": 1})

    @ray.remote(resources={"HPU": 1})
    class Actor:
        def set_habana_visible_devices(self, habana_visible_devices):
            os.environ["HABANA_VISIBLE_MODULES"] = habana_visible_devices

        def get_habana_visible_devices(self):
            return os.environ["HABANA_VISIBLE_MODULES"]

    actor = Actor.remote()
    assert ray.get(actor.get_habana_visible_devices.remote()) == "0"
    ray.get(actor.set_habana_visible_devices.remote("0,1"))
    assert ray.get(actor.get_habana_visible_devices.remote()) == "0,1"


def test_hpu_ids(shutdown_only):
    num_hpus = 3
    ray.init(num_cpus=num_hpus, resources={"HPU": num_hpus})

    def get_hpu_ids(hpus_per_worker):
        hpu_ids = ray.get_runtime_context().get_accelerator_ids()["HPU"]
        assert len(hpu_ids) == hpus_per_worker
        modules = os.environ.get("HABANA_VISIBLE_MODULES")
        if modules is not None:
            assert modules == ",".join([str(i) for i in hpu_ids])  # noqa
        for hpu_id in hpu_ids:
            assert hpu_id in [str(i) for i in range(num_hpus)]
        return hpu_ids

    f0 = ray.remote(resources={"HPU": 0})(lambda: get_hpu_ids(0))
    f1 = ray.remote(resources={"HPU": 1})(lambda: get_hpu_ids(1))

    list_of_ids = ray.get([f0.remote() for _ in range(10)])
    assert list_of_ids == 10 * [[]]
    ray.get([f1.remote() for _ in range(10)])

    # Test that actors have HABANA_VISIBLE_MODULES set properly.

    def _check_hpu_env(expected_num_hpus):
        hpu_ids = ray.get_runtime_context().get_accelerator_ids()["HPU"]
        assert len(hpu_ids) == expected_num_hpus
        if expected_num_hpus > 0:
            assert os.environ["HABANA_VISIBLE_MODULES"] == ",".join(
                [str(i) for i in hpu_ids]  # noqa
            )
        else:
            assert os.environ.get("HABANA_VISIBLE_MODULES") is None

    @ray.remote
    class Actor:
        def __init__(self, num_hpus):
            self.num_hpus = num_hpus
            _check_hpu_env(num_hpus)
            self.x = num_hpus

        def test(self):
            _check_hpu_env(self.num_hpus)
            return self.x

    a0 = Actor.remote(0)
    assert ray.get(a0.test.remote()) == 0

    a1 = Actor.options(resources={"HPU": 1}).remote(1)
    assert ray.get(a1.test.remote()) == 1


def test_hpu_with_placement_group(shutdown_only):
    num_hpus = 2
    ray.init(num_cpus=1, resources={"HPU": num_hpus})

    @ray.remote(resources={"HPU": num_hpus})
    class HPUActor:
        def __init__(self):
            pass

        def ready(self):
            hpu_ids = ray.get_runtime_context().get_accelerator_ids()["HPU"]
            assert len(hpu_ids) == num_hpus
            assert os.environ["HABANA_VISIBLE_MODULES"] == ",".join(
                [str(i) for i in hpu_ids]  # noqa
            )

    # Reserve a placement group of 1 bundle that reserves 1 CPU and 2 HPU.
    pg = placement_group([{"CPU": 1, "HPU": num_hpus}])

    # Wait until placement group is created.
    ray.get(pg.ready(), timeout=10)

    actor = HPUActor.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=pg,
        )
    ).remote()

    ray.get(actor.ready.remote(), timeout=10)


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