"""AWS SQS extractor for streaming ETL pipelines.

Consumes messages from SQS queues using boto3, 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 boto3

    _HAS_BOTO3 = True
except ImportError:
    boto3 = None  # type: ignore[assignment]
    _HAS_BOTO3 = False

logger = logging.getLogger(__name__)


def _require_boto3() -> None:
    """Raise a helpful error when boto3 is not installed."""
    if not _HAS_BOTO3:
        raise ImportError(
            "SQS extractor requires boto3. "
            "Install with: pip install pycharter[etl] or pip install boto3"
        )


class SQSExtractor(BaseExtractor):
    """Extract data from AWS SQS queues.

    Uses boto3 (available in the ``[etl]`` extra).

    Example (programmatic)::

        extractor = SQSExtractor(
            queue_url="https://sqs.us-east-1.amazonaws.com/123/orders",
        )
        pipeline = Pipeline(extractor) | Rename({...}) | PostgresLoader(...)

    Example (YAML)::

        extract:
          type: sqs
          queue_url: https://sqs.us-east-1.amazonaws.com/123456789/orders
          region: us-east-1

    Args:
        queue_url: Full SQS queue URL.
        region: AWS region name.
        max_messages_per_poll: Messages per ``receive_message`` call (max 10).
        wait_time_seconds: Long-polling wait time in seconds.
        visibility_timeout: Seconds before unacknowledged messages reappear.
        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_url: str,
        region: str = "us-east-1",
        max_messages_per_poll: int = 10,
        wait_time_seconds: int = 20,
        visibility_timeout: int = 300,
        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_boto3()
        self._queue_url = queue_url
        self._region = region
        self._max_messages_per_poll = min(max_messages_per_poll, 10)
        self._wait_time_seconds = wait_time_seconds
        self._visibility_timeout = visibility_timeout
        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._client: Any | None = None
        self._ack_tracker = AckTracker()
        self._batch_index = 0

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

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

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

        Returns:
            Configured SQSExtractor instance.
        """
        return cls(
            queue_url=config["queue_url"],
            region=config.get("region", "us-east-1"),
            max_messages_per_poll=config.get("max_messages_per_poll", 10),
            wait_time_seconds=config.get("wait_time_seconds", 20),
            visibility_timeout=config.get("visibility_timeout", 300),
            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 SQS-specific configuration.

        Args:
            extract_config: Raw configuration dictionary.

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

    # ------------------------------------------------------------------
    # Client lifecycle
    # ------------------------------------------------------------------

    def _ensure_client(self) -> Any:
        """Create the SQS client on first use."""
        if self._client is not None:
            return self._client
        self._client = boto3.client("sqs", region_name=self._region)
        logger.info("SQS client created: region=%s", self._region)
        return self._client

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

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

        Yields:
            Batches of parsed SQS messages.
        """
        cfg = self._messaging_config
        client = self._ensure_client()
        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)

        while True:
            if cfg.shutdown_event and cfg.shutdown_event.is_set():
                break

            response = await asyncio.to_thread(
                client.receive_message,
                QueueUrl=self._queue_url,
                MaxNumberOfMessages=self._max_messages_per_poll,
                WaitTimeSeconds=self._wait_time_seconds,
                VisibilityTimeout=self._visibility_timeout,
            )

            messages = response.get("Messages", [])
            if not messages:
                # Flush on empty poll if we have buffered records
                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
                continue

            for msg in messages:
                record = self._parse_message(msg)
                self._ack_tracker.track(
                    self._batch_index,
                    {
                        "receipt_handle": msg["ReceiptHandle"],
                        "message_id": msg["MessageId"],
                    },
                )
                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
                if cfg.shutdown_event and cfg.shutdown_event.is_set():
                    break

            if self._should_stop(cfg, total_records, total_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 delete messages, False to reset visibility for retry.
        """
        metadata = self._ack_tracker.get_batch_metadata(batch_index)
        if not metadata:
            return

        client = self._ensure_client()

        if success:
            # Delete in chunks of 10 (SQS batch limit)
            for i in range(0, len(metadata), 10):
                chunk = metadata[i : i + 10]
                entries = [
                    {"Id": item["message_id"], "ReceiptHandle": item["receipt_handle"]}
                    for item in chunk
                ]
                await asyncio.to_thread(
                    client.delete_message_batch,
                    QueueUrl=self._queue_url,
                    Entries=entries,
                )
            logger.debug(
                "SQS batch %d deleted (%d messages)", batch_index, len(metadata)
            )
        else:
            # Reset visibility to 0 so messages are immediately available
            for i in range(0, len(metadata), 10):
                chunk = metadata[i : i + 10]
                entries = [
                    {
                        "Id": item["message_id"],
                        "ReceiptHandle": item["receipt_handle"],
                        "VisibilityTimeout": 0,
                    }
                    for item in chunk
                ]
                await asyncio.to_thread(
                    client.change_message_visibility_batch,
                    QueueUrl=self._queue_url,
                    Entries=entries,
                )
            logger.debug(
                "SQS batch %d visibility reset (%d messages)",
                batch_index,
                len(metadata),
            )

        self._ack_tracker.clear_batch(batch_index)

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

    async def close(self) -> None:
        """Release the SQS client (no persistent connection to close)."""
        self._client = None

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

    @staticmethod
    def _should_stop(
        cfg: MessagingConfig, total_records: int, total_batches: int
    ) -> bool:
        """Check whether extraction limits have been reached."""
        if cfg.shutdown_event and cfg.shutdown_event.is_set():
            return True
        if cfg.max_records and total_records >= cfg.max_records:
            return True
        if cfg.max_batches and total_batches >= cfg.max_batches:
            return True
        return False

    def _parse_message(self, msg: dict[str, Any]) -> dict[str, Any]:
        """Parse a single SQS message into a record dict."""
        body = msg.get("Body", "")

        if self._data_format == "json":
            try:
                parsed = json.loads(body)
                if isinstance(parsed, dict):
                    parsed["_message_id"] = msg.get("MessageId")
                    return parsed
                return {
                    "data": parsed,
                    "_message_id": msg.get("MessageId"),
                }
            except (json.JSONDecodeError, TypeError):
                return {
                    "data": body,
                    "_message_id": msg.get("MessageId"),
                    "_parse_error": True,
                }

        return {
            "data": body,
            "_message_id": msg.get("MessageId"),
        }

    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
