import asyncio
import sys
from unittest.mock import MagicMock, patch

import pytest
from freezegun import freeze_time

from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
    ResourceRequestPriority,
)
from ray.train.v2._internal.execution.callback import ControllerCallback
from ray.train.v2._internal.execution.scaling_policy import (
    AUTOSCALING_REQUESTS_EXPIRE_TIME_S,
    AUTOSCALING_REQUESTS_INTERVAL_S,
    NoopDecision,
    ResizeDecision,
)
from ray.train.v2._internal.execution.scaling_policy.elastic import (
    ElasticScalingPolicy,
)
from ray.train.v2._internal.execution.worker_group import (
    WorkerGroupPollStatus,
    WorkerGroupState,
    WorkerStatus,
)
from ray.train.v2._internal.util import time_monotonic
from ray.train.v2.api.config import ScalingConfig


@pytest.fixture(autouse=True)
def mock_autoscaling_coordinator(monkeypatch):
    mock_coordinator = MagicMock()
    mock_coordinator._allocated_resources = None
    mock_coordinator.get_allocated_resources.remote = MagicMock(
        side_effect=lambda _: mock_coordinator._allocated_resources
    )

    monkeypatch.setattr(
        ElasticScalingPolicy, "_autoscaling_coordinator", mock_coordinator
    )


@pytest.fixture(autouse=True)
def patch_ray_get():
    with patch(
        "ray.get",
        side_effect=lambda x, **_: x,
    ):
        yield


def _get_mock_worker_group_status(num_workers: int) -> WorkerGroupPollStatus:
    return WorkerGroupPollStatus(
        worker_statuses={
            i: WorkerStatus(running=True, error=None) for i in range(num_workers)
        },
    )


def _get_mock_worker_group_state(
    num_workers: int, start_time: float
) -> WorkerGroupState:
    return WorkerGroupState(
        start_time=start_time,
        placement_group_handle=MagicMock(),
        workers=[MagicMock() for _ in range(num_workers)],
        sync_actor=MagicMock(),
    )


@patch.object(ElasticScalingPolicy, "GET_ALLOCATED_RESOURCES_INTERVAL_S", 0.0)
def test_non_running_worker_group_decision():
    """Test decisions being made when the worker group is initializing/restarting.
    Ensures that the policy will resize the worker group as soon as resources are available.
    """
    min_workers, max_workers = 4, 64
    resources_per_worker = {"CPU": 8, "GPU": 1}

    scaling_config = ScalingConfig(
        num_workers=(min_workers, max_workers),
        resources_per_worker=resources_per_worker,
        use_gpu=True,
    )
    policy = ElasticScalingPolicy(scaling_config)
    mock_coordinator = policy._autoscaling_coordinator

    # No resources are available at the start
    decision = policy.make_decision_for_non_running_worker_group()
    assert isinstance(decision, NoopDecision)

    # Resources for < min workers are available
    mock_coordinator._allocated_resources = [resources_per_worker] * (min_workers - 1)
    decision = policy.make_decision_for_non_running_worker_group()
    assert isinstance(decision, NoopDecision)

    # Resources for >= min workers are available
    mock_coordinator._allocated_resources = [resources_per_worker] * min_workers
    decision = policy.make_decision_for_non_running_worker_group()
    assert isinstance(decision, ResizeDecision)
    assert decision.num_workers == min_workers

    # Resources for >= max workers are available
    mock_coordinator._allocated_resources = [resources_per_worker] * max_workers
    decision = policy.make_decision_for_non_running_worker_group()
    assert isinstance(decision, ResizeDecision)
    assert decision.num_workers == max_workers


def test_before_controller_abort():
    """Test that before_controller_abort sends a cancel request to the AutoscalingCoordinator."""
    resources_per_worker = {"CPU": 4, "GPU": 1}
    scaling_config = ScalingConfig(
        num_workers=(2, 4),
        resources_per_worker=resources_per_worker,
        use_gpu=True,
    )
    policy = ElasticScalingPolicy(scaling_config)
    mock_coordinator = policy._autoscaling_coordinator

    # Call before_controller_abort and check that cancel_request is called with the requester_id
    policy.before_controller_abort()
    mock_coordinator.cancel_request.remote.assert_called_once_with(
        requester_id=policy._requester_id
    )


