"""Kafka client implementation using ``aiokafka``.

Provides :class:`KafkaConfig` for centralised configuration and
:class:`KafkaClient` which extends :class:`BrokerClient`.
"""

from __future__ import annotations

import asyncio
import json
import logging
from dataclasses import dataclass, field
from typing import Any, Dict, List, Literal, Optional, Set, TYPE_CHECKING
from uuid import uuid4

from .base import BrokerClient
from ..message.result import MessageResult, Subscription

if TYPE_CHECKING:
    from ..message.builder import MessageBuilder, Message

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# KafkaConfig — centralised configuration
# ---------------------------------------------------------------------------


@dataclass
class KafkaConfig:
    """Centralised configuration for a :class:`KafkaClient`.

    All producer and consumer settings live here.  Consumer settings are
    **defaults** that can be overridden per-subscription via the
    :class:`MessageBuilder` chain.

    Example::

        config = KafkaConfig(
            bootstrap_servers="localhost:9092",
            client_id="my-service",
            group_id="default-consumer",
        )
        client = KafkaClient(config)
    """

    # -- Connection -------------------------------------------------------
    bootstrap_servers: str
    client_id: Optional[str] = None

    # -- Producer ---------------------------------------------------------
    acks: Literal[0, 1, "all"] = "all"
    compression_type: Optional[str] = None  # "gzip" | "snappy" | "lz4" | "zstd"
    max_request_size: int = 1_048_576  # 1 MB

    # -- Consumer defaults (overridable per-subscription) -----------------
    group_id: Optional[str] = None
    auto_offset_reset: Literal["earliest", "latest"] = "latest"
    enable_auto_commit: bool = True
    auto_commit_interval_ms: int = 5000
    max_poll_records: int = 500
    session_timeout_ms: int = 10000
    max_poll_interval_ms: int = 300000  # 5 min

    # -- Security ---------------------------------------------------------
    security_protocol: Literal["PLAINTEXT", "SASL_PLAINTEXT", "SSL", "SASL_SSL"] = "PLAINTEXT"
    sasl_mechanism: Optional[str] = None  # "PLAIN" | "SCRAM-SHA-256" | "SCRAM-SHA-512"
    sasl_username: Optional[str] = None
    sasl_password: Optional[str] = None


# ---------------------------------------------------------------------------
# KafkaClient
# ---------------------------------------------------------------------------


