#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Unit tests for AIOProducer class.

"""
import asyncio
import concurrent.futures
import pytest
from unittest.mock import Mock, patch

from confluent_kafka import KafkaError, KafkaException
from confluent_kafka.experimental.aio.producer._AIOProducer import AIOProducer


class TestAIOProducer:
    """Unit tests for AIOProducer class."""

    @pytest.fixture
    def mock_producer(self):
        with patch('confluent_kafka.experimental.aio.producer._AIOProducer.confluent_kafka.Producer') as mock:
            yield mock

    @pytest.fixture
    def mock_common(self):
        with patch('confluent_kafka.experimental.aio.producer._AIOProducer._common') as mock:
            async def mock_async_call(executor, blocking_task, *args, **kwargs):
                return blocking_task(*args, **kwargs)
            mock.async_call.side_effect = mock_async_call
            yield mock

    @pytest.fixture
    def basic_config(self):
        return {'bootstrap.servers': 'localhost:9092'}

    @pytest.mark.asyncio
    async def test_constructor_behavior(self, mock_producer, mock_common, basic_config):
        custom_executor = concurrent.futures.ThreadPoolExecutor(max_workers=5)
        try:
            producer1 = AIOProducer(
                basic_config,
                max_workers=3,
                executor=custom_executor
            )

            assert producer1.executor is custom_executor
            assert producer1.executor._max_workers == 5
            assert producer1._is_closed is False
            assert hasattr(producer1, '_buffer_timeout_manager')
            assert hasattr(producer1, '_producer')

            producer2 = AIOProducer(basic_config, max_workers=2, batch_size=500)

            assert producer2.executor is not custom_executor
            assert isinstance(producer2.executor, concurrent.futures.ThreadPoolExecutor)
            assert producer2.executor._max_workers == 2

            assert producer2._batch_size == 500
            assert producer2._is_closed is False
            assert hasattr(producer2, '_buffer_timeout_manager')

            await producer2.close()
            assert producer2._is_closed is True

        finally:
            custom_executor.shutdown(wait=True)

    @pytest.mark.asyncio
    async def test_close_method(self, mock_producer, mock_common, basic_config):
        producer = AIOProducer(basic_config)
        assert producer._is_closed is False

        await producer.close()
        assert producer._is_closed is True

        producer2 = AIOProducer(basic_config)
        assert producer2._is_closed is False

        await producer2.close()
        await producer2.close()
        assert producer2._is_closed is True

    @pytest.mark.asyncio
    async def test_call_method_executor_usage(self, mock_producer, mock_common, basic_config):
        producer = AIOProducer(basic_config)

        mock_method = Mock(return_value="test_result")
        result = await producer._call(mock_method, "arg1", kwarg1="value1")

        mock_method.assert_called_once_with("arg1", kwarg1="value1")
        assert result == "test_result"

    @pytest.mark.asyncio
    async def test_produce_success(self, mock_producer, mock_common, basic_config):
        producer = AIOProducer(basic_config, batch_size=1)
        mock_msg = Mock()

        async def mock_flush_buffer(target_topic=None):
            batches = producer._batch_processor.create_batches(target_topic)
            for batch in batches:
                for future in batch.futures:
                    if not future.done():
                        future.set_result(mock_msg)
            producer._batch_processor.clear_buffer()

        with patch.object(producer, '_flush_buffer', side_effect=mock_flush_buffer):
            result_future = await producer.produce(topic="test_topic", value="test_value")
            result = await result_future
            assert result is mock_msg

        await producer.close()

    @pytest.mark.asyncio
    async def test_produce_error(self, mock_producer, mock_common, basic_config):
        producer = AIOProducer(basic_config, batch_size=1)
        mock_error = KafkaError(KafkaError._MSG_TIMED_OUT)

        async def mock_flush_buffer(target_topic=None):
            batches = producer._batch_processor.create_batches(target_topic)
            for batch in batches:
                for future in batch.futures:
                    if not future.done():
                        future.set_exception(KafkaException(mock_error))
            producer._batch_processor.clear_buffer()

        with patch.object(producer, '_flush_buffer', side_effect=mock_flush_buffer):
            result_future = await producer.produce(topic="test_topic", value="test_value")

            with pytest.raises(KafkaException):
                await result_future

        await producer.close()

    @pytest.mark.asyncio
    async def test_produce_with_delayed_callback(self, mock_producer, mock_common, basic_config):
        """Test that Future properly waits for delayed delivery callback with batching."""
        producer = AIOProducer(basic_config, batch_size=2)  # Need 2 messages to trigger flush

        batch_called = asyncio.Event()
        captured_messages = None

        def mock_produce_batch(topic, messages, partition=None):
            nonlocal captured_messages
            captured_messages = messages
            batch_called.set()
            # Don't call callbacks immediately - simulate real async behavior

        mock_producer.return_value.produce_batch.side_effect = mock_produce_batch
        mock_producer.return_value.poll.return_value = 1

        # Start first produce - won't trigger flush yet, but will return a Future
        first_future = await producer.produce(topic="test", value="test1")

        # The Future should be pending (not resolved yet)
        assert not first_future.done()

        # Add second message to trigger batch flush
        await producer.produce(topic="test", value="test2")

        # Wait for the batch operation to be called
        await asyncio.wait_for(batch_called.wait(), timeout=2.0)

        # Batch was called and should have captured messages with callbacks
        assert captured_messages is not None
        assert len(captured_messages) == 2
        assert not first_future.done()

        # Simulate delayed delivery callback (like from background polling)
        mock_msg = Mock()
        mock_msg.topic.return_value = "test"
        mock_msg.value.return_value = b"test1"

        # Call callback for first message (index 0)
        first_callback = captured_messages[0]['callback']
        first_callback(None, mock_msg)

        # Now the Future should resolve
        result = await first_future
        assert result == mock_msg

        await producer.close()

    @pytest.mark.asyncio
    async def test_buffer_timeout_background_task(self, mock_producer, mock_common, basic_config):
        """Test that buffer timeout task runs continuously in the background."""
        # Create producer with short timeout for testing
        producer = AIOProducer(basic_config, buffer_timeout=0.1)

        # Test that timeout task is created and running
        assert producer._buffer_timeout_manager._timeout_task is not None
        assert not producer._buffer_timeout_manager._timeout_task.done()
        assert producer._is_closed is False

        # Wait a bit to ensure task is running
        await asyncio.sleep(0.05)
        assert not producer._buffer_timeout_manager._timeout_task.done()

        # Close the producer
        await producer.close()

        # Verify task stops and producer is closed
        assert producer._is_closed is True
        await asyncio.sleep(0.1)  # Grace period for cleanup
        assert (producer._buffer_timeout_manager._timeout_task is None or
                producer._buffer_timeout_manager._timeout_task.done())

    @pytest.mark.asyncio
    async def test_multiple_concurrent_produce(self, mock_producer, mock_common, basic_config):
        """Test multiple concurrent produce operations with batching."""
        producer = AIOProducer(basic_config, max_workers=3, batch_size=1)  # Force immediate flush

        completed_produces = []
        batch_call_count = 0

        def mock_produce_batch(topic, messages, partition=None):
            nonlocal batch_call_count
            batch_call_count += 1

            # Simulate successful delivery for each message in batch
            for i, msg_data in enumerate(messages):
                mock_msg = Mock()
                mock_msg.topic.return_value = topic
                mock_msg.value.return_value = (
                    msg_data['value'].encode() if isinstance(
                        msg_data['value'], str) else msg_data['value'])

                completed_produces.append((topic, msg_data['value']))
                # Call the individual message callback
                if 'callback' in msg_data:
                    msg_data['callback'](None, mock_msg)

        mock_producer.return_value.produce_batch.side_effect = mock_produce_batch
        mock_producer.return_value.poll.return_value = 1

        # Start multiple produce operations concurrently
        tasks = [
            asyncio.create_task(producer.produce(topic="test", value=f"msg{i}"))
            for i in range(3)
        ]

        # All tasks should complete successfully (tests real concurrency)
        results = await asyncio.gather(*tasks)

        # Verify all operations completed
        assert len(results) == 3
        assert all(result is not None for result in results)
        assert batch_call_count == 3  # Each message triggers its own batch due to batch_size=1
        assert len(completed_produces) == 3

        # Verify all messages were produced
        produced_values = [value for topic, value in completed_produces]
        expected_values = ["msg0", "msg1", "msg2"]
        assert sorted(produced_values) == sorted(expected_values)

        await producer.close()

    @pytest.mark.asyncio
    async def test_constructor_new_implementation(self, mock_producer, mock_common, basic_config):
        producer1 = AIOProducer(basic_config)
        assert producer1._batch_size == 1000
        assert isinstance(producer1.executor, concurrent.futures.ThreadPoolExecutor)
        assert hasattr(producer1, '_loop')
        assert hasattr(producer1, '_buffer_timeout_manager')
        assert producer1._batch_processor.is_buffer_empty()
        assert producer1._is_closed is False
        await producer1.close()

        custom_executor = concurrent.futures.ThreadPoolExecutor(max_workers=8)
        try:
            producer2 = AIOProducer(
                basic_config,
                executor=custom_executor,
                batch_size=500,
                buffer_timeout=10.0
            )
            assert producer2.executor is custom_executor
            assert producer2._batch_size == 500
            assert hasattr(producer2, '_buffer_timeout_manager')
            await producer2.close()
        finally:
            custom_executor.shutdown(wait=True)

    @pytest.mark.asyncio
    async def test_lifecycle_management_new_implementation(self, mock_producer, mock_common, basic_config):
        """Test lifecycle management for current implementation."""

        # Test close method with messages in buffer
        producer = AIOProducer(basic_config)

        # Add some messages to buffer
        with patch.object(producer, '_flush_buffer'):
            await producer.produce('test', 'msg1')
            assert producer._batch_processor.get_buffer_size() == 1

            # Test close method
            await producer.close()
            assert producer._is_closed is True
            assert (producer._buffer_timeout_manager._timeout_task is None or
                    producer._buffer_timeout_manager._timeout_task.done())

    @pytest.mark.asyncio
    async def test_buffer_timeout_task_management(self, mock_producer, mock_common, basic_config):
        """Test timeout task lifecycle and weak references."""

        # Test task creation and configuration
        producer = AIOProducer(basic_config, buffer_timeout=1.0)
        assert producer._buffer_timeout_manager._timeout_task is not None
        assert not producer._buffer_timeout_manager._timeout_task.done()
        assert producer._buffer_timeout_manager._timeout == 1.0
        assert producer._is_closed is False

        # Test task stops on close
        await producer.close()
        assert producer._is_closed is True
        assert (producer._buffer_timeout_manager._timeout_task is None or
                producer._buffer_timeout_manager._timeout_task.done())

    @pytest.mark.asyncio
    async def test_buffer_timeout_behavior(self, mock_producer, mock_common, basic_config):
        """Test buffer activity tracking and timeout triggers."""

        # Test buffer activity tracking
        producer = AIOProducer(basic_config)
        initial_time = producer._buffer_timeout_manager._last_activity
        assert initial_time > 0

        # Activity updates on produce
        await asyncio.sleep(0.01)  # Ensure time difference
        with patch.object(producer, '_flush_buffer'):  # Prevent auto-flush
            await producer.produce('test', 'msg1')
        assert producer._buffer_timeout_manager._last_activity > initial_time

        await producer.close()

    @pytest.mark.asyncio
    async def test_poll_method_new_implementation(self, mock_producer, mock_common, basic_config):
        """Test poll method with different timeout scenarios."""
        producer = AIOProducer(basic_config)

        # Test timeout=0 (non-blocking)
        with patch.object(producer, '_call') as mock_call:
            await producer.poll(timeout=0)
            mock_call.assert_called_with(producer._producer.poll, 0)

        # Test positive timeout (blocking via ThreadPool)
        with patch.object(producer, '_call') as mock_call:
            await producer.poll(timeout=5)
            mock_call.assert_called_with(producer._producer.poll, 5)

        await producer.close()

    @pytest.mark.asyncio
    async def test_produce_method_batching(self, mock_producer, mock_common, basic_config):
        """Test produce method with batching behavior."""
        producer = AIOProducer(basic_config, batch_size=3)

        # Test basic produce adds to buffer
        with patch.object(producer, '_flush_buffer') as mock_flush:
            future1 = await producer.produce('topic1', 'value1', key='key1')
            assert producer._batch_processor.get_buffer_size() == 1
            assert len(producer._batch_processor._buffer_futures) == 1
            assert isinstance(future1, asyncio.Future)
            mock_flush.assert_not_called()  # Should not flush yet

        # Test batch size trigger (3rd message should trigger flush)
        with patch.object(producer, '_flush_buffer') as mock_flush:
            await producer.produce('topic1', 'value2')  # 2nd message
            await producer.produce('topic1', 'value3')  # 3rd message - should trigger flush
            mock_flush.assert_called()

        await producer.close()

    @pytest.mark.asyncio
    async def test_flush_and_purge_methods_new_implementation(self, mock_producer, mock_common, basic_config):
        """Test flush and purge methods for current implementation."""
        producer = AIOProducer(basic_config)

        # Add messages to buffer
        with patch.object(producer, '_flush_buffer'):  # Prevent auto-flush
            await producer.produce('test', 'msg1')
            await producer.produce('test', 'msg2')
        assert producer._batch_processor.get_buffer_size() == 2

        # Test purge clears buffers
        with patch.object(producer, '_call') as mock_call:
            await producer.purge()
            mock_call.assert_called_with(producer._producer.purge)

        assert producer._batch_processor.get_buffer_size() == 0
        assert len(producer._batch_processor._buffer_futures) == 0

        await producer.close()

    @pytest.mark.asyncio
    async def test_group_messages_by_topic(self, mock_producer, mock_common, basic_config):
        """Test message grouping by topic for batch processing."""
        producer = AIOProducer(basic_config)

        # Test empty buffer
        groups = producer._batch_processor._group_messages_by_topic_and_partition()
        assert groups == {}

        # Add mixed topic messages
        producer._batch_processor._message_buffer = [
            {'topic': 'topic1', 'value': 'msg1', 'user_callback': None},
            {'topic': 'topic2', 'value': 'msg2', 'user_callback': Mock()},
            {'topic': 'topic1', 'value': 'msg3', 'user_callback': None},
        ]
        producer._batch_processor._buffer_futures = [Mock(), Mock(), Mock()]

        groups = producer._batch_processor._group_messages_by_topic_and_partition()

        # Test grouping correctness
        assert len(groups) == 2
        assert ('topic1', -1) in groups and ('topic2', -1) in groups
        assert len(groups[('topic1', -1)]['messages']) == 2  # msg1, msg3
        assert len(groups[('topic2', -1)]['messages']) == 1  # msg2

        await producer.close()

    @pytest.mark.asyncio
    async def test_error_handling_new_implementation(self, mock_producer, mock_common, basic_config):
        """Test error handling in current implementation."""
        producer = AIOProducer(basic_config)

        # Test batch error propagation
        producer._batch_processor._message_buffer = [{'topic': 'test', 'value': 'msg', 'user_callback': None}]
        producer._batch_processor._buffer_futures = [asyncio.Future()]

        with patch.object(producer._batch_processor, 'flush_buffer', side_effect=Exception("Batch failed")):
            with pytest.raises(Exception, match="Batch failed"):
                await producer._flush_buffer()

        await producer.close()

    @pytest.mark.asyncio
    async def test_future_based_usage_pattern(self, mock_producer, mock_common, basic_config):
        """Test the recommended Future-based usage pattern instead of callbacks."""
        producer = AIOProducer(basic_config, batch_size=1)

        # Mock successful delivery
        mock_msg = Mock()
        mock_msg.topic.return_value = "test-topic"
        mock_msg.value.return_value = b"test-value"

        async def mock_flush_buffer(target_topic=None):
            batches = producer._batch_processor.create_batches(target_topic)
            for batch in batches:
                for future in batch.futures:
                    if not future.done():
                        future.set_result(mock_msg)
            producer._batch_processor.clear_buffer()

        with patch.object(producer, '_flush_buffer', side_effect=mock_flush_buffer):
            # Recommended usage: await the Future returned by produce()
            future = await producer.produce(topic="test-topic", value="test-value")
            result = await future

            # Verify the result
            assert result is mock_msg
            assert result.topic() == "test-topic"
            assert result.value() == b"test-value"

        await producer.close()

    @pytest.mark.asyncio
    async def test_future_based_error_handling(self, mock_producer, mock_common, basic_config):
        """Test Future-based error handling pattern."""
        producer = AIOProducer(basic_config, batch_size=1)

        # Mock delivery error
        mock_error = KafkaException(KafkaError(KafkaError._MSG_TIMED_OUT))

        async def mock_flush_buffer(target_topic=None):
            batches = producer._batch_processor.create_batches(target_topic)
            for batch in batches:
                for future in batch.futures:
                    if not future.done():
                        future.set_exception(mock_error)
            producer._batch_processor.clear_buffer()

        with patch.object(producer, '_flush_buffer', side_effect=mock_flush_buffer):
            # Recommended usage: handle exceptions via Future
            future = await producer.produce(topic="test-topic", value="test-value")

            with pytest.raises(KafkaException):
                await future

        await producer.close()

    @pytest.mark.asyncio
    async def test_future_based_concurrent_usage(self, mock_producer, mock_common, basic_config):
        """Test Future-based concurrent usage pattern."""
        producer = AIOProducer(basic_config, batch_size=1)

        # Mock successful delivery
        mock_msg = Mock()
        mock_msg.topic.return_value = "test-topic"
        mock_msg.value.return_value = b"test-value"

        async def mock_flush_buffer(target_topic=None):
            batches = producer._batch_processor.create_batches(target_topic)
            for batch in batches:
                for future in batch.futures:
                    if not future.done():
                        future.set_result(mock_msg)
            producer._batch_processor.clear_buffer()

        with patch.object(producer, '_flush_buffer', side_effect=mock_flush_buffer):
            # Recommended usage: collect Futures and await them together
            futures = []
            for i in range(5):
                future = await producer.produce(topic="test-topic", value=f"test-value-{i}")
                futures.append(future)

            # Wait for all deliveries to complete
            results = await asyncio.gather(*futures)

            # Verify all results
            assert len(results) == 5
            for result in results:
                assert result is mock_msg

        await producer.close()

    @pytest.mark.asyncio
    async def test_edge_cases_batching(self, mock_producer, mock_common, basic_config):
        """Test edge cases in batching behavior."""
        producer = AIOProducer(basic_config, batch_size=100)

        # Test large batch handling
        with patch.object(producer, '_flush_buffer') as mock_flush:
            large_batch_tasks = [
                producer.produce('test', f'msg{i}')
                for i in range(150)  # Exceeds batch_size
            ]

            # Should trigger flush automatically at 100
            await asyncio.gather(*large_batch_tasks)
            assert mock_flush.call_count >= 1  # At least one flush

        await producer.close()