def test_get_allocated_resources_interval():
    """Tests that remote calls to the AutoscalingCoordinator are spaced out by a minimum time interval."""
    min_workers, max_workers = 4, 64
    resources_per_worker = {"CPU": 8, "GPU": 1}
    get_allocated_resources_interval_s = (
        ElasticScalingPolicy.GET_ALLOCATED_RESOURCES_INTERVAL_S
    )

    scaling_config = ScalingConfig(
        num_workers=(min_workers, max_workers),
        resources_per_worker=resources_per_worker,
        use_gpu=True,
    )
    policy = ElasticScalingPolicy(scaling_config)
    mock_coordinator = policy._autoscaling_coordinator

    with freeze_time() as frozen_time:
        # No resources are available at the start
        allocated_resources = policy._get_allocated_resources()
        assert allocated_resources is None

        # Resources for < min workers are available
        frozen_time.tick(get_allocated_resources_interval_s)
        mock_coordinator._allocated_resources = [resources_per_worker] * (
            min_workers - 1
        )
        allocated_resources = policy._get_allocated_resources()
        assert allocated_resources == [resources_per_worker] * (min_workers - 1)

        # Resources for >= min workers are available, but get_allocated_resources interval
        # has not yet passed.
        mock_coordinator._allocated_resources = [resources_per_worker] * min_workers
        allocated_resources = policy._get_allocated_resources()
        assert allocated_resources == [resources_per_worker] * (min_workers - 1)

        # Resources for >= min workers are available and the get_allocated_resources
        # interval has passed.
        frozen_time.tick(get_allocated_resources_interval_s)
        mock_coordinator._allocated_resources = [resources_per_worker] * min_workers
        allocated_resources = policy._get_allocated_resources()
        assert allocated_resources == [resources_per_worker] * min_workers

        # Resources for >= max workers are available but the get_allocated_resources
        # interval has not yet passed.
        mock_coordinator._allocated_resources = [resources_per_worker] * max_workers
        allocated_resources = policy._get_allocated_resources()
        assert allocated_resources == [resources_per_worker] * min_workers

        # Resources for >= max workers are available and the get_allocated_resources
        # interval has passed.
        frozen_time.tick(get_allocated_resources_interval_s)
        allocated_resources = policy._get_allocated_resources()
        assert allocated_resources == [resources_per_worker] * max_workers


@patch.object(ElasticScalingPolicy, "GET_ALLOCATED_RESOURCES_INTERVAL_S", 0.0)
def test_running_worker_group_decision():
    """Test decisions being made when the worker group is running.
    Ensures that the policy will resize the worker group when there is a change
    in available resources.
    """
    min_workers, max_workers = 4, 64
    resources_per_worker = {"CPU": 8, "GPU": 1}

    scaling_config = ScalingConfig(
        num_workers=(min_workers, max_workers),
        resources_per_worker=resources_per_worker,
        use_gpu=True,
        # NOTE: This test just asserts the policy decisions, not the monitor interval.
        elastic_resize_monitor_interval_s=0.0,
    )
    policy = ElasticScalingPolicy(scaling_config)
    mock_coordinator = policy._autoscaling_coordinator

    # The worker group just started
    worker_group_state = _get_mock_worker_group_state(min_workers, time_monotonic())
    worker_group_status = _get_mock_worker_group_status(min_workers)

    # No change in resources
    mock_coordinator._allocated_resources = [resources_per_worker] * min_workers
    decision = policy.make_decision_for_running_worker_group(
        worker_group_state=worker_group_state,
        worker_group_status=worker_group_status,
    )
    assert isinstance(decision, NoopDecision)

    # Resources for < min workers are available
    mock_coordinator._allocated_resources = [resources_per_worker] * (min_workers - 1)
    decision = policy.make_decision_for_running_worker_group(
        worker_group_state=worker_group_state,
        worker_group_status=worker_group_status,
    )
    assert isinstance(decision, NoopDecision)

    # More resources are available.
    mock_coordinator._allocated_resources = [resources_per_worker] * max_workers
    decision = policy.make_decision_for_running_worker_group(
        worker_group_state=worker_group_state,
        worker_group_status=worker_group_status,
    )
    assert isinstance(decision, ResizeDecision)
    assert decision.num_workers == max_workers


