"""Tests for messaging extractors (Kafka, RabbitMQ, SQS) and AckableExtractor protocol."""

from __future__ import annotations

import asyncio
import json
from typing import Any, AsyncIterator
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from pycharter.etl_generator.extractors._messaging import AckTracker, MessagingConfig
from pycharter.etl_generator.protocols import AckableExtractor, Extractor
from pycharter.etl_generator.result import LoadResult

# =============================================================================
# AckTracker
# =============================================================================


class TestAckTracker:
    """Tests for AckTracker metadata tracking."""

    def test_start_batch(self):
        """start_batch initialises an empty list."""
        tracker = AckTracker()
        tracker.start_batch(0)
        assert tracker.get_batch_metadata(0) == []

    def test_track(self):
        """track appends metadata to the correct batch."""
        tracker = AckTracker()
        tracker.start_batch(0)
        tracker.track(0, {"offset": 1})
        tracker.track(0, {"offset": 2})
        assert tracker.get_batch_metadata(0) == [{"offset": 1}, {"offset": 2}]

    def test_track_auto_starts_batch(self):
        """track creates the batch if start_batch was not called."""
        tracker = AckTracker()
        tracker.track(5, "meta")
        assert tracker.get_batch_metadata(5) == ["meta"]

    def test_get_batch_metadata_missing(self):
        """get_batch_metadata returns empty list for unknown batch."""
        tracker = AckTracker()
        assert tracker.get_batch_metadata(99) == []

    def test_clear_batch(self):
        """clear_batch removes tracking data."""
        tracker = AckTracker()
        tracker.start_batch(0)
        tracker.track(0, "a")
        tracker.clear_batch(0)
        assert tracker.get_batch_metadata(0) == []

    def test_clear_batch_noop_for_missing(self):
        """clear_batch does not raise for unknown batch."""
        tracker = AckTracker()
        tracker.clear_batch(99)

    def test_multiple_batches(self):
        """Separate batches tracked independently."""
        tracker = AckTracker()
        tracker.start_batch(0)
        tracker.start_batch(1)
        tracker.track(0, "a")
        tracker.track(1, "b")
        assert tracker.get_batch_metadata(0) == ["a"]
        assert tracker.get_batch_metadata(1) == ["b"]


# =============================================================================
# MessagingConfig
# =============================================================================


class TestMessagingConfig:
    """Tests for MessagingConfig defaults and custom values."""

    def test_defaults(self):
        """Default values are sensible."""
        cfg = MessagingConfig()
        assert cfg.shutdown_event is None
        assert cfg.max_batches is None
        assert cfg.max_records is None
        assert cfg.batch_size == 1000
        assert cfg.batch_timeout == 5.0
        assert cfg.consumer_group == "pycharter"
        assert cfg.auto_acknowledge is False

    def test_custom_values(self):
        """Custom values are stored."""
        event = asyncio.Event()
        cfg = MessagingConfig(
            shutdown_event=event,
            max_batches=5,
            max_records=10000,
            batch_size=500,
            batch_timeout=2.0,
            consumer_group="my-group",
            auto_acknowledge=True,
        )
        assert cfg.shutdown_event is event
        assert cfg.max_batches == 5
        assert cfg.max_records == 10000
        assert cfg.batch_size == 500
        assert cfg.batch_timeout == 2.0
        assert cfg.consumer_group == "my-group"
        assert cfg.auto_acknowledge is True


# =============================================================================
# Kafka Extractor
# =============================================================================


def _make_kafka_message(
    topic: str, partition: int, offset: int, value: dict[str, Any]
) -> MagicMock:
    """Create a mock Kafka message."""
    msg = MagicMock()
    msg.topic = topic
    msg.partition = partition
    msg.offset = offset
    msg.value = json.dumps(value).encode("utf-8")
    return msg


def _make_topic_partition(topic: str, partition: int) -> MagicMock:
    """Create a mock TopicPartition."""
    tp = MagicMock()
    tp.topic = topic
    tp.partition = partition
    return tp


