import logging
from collections.abc import Callable
from typing import Any

from aiokafka.admin import AIOKafkaAdminClient, NewTopic
from aiokafka.errors import TopicAlreadyExistsError
from faststream.kafka import KafkaBroker

from eggai.schemas import BaseMessage
from eggai.transport.base import Transport
from eggai.transport.middleware_utils import (
    create_data_type_middleware,
    create_filter_by_data_middleware,
    create_filter_middleware,
)


class KafkaTransport(Transport):
    """
    Kafka-based transport layer adapted to use FastStream's KafkaBroker for message publishing and consumption.

    This class serves as a transport mechanism that integrates with Kafka to allow message publishing
    and consumption. It uses the FastStream KafkaBroker to interact with Kafka, offering methods to
    connect, disconnect, publish messages to Kafka topics, and subscribe to Kafka topics.

    Attributes:
        broker (KafkaBroker): The KafkaBroker instance responsible for managing Kafka connections and messaging.
    """

    def __init__(
        self,
        broker: KafkaBroker | None = None,
        bootstrap_servers: str = "localhost:19092",
        **kwargs,
    ):
        """
        Initializes the KafkaTransport with an optional KafkaBroker or creates a new one with provided bootstrap servers.

        Args:
            broker (Optional[KafkaBroker]): An existing KafkaBroker instance to use. If not provided, a new instance will
                                             be created with the specified bootstrap servers and additional parameters.
            bootstrap_servers (str): The Kafka bootstrap server addresses (default is "localhost:19092").
            **kwargs: Additional keyword arguments to pass to the KafkaBroker if a new instance is created.

        Attributes:
            bootstrap_servers (Union[str, Iterable[str]]): A list or string of `host[:port]` addresses of brokers to contact for
                                                          bootstrap. Default is `"localhost"`.\n\n
            request_timeout_ms (int): Client request timeout in milliseconds (default is 40,000 ms).
            retry_backoff_ms (int): Milliseconds to back off when retrying on errors (default is 100 ms).
            metadata_max_age_ms (int): Period after which to refresh metadata (default is 300,000 ms).
            connections_max_idle_ms (int): Close idle connections after a specified time (default is 540,000 ms).
            sasl_kerberos_service_name (str): Kerberos service name (default is `"kafka"`).
            sasl_kerberos_domain_name (Optional[str]): Kerberos domain name.
            sasl_oauth_token_provider (Optional[AbstractTokenProvider]): OAuthBearer token provider instance.
            loop (Optional[AbstractEventLoop]): Optional event loop.
            client_id (Optional[str]): A name for this client (default is `"SERVICE_NAME"`).
            acks (Union[Literal[0, 1, -1, "all"], object]): Number of acknowledgments the producer requires before considering
                                                       a request complete (default is `None`).
            key_serializer (Optional[Callable[[Any], bytes]]): Function to serialize keys (default is `None`).
            value_serializer (Optional[Callable[[Any], bytes]]): Function to serialize values (default is `None`).
            compression_type (Optional[Literal["gzip", "snappy", "lz4", "zstd"]]): Compression type (default is `None`).
            max_batch_size (int): Maximum size of buffered data per partition (default is 16 KB).
            partitioner (Callable[[bytes, List[Partition], List[Partition]], Partition]): Partitioner function for assigning messages to partitions.
            max_request_size (int): Maximum size of a request (default is 1 MB).
            linger_ms (int): Time to delay requests for batching (default is 0 ms).
            enable_idempotence (bool): Whether to enable idempotence for the producer (default is `False`).
            transactional_id (Optional[str]): Transactional ID for producing messages (default is `None`).
            transaction_timeout_ms (int): Timeout for transactions (default is 60,000 ms).
            graceful_timeout (Optional[float]): Graceful shutdown timeout (default is 15.0).
            decoder (Optional[CustomCallable]): Custom decoder for messages (default is `None`).
            parser (Optional[CustomCallable]): Custom parser for messages (default is `None`).
            dependencies (Iterable[Depends]): Dependencies to apply to all broker subscribers (default is `()`).
            middlewares (Sequence[Union["BrokerMiddleware[ConsumerRecord]", "BrokerMiddleware[Tuple[ConsumerRecord, ...]]"]]):
                         Middlewares to apply to all broker publishers/subscribers (default is `()`).
            security (Optional[BaseSecurity]): Security options for broker connection (default is `None`).
            asyncapi_url (Union[str, Iterable[str], None]): AsyncAPI server URL (default is `None`).
            protocol (Optional[str]): AsyncAPI server protocol (default is `None`).
            protocol_version (Optional[str]): AsyncAPI server protocol version (default is `"auto"`).
            description (Optional[str]): AsyncAPI server description (default is `None`).
            tags (Optional[Iterable[Union["asyncapi.Tag", "asyncapi.TagDict"]]]): AsyncAPI server tags (default is `None`).
            logger (Optional[LoggerProto]): Custom logger to pass into context (default is `EMPTY`).
            log_level (int): Log level for service messages (default is `logging.INFO`).
            log_fmt (Optional[str]): Log format (default is `None`).
        """
        self._bootstrap_servers = bootstrap_servers
        self._topics_to_create: set[str] = set()
        self._num_partitions = kwargs.pop("num_partitions", 3)
        self._running = False

        if broker:
            self.broker = broker
        else:
            # metadata_max_age_ms: refresh metadata more frequently for dynamic topics
            if "metadata_max_age_ms" not in kwargs:
                kwargs["metadata_max_age_ms"] = 10000
            self.broker = KafkaBroker(
                bootstrap_servers, log_level=logging.INFO, **kwargs
            )

    async def connect(self):
        """
        Establishes a connection to the Kafka broker by starting the KafkaBroker instance.

        This method is necessary before publishing or consuming messages. It asynchronously starts the broker
        to handle Kafka communication.

        """
        if self._running:
            return
        await self.broker.start()
        self._running = True

    async def disconnect(self):
        """
        Closes the connection to the Kafka broker by stopping the KafkaBroker instance.

        This method should be called when the transport is no longer needed to stop consuming messages
        and to release any resources held by the KafkaBroker.
        """
        self._running = False
        await self.broker.stop()

    async def _ensure_topic(self, topic: str):
        """Ensure a single topic exists and refresh producer metadata."""
        admin = AIOKafkaAdminClient(bootstrap_servers=self._bootstrap_servers)
        try:
            await admin.start()
            new_topic = NewTopic(
                name=topic,
                num_partitions=self._num_partitions,
                replication_factor=1,
            )
            await admin.create_topics([new_topic])
        except TopicAlreadyExistsError:
            pass
        except Exception:
            pass
        finally:
            await admin.close()

        # Force producer metadata refresh for the new topic
        if (
            self._running
            and hasattr(self.broker, "_producer")
            and self.broker._producer
        ):
            try:
                await self.broker._producer._producer.client.force_metadata_update()
            except Exception:
                pass

    async def ensure_topic(self, channel: str):
        """Public method to ensure a topic exists."""
        if channel not in self._topics_to_create:
            await self._ensure_topic(channel)
            self._topics_to_create.add(channel)

    async def publish(self, channel: str, message: dict[str, Any] | BaseMessage):
        """
        Publishes a message to the specified Kafka topic (channel).

        Args:
            channel (str): The name of the Kafka topic to which the message will be published.
            message (Union[Dict[str, Any], BaseMessage]): The message to publish, which can either be a dictionary
                                                         or a BaseMessage instance. The message will be serialized
                                                         before being sent.

        """
        if channel not in self._topics_to_create:
            await self._ensure_topic(channel)
            self._topics_to_create.add(channel)
        await self.broker.publish(message, topic=channel)

    async def subscribe(self, channel: str, handler, **kwargs) -> Callable:
        """
        Subscribes to a Kafka topic (channel) and sets up a handler to process incoming messages.

        Args:
            channel (str): The Kafka topic to subscribe to.
            handler (Callable): The function or coroutine that will handle messages received from the topic.
            **kwargs: Additional keyword arguments that can be used to configure the subscription.

        Keyword Args:
            filter_by_message (Callable, optional): A function to filter incoming messages based on their payload.
            group_id (Optional[str], optional): The consumer group name for dynamic partition assignment.
            auto_offset_reset (str, optional): Policy for resetting offsets ('earliest' or 'latest').

        Returns:
            Callable: A callback function that represents the subscription.
        """
        if "filter_by_message" in kwargs:
            if "middlewares" not in kwargs:
                kwargs["middlewares"] = []
            kwargs["middlewares"].append(
                create_filter_middleware(kwargs.pop("filter_by_message"))
            )

        if "data_type" in kwargs:
            data_type = kwargs.pop("data_type")

            if "middlewares" not in kwargs:
                kwargs["middlewares"] = []
            kwargs["middlewares"].append(create_data_type_middleware(data_type))

            if "filter_by_data" in kwargs:
                if "middlewares" not in kwargs:
                    kwargs["middlewares"] = []
                kwargs["middlewares"].append(
                    create_filter_by_data_middleware(
                        data_type, kwargs.pop("filter_by_data")
                    )
                )

        # Use handler_id as default group_id (preserves broadcast behavior)
        handler_id = kwargs.pop("handler_id", None)
        if "group_id" not in kwargs:
            kwargs["group_id"] = handler_id

        self._topics_to_create.add(channel)

        if "pattern" in kwargs:
            return self.broker.subscriber(**kwargs)(handler)
        else:
            return self.broker.subscriber(channel, **kwargs)(handler)