def test_monitor_recently_started_worker_group():
    """Test monitor decisions being made when the worker group is running.
    Ensures that resizing decisions are not made too soon after the worker group starts.
    """
    min_workers, max_workers = 4, 64
    monitor_interval_s = 60
    resources_per_worker = {"CPU": 8, "GPU": 1}

    scaling_config = ScalingConfig(
        num_workers=(min_workers, max_workers),
        resources_per_worker=resources_per_worker,
        use_gpu=True,
        elastic_resize_monitor_interval_s=monitor_interval_s,
    )
    policy = ElasticScalingPolicy(scaling_config)
    mock_coordinator = policy._autoscaling_coordinator

    with freeze_time() as frozen_time:
        # The worker group just started
        worker_group_state = _get_mock_worker_group_state(min_workers, time_monotonic())
        worker_group_status = _get_mock_worker_group_status(min_workers)

        # Advance time partway through the monitor interval
        frozen_time.tick(delta=monitor_interval_s / 2)

        # Even though there are new resources available, we should not resize yet
        # because the monitor interval has not passed since
        mock_coordinator._allocated_resources = [resources_per_worker] * (
            max_workers - 1
        )

        assert isinstance(
            policy.make_decision_for_running_worker_group(
                worker_group_state=worker_group_state,
                worker_group_status=worker_group_status,
            ),
            NoopDecision,
        )

        frozen_time.tick(delta=monitor_interval_s / 2)

        # The monitor interval has passed, should detect resources and resize
        decision = policy.make_decision_for_running_worker_group(
            worker_group_state=worker_group_state,
            worker_group_status=worker_group_status,
        )
        assert isinstance(decision, ResizeDecision)
        assert decision.num_workers == max_workers - 1


def test_monitor_long_running_worker_group():
    """Test monitor decisions being made when the worker group is running.
    Ensures that the resizing considerations are not made too frequently.
    """
    min_workers, max_workers = 4, 64
    monitor_interval_s = 60
    resources_per_worker = {"CPU": 8, "GPU": 1}

    scaling_config = ScalingConfig(
        num_workers=(min_workers, max_workers),
        resources_per_worker=resources_per_worker,
        use_gpu=True,
        elastic_resize_monitor_interval_s=monitor_interval_s,
    )
    policy = ElasticScalingPolicy(scaling_config)
    mock_coordinator = policy._autoscaling_coordinator

    with freeze_time() as frozen_time:
        worker_group_state = _get_mock_worker_group_state(min_workers, time_monotonic())
        worker_group_status = _get_mock_worker_group_status(min_workers)
        mock_coordinator._allocated_resources = [resources_per_worker] * min_workers

        # The worker group has been running for a while at the same size
        frozen_time.tick(monitor_interval_s * 60)

        # Consider resizing.
        decision = policy.make_decision_for_running_worker_group(
            worker_group_state=worker_group_state,
            worker_group_status=worker_group_status,
        )
        assert isinstance(decision, NoopDecision)

        # We recently considered resizing, so we should wait until the next interval
        # to consider again --> no-op even if new resources are available
        mock_coordinator._allocated_resources = [resources_per_worker] * max_workers
        frozen_time.tick(monitor_interval_s / 2)
        decision = policy.make_decision_for_running_worker_group(
            worker_group_state=worker_group_state,
            worker_group_status=worker_group_status,
        )
        assert isinstance(decision, NoopDecision)

        frozen_time.tick(monitor_interval_s / 2)
        decision = policy.make_decision_for_running_worker_group(
            worker_group_state=worker_group_state,
            worker_group_status=worker_group_status,
        )
        assert isinstance(decision, ResizeDecision)
        assert decision.num_workers == max_workers


def test_count_possible_workers():
    """Test counting the number of workers that can be started with
    available node resources."""
    resources_per_worker = {"CPU": 8, "GPU": 1}
    scaling_config = ScalingConfig(
        num_workers=(1, 8),
        use_gpu=True,
        resources_per_worker=resources_per_worker,
    )
    policy = ElasticScalingPolicy(scaling_config)

    # No resources
    assert policy._count_possible_workers([]) == 0

    # Single node
    assert policy._count_possible_workers([{"CPU": 8, "GPU": 1}]) == 1
    assert policy._count_possible_workers([{"CPU": 16, "GPU": 2}]) == 2
    assert policy._count_possible_workers([{"CPU": 16, "GPU": 1}]) == 1

    # Multinode
    assert policy._count_possible_workers([{"CPU": 7, "GPU": 1}] * 2) == 0
    assert policy._count_possible_workers([{"CPU": 9, "GPU": 2}] * 8) == 8
    assert policy._count_possible_workers([{"CPU": 16, "GPU": 2}] * 2) == 4
    assert policy._count_possible_workers([{"CPU": 8, "GPU": 1}] * 4) == 4

    # If there are excess resources, the number of workers is still capped at max_workers
    assert policy._count_possible_workers([{"CPU": 16, "GPU": 2}] * 10) == 8