class TestKafkaExtractor:
    """Tests for KafkaExtractor."""

    @patch("pycharter.etl_generator.extractors.kafka._HAS_AIOKAFKA", True)
    @patch("pycharter.etl_generator.extractors.kafka.AIOKafkaConsumer")
    @patch("pycharter.etl_generator.extractors.kafka.TopicPartition")
    @patch("pycharter.etl_generator.extractors.kafka.OffsetAndMetadata")
    async def test_from_config(self, mock_oam, mock_tp, mock_consumer_cls):
        """from_config creates extractor with correct settings."""
        from pycharter.etl_generator.extractors.kafka import KafkaExtractor

        config = {
            "topics": ["orders", "payments"],
            "bootstrap_servers": "broker:9092",
            "consumer_group": "etl-test",
            "batch_size": 500,
        }
        extractor = KafkaExtractor.from_config(config)
        assert extractor._topics == ["orders", "payments"]
        assert extractor._bootstrap_servers == "broker:9092"
        assert extractor._messaging_config.consumer_group == "etl-test"
        assert extractor._messaging_config.batch_size == 500

    @patch("pycharter.etl_generator.extractors.kafka._HAS_AIOKAFKA", True)
    @patch("pycharter.etl_generator.extractors.kafka.AIOKafkaConsumer")
    @patch("pycharter.etl_generator.extractors.kafka.TopicPartition")
    @patch("pycharter.etl_generator.extractors.kafka.OffsetAndMetadata")
    async def test_validate_config_missing_topics(
        self, mock_oam, mock_tp, mock_consumer_cls
    ):
        """validate_config raises on missing topics."""
        from pycharter.etl_generator.extractors.kafka import KafkaExtractor

        extractor = KafkaExtractor(topics=["t1"])
        with pytest.raises(ValueError, match="topics"):
            extractor.validate_config({})

    @patch("pycharter.etl_generator.extractors.kafka._HAS_AIOKAFKA", True)
    @patch("pycharter.etl_generator.extractors.kafka.AIOKafkaConsumer")
    @patch("pycharter.etl_generator.extractors.kafka.TopicPartition")
    @patch("pycharter.etl_generator.extractors.kafka.OffsetAndMetadata")
    async def test_extract_yields_batches(
        self, mock_oam, mock_tp_cls, mock_consumer_cls
    ):
        """extract yields batches from polled messages."""
        from pycharter.etl_generator.extractors.kafka import KafkaExtractor

        tp = _make_topic_partition("orders", 0)
        mock_tp_cls.return_value = tp
        msgs = [_make_kafka_message("orders", 0, i, {"id": i}) for i in range(3)]

        consumer = AsyncMock()
        consumer.getmany = AsyncMock(side_effect=[{tp: msgs}, {}])
        consumer.start = AsyncMock()
        consumer.stop = AsyncMock()
        mock_consumer_cls.return_value = consumer

        extractor = KafkaExtractor(topics=["orders"], batch_size=3, max_batches=1)
        batches = []
        async for batch in extractor.extract():
            batches.append(batch)

        assert len(batches) >= 1
        assert all(isinstance(b, list) for b in batches)
        assert batches[0][0]["id"] == 0
        await extractor.close()

    @patch("pycharter.etl_generator.extractors.kafka._HAS_AIOKAFKA", True)
    @patch("pycharter.etl_generator.extractors.kafka.AIOKafkaConsumer")
    @patch("pycharter.etl_generator.extractors.kafka.TopicPartition")
    @patch("pycharter.etl_generator.extractors.kafka.OffsetAndMetadata")
    async def test_acknowledge_commits_offsets(
        self, mock_oam_cls, mock_tp_cls, mock_consumer_cls
    ):
        """acknowledge(success=True) commits offsets to Kafka."""
        from pycharter.etl_generator.extractors.kafka import KafkaExtractor

        mock_tp_cls.side_effect = lambda topic, partition: MagicMock(
            topic=topic, partition=partition
        )
        mock_oam_cls.side_effect = lambda offset, meta: MagicMock(
            offset=offset, metadata=meta
        )

        consumer = AsyncMock()
        consumer.start = AsyncMock()
        consumer.commit = AsyncMock()
        mock_consumer_cls.return_value = consumer

        extractor = KafkaExtractor(topics=["t1"])
        extractor._consumer = consumer

        extractor._ack_tracker.start_batch(0)
        extractor._ack_tracker.track(0, {"topic": "t1", "partition": 0, "offset": 5})
        extractor._ack_tracker.track(0, {"topic": "t1", "partition": 0, "offset": 6})

        await extractor.acknowledge(0, success=True)
        consumer.commit.assert_awaited_once()

    @patch("pycharter.etl_generator.extractors.kafka._HAS_AIOKAFKA", True)
    @patch("pycharter.etl_generator.extractors.kafka.AIOKafkaConsumer")
    @patch("pycharter.etl_generator.extractors.kafka.TopicPartition")
    @patch("pycharter.etl_generator.extractors.kafka.OffsetAndMetadata")
    async def test_acknowledge_failure_seeks_back(
        self, mock_oam_cls, mock_tp_cls, mock_consumer_cls
    ):
        """acknowledge(success=False) seeks back for retry."""
        from pycharter.etl_generator.extractors.kafka import KafkaExtractor

        mock_tp_cls.side_effect = lambda topic, partition: f"{topic}-{partition}"

        consumer = AsyncMock()
        consumer.start = AsyncMock()
        consumer.seek = MagicMock()
        mock_consumer_cls.return_value = consumer

        extractor = KafkaExtractor(topics=["t1"])
        extractor._consumer = consumer

        extractor._ack_tracker.start_batch(0)
        extractor._ack_tracker.track(0, {"topic": "t1", "partition": 0, "offset": 10})
        extractor._ack_tracker.track(0, {"topic": "t1", "partition": 0, "offset": 11})

        await extractor.acknowledge(0, success=False)
        consumer.seek.assert_called_once()

    @patch("pycharter.etl_generator.extractors.kafka._HAS_AIOKAFKA", True)
    @patch("pycharter.etl_generator.extractors.kafka.AIOKafkaConsumer")
    @patch("pycharter.etl_generator.extractors.kafka.TopicPartition")
    @patch("pycharter.etl_generator.extractors.kafka.OffsetAndMetadata")
    async def test_shutdown_event_stops_extraction(
        self, mock_oam, mock_tp, mock_consumer_cls
    ):
        """Extraction stops when shutdown event is set."""
        from pycharter.etl_generator.extractors.kafka import KafkaExtractor

        shutdown = asyncio.Event()
        shutdown.set()

        consumer = AsyncMock()
        consumer.start = AsyncMock()
        consumer.stop = AsyncMock()
        mock_consumer_cls.return_value = consumer

        extractor = KafkaExtractor(topics=["t1"], shutdown_event=shutdown)
        batches = []
        async for batch in extractor.extract():
            batches.append(batch)

        assert batches == []

    @patch("pycharter.etl_generator.extractors.kafka._HAS_AIOKAFKA", True)
    @patch("pycharter.etl_generator.extractors.kafka.AIOKafkaConsumer")
    @patch("pycharter.etl_generator.extractors.kafka.TopicPartition")
    @patch("pycharter.etl_generator.extractors.kafka.OffsetAndMetadata")
    async def test_max_records_limit(self, mock_oam, mock_tp_cls, mock_consumer_cls):
        """Extraction stops after max_records."""
        from pycharter.etl_generator.extractors.kafka import KafkaExtractor

        tp = _make_topic_partition("t", 0)
        mock_tp_cls.return_value = tp
        msgs = [_make_kafka_message("t", 0, i, {"id": i}) for i in range(10)]

        consumer = AsyncMock()
        consumer.getmany = AsyncMock(side_effect=[{tp: msgs}, {}])
        consumer.start = AsyncMock()
        mock_consumer_cls.return_value = consumer

        extractor = KafkaExtractor(topics=["t"], batch_size=100, max_records=5)
        total = 0
        async for batch in extractor.extract():
            total += len(batch)

        assert total <= 5

    @patch("pycharter.etl_generator.extractors.kafka._HAS_AIOKAFKA", True)
    @patch("pycharter.etl_generator.extractors.kafka.AIOKafkaConsumer")
    @patch("pycharter.etl_generator.extractors.kafka.TopicPartition")
    @patch("pycharter.etl_generator.extractors.kafka.OffsetAndMetadata")
    async def test_close_stops_consumer(self, mock_oam, mock_tp, mock_consumer_cls):
        """close() stops the consumer."""
        from pycharter.etl_generator.extractors.kafka import KafkaExtractor

        consumer = AsyncMock()
        consumer.start = AsyncMock()
        consumer.stop = AsyncMock()
        mock_consumer_cls.return_value = consumer

        extractor = KafkaExtractor(topics=["t1"])
        extractor._consumer = consumer
        await extractor.close()

        consumer.stop.assert_awaited_once()
        assert extractor._consumer is None

    def test_missing_aiokafka_raises_import_error(self):
        """Instantiation fails with clear message when aiokafka missing."""
        with patch("pycharter.etl_generator.extractors.kafka._HAS_AIOKAFKA", False):
            from pycharter.etl_generator.extractors.kafka import KafkaExtractor

            with pytest.raises(ImportError, match="aiokafka"):
                KafkaExtractor(topics=["t1"])


