"""Apache Kafka extractor for streaming ETL pipelines.

Consumes messages from Kafka topics using aiokafka, 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:
    from aiokafka import AIOKafkaConsumer, TopicPartition
    from aiokafka.structs import OffsetAndMetadata

    _HAS_AIOKAFKA = True
except ImportError:
    AIOKafkaConsumer = None  # type: ignore[assignment,misc]
    TopicPartition = None  # type: ignore[assignment,misc]
    OffsetAndMetadata = None  # type: ignore[assignment,misc]
    _HAS_AIOKAFKA = False

logger = logging.getLogger(__name__)


def _require_aiokafka() -> None:
    """Raise a helpful error when aiokafka is not installed."""
    if not _HAS_AIOKAFKA:
        raise ImportError(
            "Kafka extractor requires aiokafka. "
            "Install with: pip install pycharter[kafka] or pip install aiokafka"
        )


class KafkaExtractor(BaseExtractor):
    """Extract data from Apache Kafka topics.

    Requires: ``pip install pycharter[kafka]`` (aiokafka).

    Example (programmatic)::

        extractor = KafkaExtractor(
            topics=["orders"],
            bootstrap_servers="localhost:9092",
            consumer_group="etl-orders",
        )
        pipeline = Pipeline(extractor) | Rename({...}) | PostgresLoader(...)

    Example (YAML)::

        extract:
          type: kafka
          topics: [orders]
          bootstrap_servers: localhost:9092
          consumer_group: etl-orders

    Args:
        topics: Kafka topic names to subscribe to.
        bootstrap_servers: Comma-separated broker addresses.
        consumer_group: Consumer group id.
        data_format: Message body format — ``'json'`` or ``'text'``.
        batch_size: Records per yielded batch.
        batch_timeout: Seconds before yielding a partial batch.
        poll_timeout: Seconds to wait on each poll call.
        auto_offset_reset: Where to start when no committed offset exists.
        security_protocol: Kafka security protocol.
        sasl_mechanism: SASL mechanism (e.g. ``'PLAIN'``, ``'SCRAM-SHA-256'``).
        sasl_username: SASL username.
        sasl_password: SASL password.
        ssl_cafile: Path to CA certificate for SSL.
        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,
        topics: list[str],
        bootstrap_servers: str = "localhost:9092",
        consumer_group: str = "pycharter",
        data_format: str = "json",
        batch_size: int = 1000,
        batch_timeout: float = 5.0,
        poll_timeout: float = 1.0,
        auto_offset_reset: str = "earliest",
        security_protocol: str = "PLAINTEXT",
        sasl_mechanism: str | None = None,
        sasl_username: str | None = None,
        sasl_password: str | None = None,
        ssl_cafile: str | None = None,
        shutdown_event: asyncio.Event | None = None,
        max_records: int | None = None,
        max_batches: int | None = None,
    ) -> None:
        _require_aiokafka()
        self._topics = topics
        self._bootstrap_servers = bootstrap_servers
        self._data_format = data_format
        self._poll_timeout = poll_timeout
        self._auto_offset_reset = auto_offset_reset
        self._security_protocol = security_protocol
        self._sasl_mechanism = sasl_mechanism
        self._sasl_username = sasl_username
        self._sasl_password = sasl_password
        self._ssl_cafile = ssl_cafile
        self._messaging_config = MessagingConfig(
            shutdown_event=shutdown_event,
            max_batches=max_batches,
            max_records=max_records,
            batch_size=batch_size,
            batch_timeout=batch_timeout,
            consumer_group=consumer_group,
        )
        self._consumer: Any | None = None
        self._ack_tracker = AckTracker()
        self._batch_index = 0

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

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

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

        Returns:
            Configured KafkaExtractor instance.
        """
        return cls(
            topics=config["topics"],
            bootstrap_servers=config.get("bootstrap_servers", "localhost:9092"),
            consumer_group=config.get("consumer_group", "pycharter"),
            data_format=config.get("data_format", "json"),
            batch_size=config.get("batch_size", 1000),
            batch_timeout=config.get("batch_timeout", 5.0),
            poll_timeout=config.get("poll_timeout", 1.0),
            auto_offset_reset=config.get("auto_offset_reset", "earliest"),
            security_protocol=config.get("security_protocol", "PLAINTEXT"),
            sasl_mechanism=config.get("sasl_mechanism"),
            sasl_username=config.get("sasl_username"),
            sasl_password=config.get("sasl_password"),
            ssl_cafile=config.get("ssl_cafile"),
            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 Kafka-specific configuration.

        Args:
            extract_config: Raw configuration dictionary.

        Raises:
            ValueError: If required fields are missing.
        """
        if not extract_config.get("topics"):
            raise ValueError("Kafka extractor requires 'topics' (list of topic names)")

    # ------------------------------------------------------------------
    # Consumer lifecycle
    # ------------------------------------------------------------------

    async def _ensure_consumer(self) -> Any:
        """Create and start the Kafka consumer on first use."""
        if self._consumer is not None:
            return self._consumer

        kwargs: dict[str, Any] = {
            "bootstrap_servers": self._bootstrap_servers,
            "group_id": self._messaging_config.consumer_group,
            "auto_offset_reset": self._auto_offset_reset,
            "enable_auto_commit": False,
            "security_protocol": self._security_protocol,
        }
        if self._sasl_mechanism:
            kwargs["sasl_mechanism"] = self._sasl_mechanism
        if self._sasl_username:
            kwargs["sasl_plain_username"] = self._sasl_username
        if self._sasl_password:
            kwargs["sasl_plain_password"] = self._sasl_password
        if self._ssl_cafile:
            kwargs["ssl_cafile"] = self._ssl_cafile

        self._consumer = AIOKafkaConsumer(*self._topics, **kwargs)
        await self._consumer.start()
        logger.info(
            "Kafka consumer started: topics=%s group=%s",
            self._topics,
            self._messaging_config.consumer_group,
        )
        return self._consumer

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

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

        Yields:
            Batches of parsed Kafka messages.
        """
        cfg = self._messaging_config
        consumer = await self._ensure_consumer()
        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)

        poll_timeout_ms = int(self._poll_timeout * 1000)

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

            data = await consumer.getmany(timeout_ms=poll_timeout_ms)

            for tp, messages in data.items():
                for msg in messages:
                    record = self._parse_message(msg)
                    self._ack_tracker.track(
                        self._batch_index,
                        {
                            "topic": tp.topic,
                            "partition": tp.partition,
                            "offset": msg.offset,
                        },
                    )
                    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

            if self._should_stop(cfg, total_records, total_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 commit offsets, False to seek back for retry.
        """
        metadata = self._ack_tracker.get_batch_metadata(batch_index)
        if not metadata:
            return

        consumer = self._consumer
        if consumer is None:
            return

        if success:
            offsets: dict[Any, Any] = {}
            for item in metadata:
                tp = TopicPartition(item["topic"], item["partition"])
                current = offsets.get(tp)
                new_offset = item["offset"] + 1
                if current is None or new_offset > current.offset:
                    offsets[tp] = OffsetAndMetadata(new_offset, "")
            await consumer.commit(offsets)
            logger.debug(
                "Kafka batch %d committed (%d offsets)", batch_index, len(offsets)
            )
        else:
            # Seek back to earliest offset in the batch for retry
            seek_targets: dict[Any, int] = {}
            for item in metadata:
                tp = TopicPartition(item["topic"], item["partition"])
                if tp not in seek_targets or item["offset"] < seek_targets[tp]:
                    seek_targets[tp] = item["offset"]
            for tp, offset in seek_targets.items():
                consumer.seek(tp, offset)
            logger.debug(
                "Kafka batch %d seek-back (%d partitions)",
                batch_index,
                len(seek_targets),
            )

        self._ack_tracker.clear_batch(batch_index)

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

    async def close(self) -> None:
        """Stop the Kafka consumer and release resources."""
        if self._consumer is not None:
            await self._consumer.stop()
            self._consumer = None
            logger.info("Kafka consumer stopped")

    # ------------------------------------------------------------------
    # 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: Any) -> dict[str, Any]:
        """Parse a single Kafka message into a record dict."""
        value = msg.value
        if isinstance(value, bytes):
            value = value.decode("utf-8", errors="replace")

        if self._data_format == "json":
            try:
                parsed = json.loads(value)
                if isinstance(parsed, dict):
                    parsed["_topic"] = msg.topic
                    parsed["_partition"] = msg.partition
                    parsed["_offset"] = msg.offset
                    return parsed
                return {
                    "data": parsed,
                    "_topic": msg.topic,
                    "_partition": msg.partition,
                    "_offset": msg.offset,
                }
            except (json.JSONDecodeError, TypeError):
                return {
                    "data": value,
                    "_topic": msg.topic,
                    "_partition": msg.partition,
                    "_offset": msg.offset,
                    "_parse_error": True,
                }

        return {
            "data": value,
            "_topic": msg.topic,
            "_partition": msg.partition,
            "_offset": msg.offset,
        }

    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