class KafkaClient(BrokerClient):
    """Async Kafka client backed by ``aiokafka``.

    Implements the :class:`BrokerClient` contract.  The producer is created
    once on :meth:`connect` and shared across all publishes.  Each
    :meth:`_subscribe` call creates a dedicated ``AIOKafkaConsumer``
    instance.
    """

    def __init__(self, config: KafkaConfig):
        self.config = config
        self._producer: Any = None
        self._active_subscriptions: Set[Subscription] = set()
        self._connected = False

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def connect(self) -> None:
        """Start the shared producer and mark the client as ready."""
        if self._connected:
            return

        try:
            from aiokafka import AIOKafkaProducer
        except ImportError:
            raise ImportError(
                "aiokafka is required for KafkaClient. "
                "Install with: pip install broker_core[kafka]"
            )

        # Silence aiokafka's verbose debug logs — only show errors.
        logging.getLogger("aiokafka").setLevel(logging.ERROR)

        producer_kwargs = self._build_producer_kwargs()
        self._producer = AIOKafkaProducer(**producer_kwargs)
        await self._producer.start()
        self._connected = True
        logger.info("KafkaClient connected to %s", self.config.bootstrap_servers)

    async def disconnect(self) -> None:
        """Stop all subscriptions, then the producer."""
        # Stop subscriptions first
        for sub in list(self._active_subscriptions):
            try:
                await sub.stop()
            except Exception:
                logger.warning("Error stopping subscription", exc_info=True)
        self._active_subscriptions.clear()

        # Stop producer
        if self._producer is not None:
            await self._producer.stop()
            self._producer = None

        self._connected = False
        logger.info("KafkaClient disconnected")

    # ------------------------------------------------------------------
    # Builder entry point
    # ------------------------------------------------------------------

    def topic(self, name: str) -> "MessageBuilder":
        """Return a :class:`MessageBuilder` scoped to *name*."""
        from ..message.builder import MessageBuilder
        return MessageBuilder(self, name)

    # ------------------------------------------------------------------
    # Publish
    # ------------------------------------------------------------------

    async def _publish(self, message: "Message") -> MessageResult:
        """Publish a single message and return the broker acknowledgment."""
        if self._producer is None:
            raise RuntimeError("KafkaClient is not connected. Call connect() first.")

        # Serialise payload
        value = None
        if message.payload is not None:
            value = json.dumps(message.payload).encode("utf-8")

        # Serialise key
        key = None
        if message.key is not None:
            key = message.key.encode("utf-8")

        # Serialise headers
        headers = []
        if message.headers:
            headers = [
                (k, v.encode("utf-8"))
                for k, v in message.headers.items()
            ]

        # Publish and wait for ack
        result = await self._producer.send_and_wait(
            topic=message.topic,
            value=value,
            key=key,
            headers=headers,
            partition=message.partition,
        )

        msg_result = MessageResult(
            topic=result.topic,
            partition=result.partition,
            offset=result.offset,
            timestamp=result.timestamp,
        )

        logger.info(
            "Kafka published: topic=%s partition=%s offset=%s key=%s",
            msg_result.topic, msg_result.partition, msg_result.offset,
            message.key or "<none>",
        )

        return msg_result

    # ------------------------------------------------------------------
    # Subscribe
    # ------------------------------------------------------------------

    async def _subscribe(
        self, message: "Message", builder: "MessageBuilder"
    ) -> Subscription:
        """Create a dedicated consumer and return a :class:`Subscription`."""
        if not self._connected:
            raise RuntimeError("KafkaClient is not connected. Call connect() first.")

        try:
            from aiokafka import AIOKafkaConsumer
        except ImportError:
            raise ImportError(
                "aiokafka is required for KafkaClient. "
                "Install with: pip install broker_core[kafka]"
            )

        # Merge consumer settings: builder overrides > config defaults
        consumer_kwargs = self._build_consumer_kwargs(builder, message.topic)

        consumer = AIOKafkaConsumer(message.topic, **consumer_kwargs)
        await consumer.start()

        sub = Subscription(consumer=consumer, topic=message.topic)
        self._active_subscriptions.add(sub)

        # Clean up self-reference when subscription stops
        original_stop = sub.stop

        async def _tracked_stop():
            await original_stop()
            self._active_subscriptions.discard(sub)
            await consumer.stop()

        sub.stop = _tracked_stop  # type: ignore[method-assign]

        return sub

    # ------------------------------------------------------------------
    # Debug helpers
    # ------------------------------------------------------------------

    def raw_message(self, message: "Message") -> str:
        """Return the Kafka message as a formatted JSON string."""
        from ..message.builder import Message as M
        return message.debug()

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

    def _build_producer_kwargs(self) -> Dict[str, Any]:
        """Build keyword arguments for ``AIOKafkaProducer``."""
        kwargs: Dict[str, Any] = {
            "bootstrap_servers": self.config.bootstrap_servers,
            "acks": self.config.acks,
            "max_request_size": self.config.max_request_size,
        }
        if self.config.client_id:
            kwargs["client_id"] = self.config.client_id
        if self.config.compression_type:
            kwargs["compression_type"] = self.config.compression_type
        if self.config.security_protocol != "PLAINTEXT":
            kwargs["security_protocol"] = self.config.security_protocol
            if self.config.sasl_mechanism:
                kwargs["sasl_mechanism"] = self.config.sasl_mechanism
            if self.config.sasl_username:
                kwargs["sasl_plain_username"] = self.config.sasl_username
            if self.config.sasl_password:
                kwargs["sasl_plain_password"] = self.config.sasl_password
        return kwargs

    def _build_consumer_kwargs(
        self, builder: "MessageBuilder", topic: str
    ) -> Dict[str, Any]:
        """Build keyword arguments for ``AIOKafkaConsumer``.

        Settings from the builder chain take precedence over the
        :class:`KafkaConfig` defaults.
        """
        def _first(*values: Any) -> Any:
            """Return the first non-None value."""
            for v in values:
                if v is not None:
                    return v
            return None

        group_id = (
            _first(
                builder._consumer_group_id,
                self.config.group_id,
            )
            or f"broker-core-{uuid4().hex[:8]}"
        )

        kwargs: Dict[str, Any] = {
            "bootstrap_servers": self.config.bootstrap_servers,
            "group_id": group_id,
            "auto_offset_reset": _first(
                builder._auto_offset_reset, self.config.auto_offset_reset
            ),
            "enable_auto_commit": _first(
                builder._enable_auto_commit, self.config.enable_auto_commit
            ),
            "auto_commit_interval_ms": _first(
                builder._auto_commit_interval_ms, self.config.auto_commit_interval_ms
            ),
            "max_poll_records": _first(
                builder._max_poll_records, self.config.max_poll_records
            ),
            "session_timeout_ms": _first(
                builder._session_timeout_ms, self.config.session_timeout_ms
            ),
            "max_poll_interval_ms": self.config.max_poll_interval_ms,
            "value_deserializer": lambda v: json.loads(v.decode("utf-8")) if v else None,
            "key_deserializer": lambda k: k.decode("utf-8") if k else None,
        }

        client_id = _first(builder._consumer_client_id, self.config.client_id)
        if client_id:
            kwargs["client_id"] = client_id

        if self.config.security_protocol != "PLAINTEXT":
            kwargs["security_protocol"] = self.config.security_protocol
            if self.config.sasl_mechanism:
                kwargs["sasl_mechanism"] = self.config.sasl_mechanism
            if self.config.sasl_username:
                kwargs["sasl_plain_username"] = self.config.sasl_username
            if self.config.sasl_password:
                kwargs["sasl_plain_password"] = self.config.sasl_password

        return kwargs