# =============================================================================
# RabbitMQ Extractor
# =============================================================================


def _make_rabbitmq_message(body: dict[str, Any], routing_key: str = "") -> MagicMock:
    """Create a mock RabbitMQ message."""
    msg = MagicMock()
    msg.body = json.dumps(body).encode("utf-8")
    msg.routing_key = routing_key
    msg.exchange = ""
    msg.ack = MagicMock()
    msg.nack = MagicMock()
    return msg


class TestRabbitMQExtractor:
    """Tests for RabbitMQExtractor."""

    @patch("pycharter.etl_generator.extractors.rabbitmq._HAS_AIO_PIKA", True)
    @patch("pycharter.etl_generator.extractors.rabbitmq.aio_pika")
    async def test_from_config(self, mock_aio_pika):
        """from_config creates extractor with correct settings."""
        from pycharter.etl_generator.extractors.rabbitmq import RabbitMQExtractor

        config = {
            "queue": "orders",
            "url": "amqp://user:pass@host/",
            "prefetch_count": 50,
        }
        extractor = RabbitMQExtractor.from_config(config)
        assert extractor._queue_name == "orders"
        assert extractor._url == "amqp://user:pass@host/"
        assert extractor._prefetch_count == 50

    @patch("pycharter.etl_generator.extractors.rabbitmq._HAS_AIO_PIKA", True)
    @patch("pycharter.etl_generator.extractors.rabbitmq.aio_pika")
    async def test_validate_config_missing_queue(self, mock_aio_pika):
        """validate_config raises on missing queue."""
        from pycharter.etl_generator.extractors.rabbitmq import RabbitMQExtractor

        extractor = RabbitMQExtractor(queue="q")
        with pytest.raises(ValueError, match="queue"):
            extractor.validate_config({})

    @patch("pycharter.etl_generator.extractors.rabbitmq._HAS_AIO_PIKA", True)
    @patch("pycharter.etl_generator.extractors.rabbitmq.aio_pika")
    async def test_acknowledge_acks_messages(self, mock_aio_pika):
        """acknowledge(success=True) acks each message."""
        from pycharter.etl_generator.extractors.rabbitmq import RabbitMQExtractor

        extractor = RabbitMQExtractor(queue="q")
        msgs = [_make_rabbitmq_message({"id": i}) for i in range(3)]
        extractor._ack_tracker.start_batch(0)
        for msg in msgs:
            extractor._ack_tracker.track(0, msg)

        await extractor.acknowledge(0, success=True)

        for msg in msgs:
            msg.ack.assert_called_once()
            msg.nack.assert_not_called()

    @patch("pycharter.etl_generator.extractors.rabbitmq._HAS_AIO_PIKA", True)
    @patch("pycharter.etl_generator.extractors.rabbitmq.aio_pika")
    async def test_acknowledge_failure_nacks_with_requeue(self, mock_aio_pika):
        """acknowledge(success=False) nacks with requeue."""
        from pycharter.etl_generator.extractors.rabbitmq import RabbitMQExtractor

        extractor = RabbitMQExtractor(queue="q")
        msgs = [_make_rabbitmq_message({"id": 1})]
        extractor._ack_tracker.start_batch(0)
        for msg in msgs:
            extractor._ack_tracker.track(0, msg)

        await extractor.acknowledge(0, success=False)

        for msg in msgs:
            msg.nack.assert_called_once_with(requeue=True)
            msg.ack.assert_not_called()

    @patch("pycharter.etl_generator.extractors.rabbitmq._HAS_AIO_PIKA", True)
    @patch("pycharter.etl_generator.extractors.rabbitmq.aio_pika")
    async def test_close_cleans_up(self, mock_aio_pika):
        """close() closes channel and connection."""
        from pycharter.etl_generator.extractors.rabbitmq import RabbitMQExtractor

        extractor = RabbitMQExtractor(queue="q")
        extractor._channel = AsyncMock()
        extractor._connection = AsyncMock()

        await extractor.close()

        extractor._channel is None
        extractor._connection is None

    def test_missing_aio_pika_raises_import_error(self):
        """Instantiation fails with clear message when aio-pika missing."""
        with patch("pycharter.etl_generator.extractors.rabbitmq._HAS_AIO_PIKA", False):
            from pycharter.etl_generator.extractors.rabbitmq import RabbitMQExtractor

            with pytest.raises(ImportError, match="aio-pika"):
                RabbitMQExtractor(queue="q")


