import asyncio
import time
import unittest
from unittest.mock import AsyncMock, MagicMock, patch

from amqp.adapter.backpressure_handler import (
    BackpressureHandler,
    ScaleRequestV1,
    ScalingRequestAlert,
)


class TestBackpressureHandler(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self.channel = AsyncMock()
        self.loop = asyncio.get_event_loop()
        self.config = MagicMock()
        self.config.amq_adapter.swarm_service_id = "test-service"
        self.config.amq_adapter.swarm_task_id = "test-task"
        self.config.backpressure.threshold_load = 100
        self.config.backpressure.time_window = 1
        self.config.backpressure.idle_duration = 2
        self.config.backpressure.cpu_monitoring_enabled = False
        self.config.backpressure.cpu_overload_duration = 1

        # Record callbacks for testing
        BackpressureHandler._callback_list = {}

        self.handler = BackpressureHandler(self.channel, self.loop, self.config)
        self.handler.exchange = AsyncMock()

        # Set do_loop to 1 to run the loop only once
        self.handler.do_loop = 1

    async def test_increase_decrease_current_load(self):
        self.handler.current_load = 0
        self.handler.increase_current_load()
        self.assertEqual(self.handler.current_load, 1)
        self.handler.decrease_current_load()
        self.assertEqual(self.handler.current_load, 0)

    async def test_update_current_load(self):
        self.handler.current_availability = 0
        self.handler.update_current_availability(5)
        self.assertEqual(self.handler.current_availability, 5)

    async def test_update_last_data_message_time(self):
        old_time = self.handler.last_backpressure_event_time
        self.handler.update_last_backpressure_event_time()
        self.assertGreater(self.handler.last_backpressure_event_time, old_time)

    @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
    async def test_update_backpressure_value(self, mock_publish):
        # Test updating with a higher maximum
        await self.handler.update_backpressure_value(50, 200)
        self.assertEqual(self.handler.current_availability, 50)
        self.assertEqual(self.handler.max_availability, 200)

        # Test updating with a lower maximum (should not change)
        await self.handler.update_backpressure_value(60, 50)
        self.assertEqual(self.handler.current_availability, 60)
        self.assertEqual(self.handler.max_availability, 200)

    @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
    async def test_check_overload_condition_normal(self, mock_publish):
        # Set up a normal state (60% available capacity)
        self.handler.max_availability = 100
        self.handler.current_availability = 60
        self.handler.last_backpressure_event = ScalingRequestAlert.UPDATE
        self.handler.last_backpressure_event_time = time.time() - 2  # 2 seconds ago

        await self.handler.check_overload_condition()

        # Should send an UPDATE event
        mock_publish.assert_called_once()
        args = mock_publish.call_args[0][0]
        self.assertEqual(args.requestType, ScalingRequestAlert.UPDATE)
        self.assertEqual(args.current_availability, 60)
        self.assertEqual(args.max_availability, 100)

    @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
    async def test_check_overload_condition_overload(self, mock_publish):
        # Set up an overload state (10% available capacity)
        self.handler.max_availability = 100
        self.handler.current_availability = 10
        self.handler.last_backpressure_event = ScalingRequestAlert.UPDATE
        self.handler.last_backpressure_event_time = time.time() - 2  # 2 seconds ago

        await self.handler.check_overload_condition()

        # Should send an OVERLOAD event
        mock_publish.assert_called_once()
        args = mock_publish.call_args[0][0]
        self.assertEqual(args.requestType, ScalingRequestAlert.OVERLOAD)
        self.assertEqual(args.current_availability, 10)
        self.assertEqual(args.max_availability, 100)

    @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
    async def test_check_overload_condition_idle(self, mock_publish):
        # Set up an idle state (90% available capacity)
        self.handler.max_availability = 100
        self.handler.current_availability = 90
        self.handler.last_backpressure_event = ScalingRequestAlert.IDLE
        # emulate fact that IDLE state is for 3 seconds longer than required minimum for sending IDLE event
        self.handler.last_backpressure_event_time = (
            time.time() - self.config.backpressure.idle_duration - 3
        )

        await self.handler.check_overload_condition()

        # Should send an IDLE event
        mock_publish.assert_called_once()
        args = mock_publish.call_args[0][0]
        self.assertEqual(args.requestType, ScalingRequestAlert.IDLE)
        self.assertEqual(args.current_availability, 90)
        self.assertEqual(args.max_availability, 100)

    @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
    async def test_check_overload_condition_auto_max(self, mock_publish):
        # Test with unset maximum that should be auto-detected
        self.handler.max_availability = -1
        self.handler.current_availability = 50
        self.handler.last_backpressure_event = ScalingRequestAlert.UPDATE
        self.handler.last_backpressure_event_time = (
            time.time() - self.config.backpressure.time_window - 2
        )

        await self.handler.check_overload_condition()

        # Should set maximum to current_availability and send UPDATE
        # self.assertEqual(self.handler.max_availability, 50)
        mock_publish.assert_called_once()
        args = mock_publish.call_args[0][0]
        self.assertEqual(args.requestType, ScalingRequestAlert.UPDATE)
        self.assertEqual(args.current_availability, 50)
        self.assertEqual(args.max_availability, -1)

        # Test with a higher value that should update the maximum
        mock_publish.reset_mock()
        await self.handler.update_backpressure_value(80, 80)

        # Should update maximum to 80
        self.assertEqual(self.handler.max_availability, 80)
        mock_publish.assert_called_once()
        args = mock_publish.call_args[0][0]
        self.assertEqual(args.current_availability, 80)
        self.assertEqual(args.max_availability, 80)

    @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
    async def test_handle_backpressure_overload_event(self, mock_publish):
        self.handler.max_availability = 100
        self.handler.current_availability = 10

        await self.handler.handle_backpressure_overload_event()

        mock_publish.assert_called_once()
        args = mock_publish.call_args[0][0]
        self.assertEqual(args.requestType, ScalingRequestAlert.OVERLOAD)
        self.assertEqual(args.current_availability, 10)
        self.assertEqual(args.max_availability, 100)

    @patch("amqp.adapter.backpressure_handler.BackpressureHandler.publish_backpressure_request")
    async def test_handle_backpressure_idle_event(self, mock_publish):
        self.handler.max_availability = 100
        self.handler.current_availability = 90

        await self.handler._handle_backpressure_idle_event()

        mock_publish.assert_called_once()
        args = mock_publish.call_args[0][0]
        self.assertEqual(args.requestType, ScalingRequestAlert.IDLE)
        self.assertEqual(args.current_availability, 90)
        self.assertEqual(args.max_availability, 100)

    @patch("asyncio.run_coroutine_threadsafe")
    async def test_publish_backpressure_request(self, mock_run):
        scaling_request = ScaleRequestV1(
            "test-service", "test-task", 100, 50, ScalingRequestAlert.UPDATE
        )

        await self.handler.publish_backpressure_request(scaling_request)

        # Check that run_coroutine_threadsafe was called
        self.assertEqual(mock_run.call_count, 1)

        # Check that the callback was recorded
        self.assertIn("_wrap_rabbit_mq_api", BackpressureHandler._callback_list)


class TestScaleRequestV1(unittest.TestCase):
    def test_to_json(self):
        request = ScaleRequestV1("test-service", "test-task", 100, 50, ScalingRequestAlert.UPDATE)
        json_str = request.to_json()

        # Check that the JSON contains the expected fields
        self.assertIn('"version": 1', json_str)
        self.assertIn('"serviceId": "test-service"', json_str)
        self.assertIn('"taskId": "test-task"', json_str)
        self.assertIn('"maxAvailability": 100', json_str)
        self.assertIn('"currentAvailability": 50', json_str)
        self.assertIn(f'"requestType": "{ScalingRequestAlert.UPDATE.value}"', json_str)


if __name__ == "__main__":
    unittest.main()