def test_count_possible_workers_with_zero_resources():
    max_workers = 4
    scaling_config = ScalingConfig(
        num_workers=(1, max_workers),
        resources_per_worker={"CPU": 0, "GPU": 0, "memory": 0},
    )
    policy = ElasticScalingPolicy(scaling_config)

    assert (
        policy._count_possible_workers([{"CPU": 1, "GPU": 1, "memory": 1}])
        == max_workers
    )


def test_request_and_clear():
    """Tests that the policy makes resource requests and clears the requests."""
    resources_per_worker = {"CPU": 8, "GPU": 1}
    policy = ElasticScalingPolicy(
        scaling_config=ScalingConfig(
            use_gpu=True, resources_per_worker=resources_per_worker, num_workers=(2, 4)
        )
    )
    assert isinstance(policy, ControllerCallback)
    mock_coordinator = policy._autoscaling_coordinator

    def assert_resource_request_called_with():
        nonlocal mock_coordinator

        mock_coordinator.request_resources.remote.assert_called_with(
            requester_id=policy._requester_id,
            resources=[resources_per_worker] * 4,
            label_selectors=None,
            expire_after_s=AUTOSCALING_REQUESTS_EXPIRE_TIME_S,
            priority=ResourceRequestPriority.HIGH,
        )

    with freeze_time() as frozen_time:
        worker_group_state = _get_mock_worker_group_state(2, time_monotonic())
        worker_group_status = _get_mock_worker_group_status(2)

        # Test request_resources is called when the controller starts.
        policy.after_controller_start(train_run_context=MagicMock())
        assert mock_coordinator.request_resources.remote.call_count == 1
        assert_resource_request_called_with()

        # Test request_resources is only called in
        # `make_decision_for_running_worker_group`,
        # if `AUTOSCALING_REQUESTS_INTERVAL_S` has passed.
        frozen_time.tick(AUTOSCALING_REQUESTS_INTERVAL_S / 2)
        policy.make_decision_for_running_worker_group(
            worker_group_state=worker_group_state,
            worker_group_status=worker_group_status,
        )
        assert mock_coordinator.request_resources.remote.call_count == 1

        frozen_time.tick(AUTOSCALING_REQUESTS_INTERVAL_S / 2)
        policy.make_decision_for_running_worker_group(
            worker_group_state=worker_group_state,
            worker_group_status=worker_group_status,
        )
        assert mock_coordinator.request_resources.remote.call_count == 2
        assert_resource_request_called_with()

    # Test cancel_request is called when the controller is shutting down.
    asyncio.run(policy.before_controller_shutdown())
    mock_coordinator.cancel_request.remote.assert_called_once()


@pytest.mark.parametrize(
    "num_autoscaler_nodes, mock_gcs_intact_slices, expected_workers",
    [
        # No resources -> 0 workers
        (0, 0, 0),
        # Autoscaler sees 3 nodes but slice requires 4 -> rounds to 0 workers
        (3, 0, 0),
        # Autoscaler sees 4 nodes and 1 intact slice -> 4 workers
        (4, 1, 4),
        # Autoscaler sees 8 nodes (2 slices worth) but only 1 slice is intact
        # (e.g. 1 host down in the other slice) -> capped at 4 workers
        (8, 1, 4),
        # Autoscaler sees 8 nodes and both slices are intact -> 8 workers
        (8, 2, 8),
        # Autoscaler sees 12 nodes (3 slices), all 3 intact -> 12 workers (max)
        (12, 3, 12),
    ],
)
@patch("ray.util.tpu.get_num_tpu_slices")
def test_count_possible_workers_tpu_slice_rounding(
    mock_get_intact_slices,
    num_autoscaler_nodes,
    mock_gcs_intact_slices,
    expected_workers,
):
    """
    Test that TPU scaling correctly floors to the nearest complete, physically
    intact TPU slice. Intact slices are counted regardless of whether they are
    currently idle or in use.
    """
    mock_get_intact_slices.return_value = mock_gcs_intact_slices

    # Scaling config for TPU v6e 4x4 between 1 and 3 slices.
    scaling_config = ScalingConfig(
        use_tpu=True,
        accelerator_type="TPU-V6E",
        topology="4x4",
        num_workers=(4, 12),
        resources_per_worker={"TPU": 4, "CPU": 1},
    )
    policy = ElasticScalingPolicy(scaling_config)

    tpu_node = {"TPU": 4, "CPU": 1, "accelerator_type:TPU-V6E": 1}
    allocated_resources = [tpu_node] * num_autoscaler_nodes

    assert policy._count_possible_workers(allocated_resources) == expected_workers


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