# =============================================================================
# SQS Extractor
# =============================================================================


class TestSQSExtractor:
    """Tests for SQSExtractor."""

    @patch("pycharter.etl_generator.extractors.sqs._HAS_BOTO3", True)
    @patch("pycharter.etl_generator.extractors.sqs.boto3")
    async def test_from_config(self, mock_boto3):
        """from_config creates extractor with correct settings."""
        from pycharter.etl_generator.extractors.sqs import SQSExtractor

        config = {
            "queue_url": "https://sqs.us-east-1.amazonaws.com/123/orders",
            "region": "us-west-2",
            "wait_time_seconds": 10,
        }
        extractor = SQSExtractor.from_config(config)
        assert extractor._queue_url == config["queue_url"]
        assert extractor._region == "us-west-2"
        assert extractor._wait_time_seconds == 10

    @patch("pycharter.etl_generator.extractors.sqs._HAS_BOTO3", True)
    @patch("pycharter.etl_generator.extractors.sqs.boto3")
    async def test_validate_config_missing_queue_url(self, mock_boto3):
        """validate_config raises on missing queue_url."""
        from pycharter.etl_generator.extractors.sqs import SQSExtractor

        extractor = SQSExtractor(queue_url="https://q")
        with pytest.raises(ValueError, match="queue_url"):
            extractor.validate_config({})

    @patch("pycharter.etl_generator.extractors.sqs._HAS_BOTO3", True)
    @patch("pycharter.etl_generator.extractors.sqs.boto3")
    async def test_acknowledge_deletes_messages(self, mock_boto3):
        """acknowledge(success=True) deletes messages."""
        from pycharter.etl_generator.extractors.sqs import SQSExtractor

        client = MagicMock()
        mock_boto3.client.return_value = client

        extractor = SQSExtractor(queue_url="https://q")
        extractor._client = client
        extractor._ack_tracker.start_batch(0)
        extractor._ack_tracker.track(0, {"receipt_handle": "rh1", "message_id": "m1"})
        extractor._ack_tracker.track(0, {"receipt_handle": "rh2", "message_id": "m2"})

        await extractor.acknowledge(0, success=True)

        client.delete_message_batch.assert_called_once()
        call_kwargs = client.delete_message_batch.call_args
        entries = call_kwargs.kwargs.get("Entries") or call_kwargs[1].get("Entries")
        assert len(entries) == 2

    @patch("pycharter.etl_generator.extractors.sqs._HAS_BOTO3", True)
    @patch("pycharter.etl_generator.extractors.sqs.boto3")
    async def test_acknowledge_failure_resets_visibility(self, mock_boto3):
        """acknowledge(success=False) resets visibility timeout."""
        from pycharter.etl_generator.extractors.sqs import SQSExtractor

        client = MagicMock()
        mock_boto3.client.return_value = client

        extractor = SQSExtractor(queue_url="https://q")
        extractor._client = client
        extractor._ack_tracker.start_batch(0)
        extractor._ack_tracker.track(0, {"receipt_handle": "rh1", "message_id": "m1"})

        await extractor.acknowledge(0, success=False)

        client.change_message_visibility_batch.assert_called_once()

    @patch("pycharter.etl_generator.extractors.sqs._HAS_BOTO3", True)
    @patch("pycharter.etl_generator.extractors.sqs.boto3")
    async def test_max_messages_per_poll_capped_at_10(self, mock_boto3):
        """max_messages_per_poll is capped at SQS limit of 10."""
        from pycharter.etl_generator.extractors.sqs import SQSExtractor

        extractor = SQSExtractor(queue_url="https://q", max_messages_per_poll=20)
        assert extractor._max_messages_per_poll == 10

    @patch("pycharter.etl_generator.extractors.sqs._HAS_BOTO3", True)
    @patch("pycharter.etl_generator.extractors.sqs.boto3")
    async def test_close_is_noop(self, mock_boto3):
        """close() clears client reference."""
        from pycharter.etl_generator.extractors.sqs import SQSExtractor

        extractor = SQSExtractor(queue_url="https://q")
        extractor._client = MagicMock()
        await extractor.close()
        assert extractor._client is None

    def test_missing_boto3_raises_import_error(self):
        """Instantiation fails with clear message when boto3 missing."""
        with patch("pycharter.etl_generator.extractors.sqs._HAS_BOTO3", False):
            from pycharter.etl_generator.extractors.sqs import SQSExtractor

            with pytest.raises(ImportError, match="boto3"):
                SQSExtractor(queue_url="https://q")

    @patch("pycharter.etl_generator.extractors.sqs._HAS_BOTO3", True)
    @patch("pycharter.etl_generator.extractors.sqs.boto3")
    async def test_extract_yields_batches(self, mock_boto3):
        """extract yields batches from polled SQS messages."""
        from pycharter.etl_generator.extractors.sqs import SQSExtractor

        client = MagicMock()
        mock_boto3.client.return_value = client

        sqs_messages = [
            {
                "MessageId": f"m{i}",
                "ReceiptHandle": f"rh{i}",
                "Body": json.dumps({"id": i}),
            }
            for i in range(3)
        ]

        # Return messages on first call, empty on second to trigger flush
        client.receive_message.side_effect = [
            {"Messages": sqs_messages},
            {"Messages": []},
            {"Messages": []},
        ]

        extractor = SQSExtractor(queue_url="https://q", batch_size=3, max_batches=1)

        batches = []
        async for batch in extractor.extract():
            batches.append(batch)

        assert len(batches) >= 1
        assert batches[0][0]["id"] == 0

    @patch("pycharter.etl_generator.extractors.sqs._HAS_BOTO3", True)
    @patch("pycharter.etl_generator.extractors.sqs.boto3")
    async def test_shutdown_event_stops_extraction(self, mock_boto3):
        """Extraction stops when shutdown event is set."""
        from pycharter.etl_generator.extractors.sqs import SQSExtractor

        shutdown = asyncio.Event()
        shutdown.set()

        extractor = SQSExtractor(queue_url="https://q", shutdown_event=shutdown)
        batches = []
        async for batch in extractor.extract():
            batches.append(batch)

        assert batches == []


