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

from amqp.adapter.backpressure_handler import BackpressureHandler
from amqp.adapter.cleverthis_service_adapter import CleverThisServiceAdapter
from amqp.adapter.data_message_factory import DataMessageFactory
from amqp.adapter.file_handler import FileHandler
from amqp.adapter.service_message_factory import ServiceMessageFactory
from amqp.config.amq_configuration import AMQConfiguration
from amqp.model.model import (
    DataMessage,
    DataMessageMagicByte,
    DataResponse,
    MultipartDataMessage,
)
from amqp.model.snowflake_id import SnowflakeId
from amqp.rabbitmq.rabbit_mq_client import RabbitMQClient
from amqp.service.amq_message_handler import (
    DataMessageHandler,
    collect_trace_info,
    get_default_trace_state,
    set_text,
)


class TestDataMessageHandler(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        # Mock dependencies
        self.service_adapter = MagicMock(spec=CleverThisServiceAdapter)
        self.service_adapter.on_message = AsyncMock()

        self.tracer = MagicMock()
        self.tracer.start_span = MagicMock()

        self.message_factory = MagicMock(spec=DataMessageFactory)
        self.service_message_factory = MagicMock(spec=ServiceMessageFactory)
        self.reply_to_exchange = MagicMock()
        self.reply_to_exchange.publish = AsyncMock()
        self.reply_to_exchange.name = "test-exchange"

        self.rabbit_mq_client = MagicMock(spec=RabbitMQClient)
        self.rabbit_mq_client.connection = MagicMock()

        self.amq_configuration = MagicMock(spec=AMQConfiguration)
        self.amq_configuration.dispatch = MagicMock()
        self.amq_configuration.dispatch.rabbit_mq_url = "amqp://guest:guest@localhost:5672/"
        self.amq_configuration.dispatch.download_dir = "/tmp/test_downloads"

        self.loop = asyncio.get_event_loop()

        self.backpressure_handler = MagicMock(spec=BackpressureHandler)
        self.backpressure_handler.increase_current_load = MagicMock()
        self.backpressure_handler.decrease_current_load = MagicMock()
        self.backpressure_handler.check_overload_condition = AsyncMock()
        self.backpressure_handler.current_availability = 5
        self.backpressure_handler.max_availability = 10

        self.file_handler = MagicMock(spec=FileHandler)
        self.file_handler.on_message = AsyncMock(return_value={})

        # Create the handler
        self.handler = DataMessageHandler(
            service_adapter=self.service_adapter,
            tracer=self.tracer,
            message_factory=self.message_factory,
            service_message_factory=self.service_message_factory,
            reply_to_exchange=self.reply_to_exchange,
            rabbit_mq_client=self.rabbit_mq_client,
            amq_configuration=self.amq_configuration,
            loop=self.loop,
            backpressure_handler=self.backpressure_handler,
            file_handler=self.file_handler,
        )

        # Create a mock message
        self.mock_message = MagicMock()
        self.mock_message.delivery_tag = 123
        self.mock_message.consumer_tag = "consumer-tag"
        self.mock_message.reply_to = "reply-queue"
        self.mock_message.correlation_id = "correlation-id"
        self.mock_message.headers = {"clevermicro.addressing": 0}
        self.mock_message.properties = {}
        self.mock_message.body = b"test-body"
        self.mock_message.ack = AsyncMock()
        self.mock_message.nack = AsyncMock()

        # Create a mock data message
        self.mock_data_message = MagicMock(spec=DataMessage)
        self.mock_data_message.magic.return_value = DataMessageMagicByte.HTTP_REQUEST.value
        self.mock_data_message.id.return_value = SnowflakeId.from_hex("1234567890abcdef1234")
        self.mock_data_message.method.return_value = "GET"
        self.mock_data_message.path.return_value = "/test/path"
        self.mock_data_message.trace_info.return_value = {
            "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
        }
        self.mock_data_message.body.return_value = b'{"test": "data"}'

        # Create a mock data response
        self.mock_data_response = MagicMock(spec=DataResponse)
        self.mock_data_response.id.return_value = "response-id"
        self.mock_data_response.content_type.return_value = "application/json"

        # Patch DataMessageFactory.from_bytes
        self.from_bytes_patch = patch(
            "amqp.adapter.data_message_factory.DataMessageFactory.from_bytes",
            return_value=self.mock_data_message,
        )
        self.mock_from_bytes = self.from_bytes_patch.start()

        # Patch DataMessageFactory.from_stream
        self.from_stream_patch = patch(
            "amqp.adapter.data_message_factory.DataMessageFactory.from_stream",
            return_value=self.mock_data_message,
        )
        self.mock_from_stream = self.from_stream_patch.start()

        # Patch DataResponseFactory.serialize
        self.serialize_patch = patch(
            "amqp.adapter.data_response_factory.DataResponseFactory.serialize",
            return_value=b"serialized-response",
        )
        self.mock_serialize = self.serialize_patch.start()

        # Patch DataResponseFactory.from_bytes
        self.response_from_bytes_patch = patch(
            "amqp.adapter.data_response_factory.DataResponseFactory.from_bytes",
            return_value=self.mock_data_response,
        )
        self.mock_response_from_bytes = self.response_from_bytes_patch.start()

        # Patch asyncio.run_coroutine_threadsafe
        self.run_coroutine_patch = patch("asyncio.run_coroutine_threadsafe")
        self.mock_run_coroutine = self.run_coroutine_patch.start()

        # Patch time.time
        self.time_patch = patch("time.time", return_value=1234567890.0)
        self.mock_time = self.time_patch.start()

        # Patch os.remove
        self.os_remove_patch = patch("os.remove")
        self.mock_os_remove = self.os_remove_patch.start()

    def tearDown(self):
        self.from_bytes_patch.stop()
        self.from_stream_patch.stop()
        self.serialize_patch.stop()
        self.response_from_bytes_patch.stop()
        self.run_coroutine_patch.stop()
        self.time_patch.stop()
        self.os_remove_patch.stop()

    async def test_inbound_data_message_callback_no_thread_success(self):
        """Test successful processing of an inbound message."""
        # Setup
        self.service_adapter.on_message.return_value = self.mock_data_response

        # Execute
        await self.handler.inbound_data_message_callback_no_thread(self.mock_message)

        # Verify
        self.mock_from_bytes.assert_called_once_with(self.mock_message.body)
        self.tracer.start_span.assert_called_once()
        self.backpressure_handler.increase_current_load.assert_called_once()
        self.service_adapter.on_message.assert_called_once()
        self.reply_to_exchange.publish.assert_called_once()
        self.backpressure_handler.decrease_current_load.assert_called_once()
        self.mock_message.ack.assert_called_once_with(multiple=False)

    async def test_inbound_data_message_callback_no_thread_unknown_magic(self):
        """Test handling of a message with unknown magic value."""
        # Setup
        self.mock_data_message.magic.return_value = "UNKNOWN"

        # Execute
        await self.handler.inbound_data_message_callback_no_thread(self.mock_message)

        # Verify
        self.mock_message.nack.assert_called_once_with(requeue=False)
        # self.mock_message.ack.assert_called_once_with(multiple=False)

    async def test_inbound_data_message_callback_no_thread_no_message(self):
        """Test handling when no valid message is present."""
        # Setup
        self.mock_from_bytes.return_value = None

        # Execute
        await self.handler.inbound_data_message_callback_no_thread(self.mock_message)

        # Verify
        self.mock_message.ack.assert_called_once_with(multiple=False)

    async def test_inbound_data_message_callback_as_thread(self):
        """Test the thread-based callback method."""
        # Setup
        with patch("threading.Thread") as mock_thread:
            mock_thread_instance = MagicMock()
            mock_thread.return_value = mock_thread_instance

            # Execute
            await self.handler.inbound_data_message_callback_as_thread(self.mock_message)

            # Verify
            mock_thread.assert_called_once()
            mock_thread_instance.start.assert_called_once()

    async def test_reconstitute_data_message_standard(self):
        """Test reconstituting a standard data message."""
        # Execute
        result = await self.handler.reconstitute_data_message(self.mock_message)

        # Verify
        self.assertEqual(result, self.mock_data_message)
        self.mock_from_bytes.assert_called_once_with(self.mock_message.body)

    async def test_reconstitute_data_message_stream(self):
        """Test reconstituting a message from a stream."""
        # Setup
        self.mock_message.headers = {"clevermicro.addressing": 1}

        # Execute
        result = await self.handler.reconstitute_data_message(self.mock_message)

        # Verify
        self.assertEqual(MultipartDataMessage, result.__class__)
        self.mock_from_stream.assert_called_once()
        self.file_handler.on_message.assert_called_once()

    async def test_run_http_request_magic_success(self):
        """Test successful processing of an HTTP request."""
        # Setup
        self.service_adapter.on_message.return_value = self.mock_data_response

        # Execute
        await self.handler.run_http_request_magic(
            self.mock_message, self.mock_data_message, self.loop
        )

        # Verify
        self.backpressure_handler.increase_current_load.assert_called_once()
        self.backpressure_handler.check_overload_condition.assert_called()
        self.service_adapter.on_message.assert_called_once()
        self.reply_to_exchange.publish.assert_called_once()
        self.backpressure_handler.decrease_current_load.assert_called_once()

    async def test_run_http_request_magic_multipart(self):
        """Test processing a multipart message with file cleanup."""
        # Setup
        self.service_adapter.on_message.return_value = self.mock_data_response
        multipart_message = MultipartDataMessage(
            self.mock_data_message, {"file1": "/tmp/file1.txt"}
        )

        # Execute
        await self.handler.run_http_request_magic(self.mock_message, multipart_message, self.loop)

        # Verify
        self.mock_os_remove.assert_called_once_with("/tmp/file1.txt")

    async def test_run_http_request_magic_service_exception(self):
        """Test handling of an exception from the service adapter."""
        # Setup
        self.service_adapter.on_message.side_effect = Exception("Service error")

        # Execute
        await self.handler.run_http_request_magic(
            self.mock_message, self.mock_data_message, self.loop
        )

        # Verify
        self.backpressure_handler.decrease_current_load.assert_called_once()
        self.backpressure_handler.check_overload_condition.assert_called()

    async def test_run_http_request_magic_reply_exception(self):
        """Test handling of an exception when sending the reply."""
        # Setup
        self.service_adapter.on_message.return_value = self.mock_data_response
        self.reply_to_exchange.publish.side_effect = Exception("Reply error")

        # Execute
        await self.handler.run_http_request_magic(
            self.mock_message, self.mock_data_message, self.loop
        )

        # Verify
        # [no ACK at this level, ACK is in parent call] self.mock_message.ack.assert_called_once_with(requeue=False)
        self.backpressure_handler.decrease_current_load.assert_called_once()

    async def test_send_reply_success(self):
        """Test successful sending of a reply."""
        # Execute
        await self.handler.send_reply("reply-queue", self.mock_data_response)

        # Verify
        self.mock_run_coroutine.assert_called_once()

    async def test_send_reply_no_reply_to(self):
        """Test handling when no reply_to is provided."""
        # Execute
        await self.handler.send_reply(None, self.mock_data_response)

        # Verify
        self.mock_run_coroutine.assert_not_called()

    async def test_reply_received_callback_with_future(self):
        """Test handling a reply when there's a matching future."""
        # Setup
        future = asyncio.Future()
        self.handler.outstanding["correlation-id"] = future

        # Execute
        await self.handler.reply_received_callback(self.mock_message)

        # Verify
        self.assertTrue(future.done())
        self.assertEqual(future.result(), self.mock_data_response.body())
        self.mock_message.ack.assert_called_once_with(multiple=False)

    async def test_reply_received_callback_no_future(self):
        """Test handling a reply when there's no matching future."""
        # Execute
        await self.handler.reply_received_callback(self.mock_message)

        # Verify
        self.mock_message.ack.assert_called_once_with(multiple=False)

    def test_ensure_trace_info(self):
        """Test ensuring trace info is properly set."""
        # Execute
        self.handler.ensure_trace_info(self.mock_data_message)

        # Verify
        self.tracer.start_span.assert_called_once()

    def test_record_duration(self):
        """Test recording the duration of message processing."""
        # Setup
        start_time = 1234567880.0  # 10 seconds before mock_time

        # Execute
        self.handler.record_duration(start_time, 123)

        # No specific assertions needed as this is mostly logging

    def test_collect_trace_info(self):
        """Test collecting trace info."""
        # Execute
        result = collect_trace_info()

        # Verify it returns a dictionary
        self.assertIsInstance(result, dict)

    def test_get_default_trace_state(self):
        """Test getting default trace state."""
        # Execute
        result = get_default_trace_state()

        # Verify it returns a TraceState
        self.assertIsNotNone(result)

    def test_set_text(self):
        """Test setting text in a carrier."""
        # Setup
        carrier = {}

        # Execute
        set_text(carrier, "key", "value")

        # Verify
        self.assertEqual(carrier["key"], "value")

    async def test_reconstitute_data_message_stream_none_result(self):
        """Test reconstituting a message from a stream when the result is None."""
        # Setup
        self.mock_message.headers = {"clevermicro.addressing": 1}
        self.mock_from_stream.return_value = None

        # Execute
        result = await self.handler.reconstitute_data_message(self.mock_message)

        # Verify
        self.assertIsNone(result)
        self.mock_from_stream.assert_called_once()
        self.file_handler.on_message.assert_not_called()

    async def test_run_http_request_magic_file_removal_exception(self):
        """Test handling of an exception when removing files."""
        # Setup
        self.service_adapter.on_message.return_value = self.mock_data_response
        multipart_message = MultipartDataMessage(
            self.mock_data_message, {"file1": "/tmp/file1.txt"}
        )
        self.mock_os_remove.side_effect = OSError("File removal error")

        # Execute
        await self.handler.run_http_request_magic(self.mock_message, multipart_message, self.loop)

        # Verify
        self.mock_os_remove.assert_called_once_with("/tmp/file1.txt")
        # Should continue without raising an exception

    async def test_send_reply_content_type_list(self):
        """Test sending a reply with content type as a list."""
        # Setup
        self.mock_data_response.content_type.return_value = ["application/json", "text/plain"]

        # Execute
        await self.handler.send_reply("reply-queue", self.mock_data_response)

        # Verify
        self.mock_run_coroutine.assert_called_once()
        # We verify that Message is created with correct content type by inspecting the RabbitMQ call args if available.
        self.assertEqual(
            "application/json", self.reply_to_exchange.publish.call_args[1]["message"].content_type
        )


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