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

from aio_pika import Message

from amqp.config.amq_configuration import AMQConfiguration
from amqp.rabbitmq.user_management_service_client import (
    UserManagementServiceClient,
    UserManagementServiceException
)


class TestUserManagementServiceClient(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        # Mock configuration
        self.config = MagicMock(spec=AMQConfiguration)
        self.config.dispatch = MagicMock()
        self.config.dispatch.amq_host = "localhost"
        self.config.dispatch.amq_port = 5672
        self.config.dispatch.rabbit_mq_user = "guest"
        self.config.dispatch.rabbit_mq_password = "guest"
        
        # Create client
        self.client = UserManagementServiceClient(self.config)
        
        # Mock connection and channel
        self.client.connection = MagicMock()
        self.client.channel = MagicMock()
        self.client.exchange = MagicMock()
        self.client.exchange.publish = AsyncMock()
        self.client.response_queue = MagicMock()
        self.client.initialized = True
        
    async def test_initialize(self):
        # Mock aio_pika.connect_robust
        with patch('aio_pika.connect_robust', new_callable=AsyncMock) as mock_connect:
            # Mock connection
            mock_connection = MagicMock()
            mock_connect.return_value = mock_connection
            
            # Mock channel - create an AsyncMock that returns the channel when awaited
            mock_channel = AsyncMock()
            mock_connection.channel = AsyncMock()
            mock_connection.channel.return_value = mock_channel
            mock_channel.set_qos = AsyncMock()
            
            # Mock exchange
            mock_exchange = MagicMock()
            mock_channel.declare_exchange = AsyncMock(return_value=mock_exchange)
            
            # Mock queue
            mock_queue = MagicMock()
            mock_channel.declare_queue = AsyncMock(return_value=mock_queue)
            mock_queue.consume = AsyncMock(return_value="consumer-tag")
            
            # Reset client
            self.client.initialized = False
            self.client.connection = None
            self.client.channel = None
            self.client.exchange = None
            self.client.response_queue = None
            
            # Call initialize
            await self.client.initialize()
            
            # Verify
            mock_connect.assert_called_once_with(
                host=self.config.dispatch.amq_host,
                port=self.config.dispatch.amq_port,
                login=self.config.dispatch.rabbit_mq_user,
                password=self.config.dispatch.rabbit_mq_password,
                loop=self.client.loop
            )
            mock_connection.channel.assert_called_once()
            mock_channel.set_qos.assert_called_once_with(prefetch_count=1)
            mock_channel.declare_exchange.assert_called_once()
            mock_channel.declare_queue.assert_called_once()
            mock_queue.consume.assert_called_once()
            self.assertTrue(self.client.initialized)
    
    async def test_call_service_success(self):
        # Create a future for the response
        future = asyncio.Future()
        future.set_result({"type": "OK", "result": [{"name": "John Doe"}]})
        
        # Mock the futures dictionary
        with patch.dict(self.client.futures, {}, clear=True):
            # Mock uuid.uuid4
            with patch('uuid.uuid4', return_value="test-correlation-id"):
                # Mock loop.create_future
                with patch.object(self.client.loop, 'create_future', return_value=future):
                    # Call service
                    result = await self.client.call_service(
                        service=UserManagementServiceClient.TOKEN_SERVICE,
                        method="validateToken",
                        args={"token": "test-token"}
                    )
                    
                    # Verify
                    self.client.exchange.publish.assert_called_once()
                    self.assertEqual(result, [{"name": "John Doe"}])
    
    async def test_call_service_error(self):
        # Create a future for the response
        future = asyncio.Future()
        future.set_result({
            "type": "ERROR",
            "exception": "cleverthis.clevermicro.auth.invalid_token",
            "message": "Invalid token",
            "additional_info": {"token": "test-token"}
        })
        
        # Mock the futures dictionary
        with patch.dict(self.client.futures, {}, clear=True):
            # Mock uuid.uuid4
            with patch('uuid.uuid4', return_value="test-correlation-id"):
                # Mock loop.create_future
                with patch.object(self.client.loop, 'create_future', return_value=future):
                    # Call service and expect exception
                    with self.assertRaises(UserManagementServiceException) as context:
                        await self.client.call_service(
                            service=UserManagementServiceClient.TOKEN_SERVICE,
                            method="validateToken",
                            args={"token": "test-token"}
                        )
                    
                    # Verify exception
                    exception = context.exception
                    self.assertEqual(exception.exception_type, "cleverthis.clevermicro.auth.invalid_token")
                    self.assertEqual(exception.message, "Invalid token")
                    self.assertEqual(exception.additional_info, {"token": "test-token"})
    
    async def test_on_response_callback(self):
        # Create a future
        future = asyncio.Future()
        
        # Add future to futures dictionary
        self.client.futures["test-correlation-id"] = future
        
        # Create a mock message
        message = MagicMock()
        message.correlation_id = "test-correlation-id"
        message.body = json.dumps({"type": "OK", "result": [{"name": "John Doe"}]}).encode('utf-8')
        message.ack = AsyncMock()
        
        # Call callback
        await self.client._on_response_callback(message)
        
        # Verify
        message.ack.assert_called_once()
        self.assertTrue(future.done())
        self.assertEqual(future.result(), {"type": "OK", "result": [{"name": "John Doe"}]})
    
    async def test_on_response_callback_no_correlation_id(self):
        # Create a mock message with no correlation ID
        message = MagicMock()
        message.correlation_id = None
        message.ack = AsyncMock()
        
        # Call callback
        await self.client._on_response_callback(message)
        
        # Verify
        message.ack.assert_called_once()
    
    async def test_on_response_callback_no_future(self):
        # Create a mock message with unknown correlation ID
        message = MagicMock()
        message.correlation_id = "unknown-correlation-id"
        message.ack = AsyncMock()
        
        # Call callback
        await self.client._on_response_callback(message)
        
        # Verify
        message.ack.assert_called_once()
    
    async def test_validate_token(self):
        # Mock call_service
        with patch.object(self.client, 'call_service', new_callable=AsyncMock) as mock_call:
            mock_call.return_value = [{"sub": "user-123", "name": "John Doe"}]
            
            # Call validate_token
            result = await self.client.validate_token("test-token")
            
            # Verify
            mock_call.assert_called_once_with(
                service=UserManagementServiceClient.TOKEN_SERVICE,
                method="validateToken",
                args={"token": "test-token"}
            )
            self.assertEqual(result, [{"sub": "user-123", "name": "John Doe"}])
    
    async def test_query_user_by_id(self):
        # Mock call_service
        with patch.object(self.client, 'call_service', new_callable=AsyncMock) as mock_call:
            mock_call.return_value = [{"id": "user-123", "name": "John Doe", "email": "john@example.com"}]
            
            # Call query_user_by_id
            result = await self.client.query_user_by_id("user-123")
            
            # Verify
            mock_call.assert_called_once_with(
                service=UserManagementServiceClient.USER_SERVICE,
                method="queryUserById",
                args={"userId": "user-123"}
            )
            self.assertEqual(result, [{"id": "user-123", "name": "John Doe", "email": "john@example.com"}])
    
    async def test_close(self):
        # Mock connection
        self.client.connection.is_closed = False
        self.client.connection.close = AsyncMock()
        
        # Call close
        await self.client.close()
        
        # Verify
        self.client.connection.close.assert_called_once()
        self.assertFalse(self.client.initialized)


class TestGetUserInfoFromToken(unittest.IsolatedAsyncioTestCase):
    async def test_get_user_info_from_token_success(self):
        # Mock client
        client = MagicMock(spec=UserManagementServiceClient)
        client.validate_token = AsyncMock(return_value=[{"sub": "user-123"}])
        client.query_user_by_id = AsyncMock(return_value=[{"id": "user-123", "name": "John Doe"}])
        
        # Import the function
        from amqp.rabbitmq.user_management_sample import get_user_info_from_token
        
        # Call function
        result = await get_user_info_from_token(client, "test-token")
        
        # Verify
        client.validate_token.assert_called_once_with("test-token")
        client.query_user_by_id.assert_called_once_with("user-123")
        self.assertEqual(result, [{"id": "user-123", "name": "John Doe"}])
    
    async def test_get_user_info_from_token_no_user_id(self):
        # Mock client
        client = MagicMock(spec=UserManagementServiceClient)
        client.validate_token = AsyncMock(return_value=[{"name": "John Doe"}])  # No sub claim
        
        # Import the function
        from amqp.rabbitmq.user_management_sample import get_user_info_from_token
        
        # Call function and expect exception
        with self.assertRaises(UserManagementServiceException) as context:
            await get_user_info_from_token(client, "test-token")
        
        # Verify exception
        exception = context.exception
        self.assertEqual(exception.exception_type, "cleverthis.clevermicro.auth.invalid_token")
        self.assertEqual(exception.message, "Token does not contain user ID (sub claim)")
    
    async def test_get_user_info_from_token_service_exception(self):
        # Mock client
        client = MagicMock(spec=UserManagementServiceClient)
        client.validate_token = AsyncMock(
            side_effect=UserManagementServiceException(
                "cleverthis.clevermicro.auth.invalid_token",
                "Invalid token",
                {}
            )
        )
        
        # Import the function
        from amqp.rabbitmq.user_management_sample import get_user_info_from_token
        
        # Call function and expect exception
        with self.assertRaises(UserManagementServiceException) as context:
            await get_user_info_from_token(client, "test-token")
        
        # Verify exception
        exception = context.exception
        self.assertEqual(exception.exception_type, "cleverthis.clevermicro.auth.invalid_token")
        self.assertEqual(exception.message, "Invalid token")


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