# =============================================================================
# AckableExtractor Protocol
# =============================================================================


class TestAckableExtractorProtocol:
    """Tests for AckableExtractor protocol checks."""

    @patch("pycharter.etl_generator.extractors.kafka._HAS_AIOKAFKA", True)
    @patch("pycharter.etl_generator.extractors.kafka.AIOKafkaConsumer")
    @patch("pycharter.etl_generator.extractors.kafka.TopicPartition")
    @patch("pycharter.etl_generator.extractors.kafka.OffsetAndMetadata")
    def test_kafka_is_ackable(self, mock_oam, mock_tp, mock_consumer):
        """KafkaExtractor has acknowledge method (ackable interface)."""
        from pycharter.etl_generator.extractors.kafka import KafkaExtractor

        extractor = KafkaExtractor(topics=["t1"])
        assert hasattr(extractor, "acknowledge")

    @patch("pycharter.etl_generator.extractors.rabbitmq._HAS_AIO_PIKA", True)
    @patch("pycharter.etl_generator.extractors.rabbitmq.aio_pika")
    def test_rabbitmq_is_ackable(self, mock_aio_pika):
        """RabbitMQExtractor has acknowledge method (ackable interface)."""
        from pycharter.etl_generator.extractors.rabbitmq import RabbitMQExtractor

        extractor = RabbitMQExtractor(queue="q")
        assert hasattr(extractor, "acknowledge")

    @patch("pycharter.etl_generator.extractors.sqs._HAS_BOTO3", True)
    @patch("pycharter.etl_generator.extractors.sqs.boto3")
    def test_sqs_is_ackable(self, mock_boto3):
        """SQSExtractor has acknowledge method (ackable interface)."""
        from pycharter.etl_generator.extractors.sqs import SQSExtractor

        extractor = SQSExtractor(queue_url="https://q")
        assert hasattr(extractor, "acknowledge")

    def test_regular_extractor_not_ackable(self):
        """Regular extractors do not have acknowledge."""
        from pycharter.etl_generator.extractors.http import HTTPExtractor

        extractor = HTTPExtractor(url="https://example.com")
        assert not hasattr(extractor, "acknowledge")


