"""RabbitMQ extractor for streaming ETL pipelines.

Consumes messages from RabbitMQ queues using aio-pika, yielding batches
of parsed records with deferred acknowledgment support.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, AsyncIterator

from pycharter.etl_generator.extractors._messaging import AckTracker, MessagingConfig
from pycharter.etl_generator.extractors._streaming import BatchAccumulator
from pycharter.etl_generator.extractors.base import BaseExtractor

try:
    import aio_pika

    _HAS_AIO_PIKA = True
except ImportError:
    aio_pika = None  # type: ignore[assignment]
    _HAS_AIO_PIKA = False

logger = logging.getLogger(__name__)


def _require_aio_pika() -> None:
    """Raise a helpful error when aio-pika is not installed."""
    if not _HAS_AIO_PIKA:
        raise ImportError(
            "RabbitMQ extractor requires aio-pika. "
            "Install with: pip install pycharter[rabbitmq] or pip install aio-pika"
        )


class RabbitMQExtractor(BaseExtractor):
    """Extract data from RabbitMQ queues.

    Requires: ``pip install pycharter[rabbitmq]`` (aio-pika).

    Example (programmatic)::

        extractor = RabbitMQExtractor(
            queue="orders",
            url="amqp://guest:guest@localhost/",
        )
        pipeline = Pipeline(extractor) | Rename({...}) | PostgresLoader(...)

    Example (YAML)::

        extract:
          type: rabbitmq
          queue: orders
          url: amqp://guest:guest@localhost/

    Args:
        queue: Queue name to consume from.
        url: AMQP connection URL.
        exchange: Exchange name (empty string for default exchange).
        routing_key: Routing key for queue binding.
        prefetch_count: Maximum unacknowledged messages per channel.
        data_format: Message body format — ``'json'`` or ``'text'``.
        batch_size: Records per yielded batch.
        batch_timeout: Seconds before yielding a partial batch.
        shutdown_event: External event to signal graceful shutdown.
        max_records: Stop after this many total records.
        max_batches: Stop after yielding this many batches.
    """

    def __init__(
        self,
        queue: str,
        url: str = "amqp://guest:guest@localhost/",
        exchange: str = "",
        routing_key: str = "",
        prefetch_count: int = 100,
        data_format: str = "json",
        batch_size: int = 1000,
        batch_timeout: float = 5.0,
        shutdown_event: asyncio.Event | None = None,
        max_records: int | None = None,
        max_batches: int | None = None,
    ) -> None:
        _require_aio_pika()
        self._queue_name = queue
        self._url = url
        self._exchange = exchange
        self._routing_key = routing_key
        self._prefetch_count = prefetch_count
        self._data_format = data_format
        self._messaging_config = MessagingConfig(
            shutdown_event=shutdown_event,
            max_batches=max_batches,
            max_records=max_records,
            batch_size=batch_size,
            batch_timeout=batch_timeout,
        )
        self._connection: Any | None = None
        self._channel: Any | None = None
        self._ack_tracker = AckTracker()
        self._batch_index = 0

    # ------------------------------------------------------------------
    # Factory helpers
    # ------------------------------------------------------------------

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> RabbitMQExtractor:
        """Create a RabbitMQExtractor from a configuration dictionary.

        Args:
            config: Extractor configuration with RabbitMQ-specific keys.

        Returns:
            Configured RabbitMQExtractor instance.
        """
        return cls(
            queue=config["queue"],
            url=config.get("url", "amqp://guest:guest@localhost/"),
            exchange=config.get("exchange", ""),
            routing_key=config.get("routing_key", ""),
            prefetch_count=config.get("prefetch_count", 100),
            data_format=config.get("data_format", "json"),
            batch_size=config.get("batch_size", 1000),
            batch_timeout=config.get("batch_timeout", 5.0),
            shutdown_event=config.get("shutdown_event"),
            max_records=config.get("max_records"),
            max_batches=config.get("max_batches"),
        )

    def validate_config(self, extract_config: dict[str, Any]) -> None:
        """Validate RabbitMQ-specific configuration.

        Args:
            extract_config: Raw configuration dictionary.

        Raises:
            ValueError: If required fields are missing.
        """
        if not extract_config.get("queue"):
            raise ValueError("RabbitMQ extractor requires 'queue'")

    # ------------------------------------------------------------------
    # Connection lifecycle
    # ------------------------------------------------------------------

    async def _ensure_connection(self) -> tuple[Any, Any]:
        """Create connection and channel on first use."""
        if self._connection is not None and self._channel is not None:
            return self._connection, self._channel

        self._connection = await aio_pika.connect_robust(self._url)
        self._channel = await self._connection.channel()
        await self._channel.set_qos(prefetch_count=self._prefetch_count)
        logger.info(
            "RabbitMQ connected: queue=%s prefetch=%d",
            self._queue_name,
            self._prefetch_count,
        )
        return self._connection, self._channel

    # ------------------------------------------------------------------
    # Extract
    # ------------------------------------------------------------------

    async def extract(self, **params: Any) -> AsyncIterator[list[dict[str, Any]]]:
        """Consume messages from RabbitMQ and yield record batches.

        Yields:
            Batches of parsed RabbitMQ messages.
        """
        cfg = self._messaging_config
        _, channel = await self._ensure_connection()
        queue = await channel.declare_queue(self._queue_name, passive=True)
        accumulator = BatchAccumulator(cfg.batch_size, cfg.batch_timeout)
        total_records = 0
        total_batches = 0
        self._batch_index = 0
        self._ack_tracker.start_batch(self._batch_index)

        async with queue.iterator() as queue_iter:
            async for message in queue_iter:
                if cfg.shutdown_event and cfg.shutdown_event.is_set():
                    break

                record = self._parse_message(message)
                self._ack_tracker.track(self._batch_index, message)
                batch = await accumulator.add(record)
                total_records += 1

                if batch:
                    yield batch
                    total_batches += 1
                    self._batch_index += 1
                    self._ack_tracker.start_batch(self._batch_index)

                if cfg.max_records and total_records >= cfg.max_records:
                    break
                if cfg.max_batches and total_batches >= cfg.max_batches:
                    break

                # Flush partial batch on timeout
                if accumulator.should_flush_timeout():
                    remaining = await accumulator.flush()
                    if remaining:
                        yield remaining
                        total_batches += 1
                        self._batch_index += 1
                        self._ack_tracker.start_batch(self._batch_index)

                        if cfg.max_batches and total_batches >= cfg.max_batches:
                            break

        # Final flush
        remaining = await accumulator.flush()
        if remaining:
            yield remaining

    # ------------------------------------------------------------------
    # Acknowledgment
    # ------------------------------------------------------------------

    async def acknowledge(self, batch_index: int, success: bool) -> None:
        """Acknowledge a batch after processing.

        Args:
            batch_index: The batch number (0-indexed) to acknowledge.
            success: True to ack messages, False to nack with requeue.
        """
        messages = self._ack_tracker.get_batch_metadata(batch_index)
        if not messages:
            return

        for msg in messages:
            if success:
                msg.ack()
            else:
                msg.nack(requeue=True)

        action = "acked" if success else "nacked"
        logger.debug(
            "RabbitMQ batch %d %s (%d messages)",
            batch_index,
            action,
            len(messages),
        )
        self._ack_tracker.clear_batch(batch_index)

    # ------------------------------------------------------------------
    # Cleanup
    # ------------------------------------------------------------------

    async def close(self) -> None:
        """Close the RabbitMQ channel and connection."""
        if self._channel is not None:
            await self._channel.close()
            self._channel = None
        if self._connection is not None:
            await self._connection.close()
            self._connection = None
            logger.info("RabbitMQ connection closed")

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _parse_message(self, message: Any) -> dict[str, Any]:
        """Parse a single RabbitMQ message into a record dict."""
        body = message.body
        if isinstance(body, bytes):
            body = body.decode("utf-8", errors="replace")

        if self._data_format == "json":
            try:
                parsed = json.loads(body)
                if isinstance(parsed, dict):
                    parsed["_routing_key"] = message.routing_key
                    parsed["_exchange"] = message.exchange
                    return parsed
                return {
                    "data": parsed,
                    "_routing_key": message.routing_key,
                    "_exchange": message.exchange,
                }
            except (json.JSONDecodeError, TypeError):
                return {
                    "data": body,
                    "_routing_key": message.routing_key,
                    "_exchange": message.exchange,
                    "_parse_error": True,
                }

        return {
            "data": body,
            "_routing_key": message.routing_key,
            "_exchange": message.exchange,
        }

    async def extract_streaming(
        self,
        extract_config: dict[str, Any],
        params: dict[str, Any],
        headers: dict[str, Any],
        contract_dir: Any | None = None,
        batch_size: int = 1000,
        max_records: int | None = None,
        config_context: dict[str, Any] | None = None,
    ) -> AsyncIterator[list[dict[str, Any]]]:
        """Extract via config-driven path (delegates to extract)."""
        async for batch in self.extract(**params):
            yield batch
