"""Shared infrastructure for message queue extractors.

Provides common configuration and acknowledgment tracking used by
Kafka, RabbitMQ, and SQS extractors.
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from typing import Any


@dataclass(frozen=True, slots=True)
class MessagingConfig:
    """Common configuration for message queue extractors.

    Attributes:
        shutdown_event: External event to signal graceful shutdown.
        max_batches: Stop after yielding this many batches.
        max_records: Stop after processing this many total records.
        batch_size: Maximum records per batch before yielding.
        batch_timeout: Seconds before yielding a partial batch.
        consumer_group: Consumer group identifier.
        auto_acknowledge: If True, acknowledge immediately (at-most-once).
    """

    shutdown_event: asyncio.Event | None = None
    max_batches: int | None = None
    max_records: int | None = None
    batch_size: int = 1000
    batch_timeout: float = 5.0
    consumer_group: str = "pycharter"
    auto_acknowledge: bool = False


class AckTracker:
    """Tracks message metadata per batch for deferred acknowledgment.

    Each extractor stores connector-specific metadata (offsets, delivery
    tags, receipt handles) keyed by batch index.  After a batch is
    successfully loaded the extractor retrieves the metadata and
    acknowledges the messages with the broker.
    """

    def __init__(self) -> None:
        self._batches: dict[int, list[Any]] = {}

    def start_batch(self, batch_index: int) -> None:
        """Initialise tracking for a new batch.

        Args:
            batch_index: Zero-based batch number.
        """
        self._batches[batch_index] = []

    def track(self, batch_index: int, metadata: Any) -> None:
        """Record metadata for a single message within a batch.

        Args:
            batch_index: Batch the message belongs to.
            metadata: Connector-specific message metadata.
        """
        if batch_index not in self._batches:
            self._batches[batch_index] = []
        self._batches[batch_index].append(metadata)

    def get_batch_metadata(self, batch_index: int) -> list[Any]:
        """Return all tracked metadata for *batch_index*.

        Args:
            batch_index: Batch to retrieve.

        Returns:
            List of metadata items (empty list if batch not tracked).
        """
        return self._batches.get(batch_index, [])

    def clear_batch(self, batch_index: int) -> None:
        """Remove tracking data for *batch_index*.

        Args:
            batch_index: Batch to clear.
        """
        self._batches.pop(batch_index, None)