# =============================================================================
# Pipeline Ack Integration
# =============================================================================


class _MockAckableExtractor:
    """Mock extractor with acknowledge support for pipeline integration tests."""

    def __init__(self, batches: list[list[dict[str, Any]]]) -> None:
        self._batches = batches
        self.ack_calls: list[tuple[int, bool]] = []

    async def extract(self, **params: Any) -> AsyncIterator[list[dict[str, Any]]]:
        """Yield pre-set batches."""
        for batch in self._batches:
            yield batch

    async def acknowledge(self, batch_index: int, success: bool) -> None:
        """Record acknowledgment call."""
        self.ack_calls.append((batch_index, success))


class _MockPlainExtractor:
    """Mock extractor without acknowledge for pipeline integration tests."""

    def __init__(self, batches: list[list[dict[str, Any]]]) -> None:
        self._batches = batches

    async def extract(self, **params: Any) -> AsyncIterator[list[dict[str, Any]]]:
        """Yield pre-set batches."""
        for batch in self._batches:
            yield batch


class TestPipelineAckIntegration:
    """Tests for pipeline acknowledge integration."""

    async def test_pipeline_calls_acknowledge_after_load(self):
        """Pipeline calls acknowledge after successful load."""
        from pycharter.etl_generator.pipeline import Pipeline

        extractor = _MockAckableExtractor(batches=[[{"id": 1}], [{"id": 2}]])
        loader = AsyncMock()
        loader.load = AsyncMock(return_value=LoadResult(success=True, rows_loaded=1))

        pipeline = Pipeline(extractor)
        pipeline.loader = loader

        result = await pipeline.run()

        assert result.success
        assert len(extractor.ack_calls) == 2
        assert extractor.ack_calls[0] == (0, True)
        assert extractor.ack_calls[1] == (1, True)

    async def test_pipeline_skips_ack_for_non_ackable(self):
        """Pipeline does not call acknowledge on plain extractors."""
        from pycharter.etl_generator.pipeline import Pipeline

        extractor = _MockPlainExtractor(batches=[[{"id": 1}]])
        loader = AsyncMock()
        loader.load = AsyncMock(return_value=LoadResult(success=True, rows_loaded=1))

        pipeline = Pipeline(extractor)
        pipeline.loader = loader

        result = await pipeline.run()
        assert result.success

    async def test_dry_run_acknowledges_with_success(self):
        """dry_run mode acknowledges with success=True."""
        from pycharter.etl_generator.pipeline import Pipeline

        extractor = _MockAckableExtractor(batches=[[{"id": 1}]])

        pipeline = Pipeline(extractor)
        result = await pipeline.run(dry_run=True)

        assert result.success
        assert len(extractor.ack_calls) == 1
        assert extractor.ack_calls[0] == (0, True)

    async def test_pipeline_ack_with_load_failure(self):
        """Pipeline sends success=False to acknowledge when load fails."""
        from pycharter.etl_generator.pipeline import Pipeline

        extractor = _MockAckableExtractor(batches=[[{"id": 1}]])
        loader = AsyncMock()
        loader.load = AsyncMock(
            return_value=LoadResult(success=False, rows_loaded=0, error="DB down")
        )

        pipeline = Pipeline(extractor)
        pipeline.loader = loader

        result = await pipeline.run()

        assert len(extractor.ack_calls) == 1
        assert extractor.ack_calls[0] == (0, False)

    async def test_pipeline_ack_exception_does_not_break_pipeline(self):
        """Pipeline continues even if acknowledge raises an exception."""
        from pycharter.etl_generator.pipeline import Pipeline

        extractor = _MockAckableExtractor(batches=[[{"id": 1}], [{"id": 2}]])
        # Make acknowledge raise on first call only
        original_ack = extractor.acknowledge

        call_count = 0

        async def _failing_ack(batch_index: int, success: bool) -> None:
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                raise RuntimeError("Ack failed")
            await original_ack(batch_index, success)

        extractor.acknowledge = _failing_ack

        loader = AsyncMock()
        loader.load = AsyncMock(return_value=LoadResult(success=True, rows_loaded=1))

        pipeline = Pipeline(extractor)
        pipeline.loader = loader

        result = await pipeline.run()

        # Pipeline still processes both batches despite ack failure
        assert result.batches_processed == 2


# =============================================================================
# Config Models
# =============================================================================


class TestMessagingConfigModels:
    """Tests for Kafka, RabbitMQ, and SQS Pydantic config models."""

    def test_kafka_config_defaults(self):
        """KafkaExtractConfig defaults are set correctly."""
        from pycharter.etl_generator.config_models import KafkaExtractConfig

        cfg = KafkaExtractConfig(topics=["t1"])
        assert cfg.type == "kafka"
        assert cfg.bootstrap_servers == "localhost:9092"
        assert cfg.consumer_group == "pycharter"
        assert cfg.batch_size == 1000

    def test_kafka_config_requires_topics(self):
        """KafkaExtractConfig requires topics."""
        from pydantic import ValidationError

        from pycharter.etl_generator.config_models import KafkaExtractConfig

        with pytest.raises(ValidationError):
            KafkaExtractConfig()

    def test_rabbitmq_config_defaults(self):
        """RabbitMQExtractConfig defaults are set correctly."""
        from pycharter.etl_generator.config_models import RabbitMQExtractConfig

        cfg = RabbitMQExtractConfig(queue="q")
        assert cfg.type == "rabbitmq"
        assert cfg.url == "amqp://guest:guest@localhost/"
        assert cfg.prefetch_count == 100

    def test_rabbitmq_config_requires_queue(self):
        """RabbitMQExtractConfig requires queue."""
        from pydantic import ValidationError

        from pycharter.etl_generator.config_models import RabbitMQExtractConfig

        with pytest.raises(ValidationError):
            RabbitMQExtractConfig()

    def test_sqs_config_defaults(self):
        """SQSExtractConfig defaults are set correctly."""
        from pycharter.etl_generator.config_models import SQSExtractConfig

        cfg = SQSExtractConfig(queue_url="https://q")
        assert cfg.type == "sqs"
        assert cfg.region == "us-east-1"
        assert cfg.wait_time_seconds == 20
        assert cfg.visibility_timeout == 300
        assert cfg.max_messages_per_poll == 10

    def test_sqs_config_requires_queue_url(self):
        """SQSExtractConfig requires queue_url."""
        from pydantic import ValidationError

        from pycharter.etl_generator.config_models import SQSExtractConfig

        with pytest.raises(ValidationError):
            SQSExtractConfig()

    def test_parse_extract_config_kafka(self):
        """parse_extract_config handles kafka type."""
        from pycharter.etl_generator.config_models import parse_extract_config

        config = {"type": "kafka", "topics": ["t1"]}
        result = parse_extract_config(config)
        assert result.type == "kafka"

    def test_parse_extract_config_rabbitmq(self):
        """parse_extract_config handles rabbitmq type."""
        from pycharter.etl_generator.config_models import parse_extract_config

        config = {"type": "rabbitmq", "queue": "q"}
        result = parse_extract_config(config)
        assert result.type == "rabbitmq"

    def test_parse_extract_config_sqs(self):
        """parse_extract_config handles sqs type."""
        from pycharter.etl_generator.config_models import parse_extract_config

        config = {"type": "sqs", "queue_url": "https://q"}
        result = parse_extract_config(config)
        assert result.type == "sqs"


# =============================================================================
# Factory Registration
# =============================================================================


class TestExtractorFactoryMessaging:
    """Tests for ExtractorFactory messaging type registration."""

    def test_factory_lists_messaging_types(self):
        """ExtractorFactory.list_types() includes kafka, rabbitmq, sqs."""
        from pycharter.etl_generator.extractors.factory import ExtractorFactory

        types = ExtractorFactory.list_types()
        assert "kafka" in types
        assert "rabbitmq" in types
        assert "sqs" in types
