import asyncio
import logging
from collections import defaultdict
from collections.abc import Sequence
from ssl import SSLContext
from typing import Any, Optional, Union

import async_timeout

from aiokafka import __version__
from aiokafka.client import AIOKafkaClient
from aiokafka.errors import (
    IncompatibleBrokerVersion,
    LeaderNotAvailableError,
    NotControllerError,
    NotLeaderForPartitionError,
    for_code,
)
from aiokafka.protocol.admin import (
    AlterConfigsRequest,
    ApiVersionRequest_v0,
    CreatePartitionsRequest,
    CreateTopicsRequest,
    DeleteRecordsRequest,
    DeleteTopicsRequest,
    DescribeConfigsRequest,
    DescribeGroupsRequest,
    ListGroupsRequest,
)
from aiokafka.protocol.api import Request, Response
from aiokafka.protocol.commit import GroupCoordinatorRequest, OffsetFetchRequest
from aiokafka.protocol.metadata import MetadataRequest
from aiokafka.structs import OffsetAndMetadata, TopicPartition

from .config_resource import ConfigResource, ConfigResourceType
from .new_partitions import NewPartitions
from .new_topic import NewTopic
from .records_to_delete import RecordsToDelete

log = logging.getLogger(__name__)


class AIOKafkaAdminClient:
    """A class for administering the Kafka cluster.

    .. note::

        This class is considered **experimental**, so beware that it is subject
        to changes even in patch releases.

    Keyword Arguments:
        bootstrap_servers: 'host[:port]' string (or list of 'host[:port]'
            strings) that the consumer should contact to bootstrap initial
            cluster metadata. This does not have to be the full node list.
            It just needs to have at least one broker that will respond to a
            Metadata API Request. Default port is 9092. If no servers are
            specified, will default to localhost:9092.
        client_id (str): a name for this client. This string is passed in
            each request to servers and can be used to identify specific
            server-side log entries that correspond to this client. Also
            submitted to GroupCoordinator for logging with respect to
            consumer group administration. Default: 'aiokafka-{version}'
        request_timeout_ms (int): Client request timeout in milliseconds.
            Default: 40000.
        connections_max_idle_ms: Close idle connections after the number of
            milliseconds specified by this config. The broker closes idle
            connections after connections.max.idle.ms, so this avoids hitting
            unexpected socket disconnected errors on the client.
            Default: 540000
        retry_backoff_ms (int): Milliseconds to backoff when retrying on
            errors. Default: 100.
        metadata_max_age_ms (int): The period of time in milliseconds after
            which we force a refresh of metadata even if we haven't seen any
            partition leadership changes to proactively discover any new
            brokers or partitions. Default: 300000
        security_protocol (str): Protocol used to communicate with brokers.
            Valid values are: PLAINTEXT, SSL, SASL_PLAINTEXT, SASL_SSL.
            Default: PLAINTEXT.
        ssl_context (ssl.SSLContext): Pre-configured SSLContext for wrapping
            socket connections. If provided, all other ssl_* configurations
            will be ignored. Default: None.
        api_version (str): Specify which kafka API version to use.
            AIOKafka supports Kafka API versions >=0.9 only.
            If set to 'auto', will attempt to infer the broker version by
            probing various APIs. Default: auto
    """

    def __init__(
        self,
        *,
        loop=None,
        bootstrap_servers: Union[str, list[str]] = "localhost",
        client_id: str = "aiokafka-" + __version__,
        request_timeout_ms: int = 40000,
        connections_max_idle_ms: int = 540000,
        retry_backoff_ms: int = 100,
        metadata_max_age_ms: int = 300000,
        security_protocol: str = "PLAINTEXT",
        ssl_context: Optional[SSLContext] = None,
        api_version: str = "auto",
        sasl_mechanism: str = "PLAIN",
        sasl_plain_username: Optional[str] = None,
        sasl_plain_password: Optional[str] = None,
        sasl_kerberos_service_name: str = "kafka",
        sasl_kerberos_domain_name: Optional[str] = None,
        sasl_oauth_token_provider: Optional[str] = None,
    ):
        self._closed = False
        self._started = False
        self._version_info = {}
        self._request_timeout_ms = request_timeout_ms
        self._client = AIOKafkaClient(
            loop=loop,
            bootstrap_servers=bootstrap_servers,
            client_id=client_id,
            metadata_max_age_ms=metadata_max_age_ms,
            request_timeout_ms=request_timeout_ms,
            retry_backoff_ms=retry_backoff_ms,
            api_version=api_version,
            ssl_context=ssl_context,
            security_protocol=security_protocol,
            connections_max_idle_ms=connections_max_idle_ms,
            sasl_mechanism=sasl_mechanism,
            sasl_plain_username=sasl_plain_username,
            sasl_plain_password=sasl_plain_password,
            sasl_kerberos_service_name=sasl_kerberos_service_name,
            sasl_kerberos_domain_name=sasl_kerberos_domain_name,
            sasl_oauth_token_provider=sasl_oauth_token_provider,
        )

    async def close(self):
        """Close the AIOKafkaAdminClient connection to the Kafka broker."""
        if not hasattr(self, "_closed") or self._closed:
            log.info("AIOKafkaAdminClient already closed.")
            return

        await self._client.close()
        self._closed = True
        log.debug("AIOKafkaAdminClient is now closed.")

    async def _send_request(
        self,
        request: Request,
        node_id: Optional[int] = None,
    ) -> Response:
        if node_id is None:
            node_id = self._client.get_random_node()
        return await self._client.send(node_id, request)

    async def _get_version_info(self):
        resp = await self._send_request(ApiVersionRequest_v0())
        for api_key, min_version, max_version in resp.api_versions:
            self._version_info[api_key] = (min_version, max_version)

    async def start(self):
        if self._started:
            return
        await self._client.bootstrap()
        await self._get_version_info()
        log.debug("AIOKafkaAdminClient started")
        self._started = True

    def _matching_api_version(self, operation: Sequence[type[Request]]) -> int:
        """Find the latest version of the protocol operation
        supported by both this library and the broker.

        This resolves to the lesser of either the latest api
        version this library supports, or the max version
        supported by the broker.

        :param operation: A list of protocol operation versions from
        aiokafka.protocol.
        :return: The max matching version number between client and broker.
        """
        api_key = operation[0].API_KEY
        if not self._version_info or api_key not in self._version_info:
            raise IncompatibleBrokerVersion(
                f"Kafka broker does not support the '{operation[0].__name__}' "
                "Kafka protocol."
            )
        min_version, max_version = self._version_info[api_key]
        version = min(len(operation) - 1, max_version)
        if version < min_version:
            raise IncompatibleBrokerVersion(
                f"No version of the '{operation[0].__name__}' Kafka protocol "
                "is supported by both the client and broker."
            )
        return version

    async def _send_request_to_node(self, node_id: int, request: Request) -> Response:
        async with async_timeout.timeout(self._client._request_timeout_ms / 1000):
            while True:
                ready = await self._client.ready(node_id)
                if ready:
                    break
                await asyncio.sleep(self._client._retry_backoff)

        return await self._client.send(node_id, request)

    async def _send_to_controller(self, request: Request) -> Response:
        # With "auto" api_version the first request is sent with minimal
        # version, so the controller is not returned in metadata.
        if self._client.cluster.controller is None:
            await self._client.force_metadata_update()

        # 2 attempts in case cluster metadata is outdated
        try:
            return await self._send_request_to_node(
                self._client.cluster.controller.nodeId, request
            )
        except NotControllerError:
            await self._client.force_metadata_update()
            return await self._send_request_to_node(
                self._client.cluster.controller.nodeId, request
            )

    @staticmethod
    def _convert_new_topic_request(new_topic):
        return (
            new_topic.name,
            new_topic.num_partitions,
            new_topic.replication_factor,
            list(new_topic.replica_assignments.items()),
            list(new_topic.topic_configs.items()),
        )

    async def create_topics(
        self,
        new_topics: list[NewTopic],
        timeout_ms: Optional[int] = None,
        validate_only: bool = False,
    ) -> Response:
        """Create new topics in the cluster.

        :param new_topics: A list of NewTopic objects.
        :param timeout_ms: Milliseconds to wait for new topics to be created
            before the broker returns.
        :param validate_only: If True, don't actually create new topics.
            Not supported by all versions. Default: False
        :return: Appropriate version of CreateTopicResponse class.
        """
        version = self._matching_api_version(CreateTopicsRequest)
        topics = [self._convert_new_topic_request(nt) for nt in new_topics]
        log.debug("Attempting to send create topic request for %r", new_topics)
        timeout_ms = timeout_ms or self._request_timeout_ms
        if version == 0:
            if validate_only:
                raise IncompatibleBrokerVersion(
                    "validate_only requires CreateTopicsRequest >= v1, "
                    f"which is not supported by Kafka {self._client.api_version}."
                )
            request = CreateTopicsRequest[version](
                create_topic_requests=topics,
                timeout=timeout_ms,
            )
        elif version <= 3:
            request = CreateTopicsRequest[version](
                create_topic_requests=topics,
                timeout=timeout_ms,
                validate_only=validate_only,
            )
        else:
            raise NotImplementedError(
                f"Support for CreateTopics v{version} has not yet been added "
                "to AIOKafkaAdminClient."
            )
        return await self._send_to_controller(request)

    async def delete_topics(
        self,
        topics: list[str],
        timeout_ms: Optional[int] = None,
    ) -> Response:
        """Delete topics from the cluster.

        :param topics: A list of topic name strings.
        :param timeout_ms: Milliseconds to wait for topics to be deleted
            before the broker returns.
        :return: Appropriate version of DeleteTopicsResponse class.
        """
        version = self._matching_api_version(DeleteTopicsRequest)
        req_cls = DeleteTopicsRequest[version]
        request = req_cls(topics, timeout_ms or self._request_timeout_ms)
        return await self._send_to_controller(request)

    async def _get_cluster_metadata(
        self,
        topics: Optional[list[str]] = None,
    ) -> Response:
        """
        Retrieve cluster metadata
        :param topics List of topic names, None means "get all topics"
        :return MetadataResponse
        """
        req_cls = MetadataRequest[self._matching_api_version(MetadataRequest)]
        request = req_cls(topics=topics)
        return await self._send_request(request)

    async def list_topics(self) -> list[str]:
        metadata = await self._get_cluster_metadata(topics=None)
        obj = metadata.to_object()
        return [t["topic"] for t in obj["topics"]]

    async def describe_topics(
        self,
        topics: Optional[list[str]] = None,
    ) -> list[Any]:
        metadata = await self._get_cluster_metadata(topics=topics)
        obj = metadata.to_object()
        return obj["topics"]

    async def describe_cluster(self) -> dict[str, Any]:
        metadata = await self._get_cluster_metadata()
        obj = metadata.to_object()
        obj.pop("topics")  # We have 'describe_topics' for this
        return obj

    async def describe_configs(
        self,
        config_resources: list[ConfigResource],
        include_synonyms: bool = False,
    ) -> list[Response]:
        """Fetch configuration parameters for one or more Kafka resources.

        :param config_resources: An list of ConfigResource objects.
            Any keys in ConfigResource.configs dict will be used to filter the
            result. Setting the configs dict to None will get all values. An
            empty dict will get zero values (as per Kafka protocol).
        :param include_synonyms: If True, return synonyms in response. Not
            supported by all versions. Default: False.
        :return: List of appropriate version of DescribeConfigsResponse class.
        """

        futures = []
        version = self._matching_api_version(DescribeConfigsRequest)
        if version == 0 and include_synonyms:
            raise IncompatibleBrokerVersion(
                "include_synonyms requires DescribeConfigsRequest >= v1,"
                f" which is not supported by Kafka {self._client.api_version}."
            )
        broker_res, topic_res = self._convert_config_resources(
            config_resources,
            "describe",
        )
        req_cls = DescribeConfigsRequest[version]
        for broker_id in broker_res:
            if version == 0:
                req = req_cls(resources=broker_res[broker_id])
            else:
                req = req_cls(
                    resources=broker_res[broker_id],
                    include_synonyms=include_synonyms,
                )
            futures.append(self._send_request(req, broker_id))
        if topic_res:
            if version == 0:
                req = req_cls(topic_res)
            else:
                req = req_cls(topic_res, include_synonyms)
            futures.append(self._send_request(req))
        return await asyncio.gather(*futures)

    async def alter_configs(
        self, config_resources: list[ConfigResource]
    ) -> list[Response]:
        """Alter configuration parameters of one or more Kafka resources.
        :param config_resources: A list of ConfigResource objects.
        :return: Appropriate version of AlterConfigsResponse class.
        """
        futures = []
        version = self._matching_api_version(AlterConfigsRequest)
        broker_resources, topic_resources = self._convert_config_resources(
            config_resources,
            "alter",
        )
        req_cls = AlterConfigsRequest[version]
        futures.append(self._send_request(req_cls(resources=topic_resources)))
        for broker_id in broker_resources:
            req = req_cls(resources=broker_resources[broker_id])
            futures.append(self._send_request(req, broker_id))
        return await asyncio.gather(*futures)

    @staticmethod
    def _convert_describe_config_resource_request(config_resource):
        return (
            config_resource.resource_type,
            config_resource.name,
            list(config_resource.configs.keys()) if config_resource.configs else None,
        )

    @staticmethod
    def _convert_alter_config_resource_request(config_resource):
        return (
            config_resource.resource_type,
            config_resource.name,
            list(config_resource.configs.items()),
        )

    @classmethod
    def _convert_config_resources(
        cls,
        config_resources: list[ConfigResource],
        op_type: str = "describe",
    ) -> tuple[dict[int, Any], list[Any]]:
        broker_resources = defaultdict(list)
        topic_resources = []
        if op_type == "describe":
            convert_func = cls._convert_describe_config_resource_request
        else:
            convert_func = cls._convert_alter_config_resource_request
        for config_resource in config_resources:
            resource = convert_func(config_resource)
            if config_resource.resource_type == ConfigResourceType.BROKER:
                broker_resources[int(resource[1])].append(resource)
            else:
                topic_resources.append(resource)
        return broker_resources, topic_resources

    @staticmethod
    def _convert_topic_partitions(topic_partitions: dict[str, NewPartitions]):
        return [
            (topic_name, (new_part.total_count, new_part.new_assignments))
            for topic_name, new_part in topic_partitions.items()
        ]

    async def create_partitions(
        self,
        topic_partitions: dict[str, NewPartitions],
        timeout_ms: Optional[int] = None,
        validate_only: bool = False,
    ) -> Response:
        """Create additional partitions for an existing topic.

        :param topic_partitions: A map of topic name strings to NewPartition
         objects.
        :param timeout_ms: Milliseconds to wait for new partitions to be
            created before the broker returns.
        :param validate_only: If True, don't actually create new partitions.
            Default: False
        :return: Appropriate version of CreatePartitionsResponse class.
        """
        version = self._matching_api_version(CreatePartitionsRequest)
        req_class = CreatePartitionsRequest[version]
        converted_partitions = self._convert_topic_partitions(topic_partitions)
        req = req_class(
            topic_partitions=converted_partitions,
            timeout=timeout_ms or self._request_timeout_ms,
            validate_only=validate_only,
        )
        resp = await self._send_to_controller(req)
        for topic, code, message in resp.topic_errors:
            if code:
                err_cls = for_code(code)
                raise err_cls(f"Could not create partitions for {topic}: {message}")
        return resp

    async def describe_consumer_groups(
        self,
        group_ids: list[str],
        group_coordinator_id: Optional[int] = None,
        include_authorized_operations: bool = False,
    ) -> list[Response]:
        """Describe a set of consumer groups.

        Any errors are immediately raised.

        :param group_ids: A list of consumer group IDs. These are typically the
            group names as strings.
        :param group_coordinator_id: The node_id of the groups' coordinator
            broker. If set to None, it will query the cluster for each group to
            find that group's coordinator. Explicitly specifying this can be
            useful for avoiding extra network round trips if you already know
            the group coordinator. This is only useful when all the group_ids
            have the same coordinator, otherwise it will error. Default: None.
        :param include_authorized_operations: Whether or not to include
            information about the operations a group is allowed to perform.
            Only supported on API version >= v3. Default: False.
        :return: A list of group descriptions. For now the group descriptions
            are the raw results from the DescribeGroupsResponse.
        """
        version = self._matching_api_version(DescribeGroupsRequest)
        if version < 3 and include_authorized_operations:
            raise IncompatibleBrokerVersion(
                "include_authorized_operations requests "
                "DescribeGroupsRequest >= v3, which is not "
                f"supported by Kafka {version}"
            )
        req_class = DescribeGroupsRequest[version]
        futures = []
        node_to_groups = defaultdict(set)
        for group_id in group_ids:
            if group_coordinator_id is None:
                node_id = await self.find_coordinator(group_id)
            else:
                node_id = group_coordinator_id
            node_to_groups[node_id].add(group_id)
        for node_id, groups in node_to_groups.items():
            if include_authorized_operations:
                req = req_class(
                    groups=list(groups),
                    include_authorized_operations=include_authorized_operations,
                )
            else:
                req = req_class(groups=list(groups))
            future = self._send_request(req, node_id)
            futures.append(future)
        results = await asyncio.gather(*futures)
        return results

    async def list_consumer_groups(
        self,
        broker_ids: Optional[list[int]] = None,
    ) -> list[tuple[Any, ...]]:
        """List all consumer groups known to the cluster.

        This returns a list of Consumer Group tuples. The tuples are
        composed of the consumer group name and the consumer group protocol
        type.

        Only consumer groups that store their offsets in Kafka are returned.
        The protocol type will be an empty string for groups created using
        Kafka < 0.9 APIs because, although they store their offsets in Kafka,
        they don't use Kafka for group coordination. For groups created using
        Kafka >= 0.9, the protocol type will typically be "consumer".

        As soon as any error is encountered, it is immediately raised.

        :param broker_ids: A list of broker node_ids to query for consumer
            groups. If set to None, will query all brokers in the cluster.
            Explicitly specifying broker(s) can be useful for determining which
            consumer groups are coordinated by those broker(s). Default: None
        :return list: List of tuples of Consumer Groups.
        :exception GroupCoordinatorNotAvailableError: The coordinator is not
            available, so cannot process requests.
        :exception GroupLoadInProgressError: The coordinator is loading and
            hence can't process requests.
        """
        if broker_ids is None:
            metadata = await self._get_cluster_metadata()
            broker_ids = [broker[0] for broker in metadata.brokers]
        consumer_groups = set()
        for broker_id in broker_ids:
            response = await self._send_request(
                ListGroupsRequest[self._matching_api_version(ListGroupsRequest)](),
                broker_id,
            )
            if response.error_code:
                raise for_code(response.error_code)("Error listing consumer groups")
            consumer_groups.update(response.groups)
        return list(consumer_groups)

    async def find_coordinator(self, group_id: str, coordinator_type: int = 0) -> int:
        """Find the broker id for a given consumer group

        :param group_id: str the group id
        :param coordinator_type: int the type of coordinator:
        0 for group, 1 for transaction. Defaults to group.
        Only supported by version 1 and up

        :return int: the acting coordinator broker id
        """
        version = self._matching_api_version(GroupCoordinatorRequest)
        if version == 0 and coordinator_type:
            raise IncompatibleBrokerVersion(
                "Cannot query for transaction id on current broker version"
            )
        req_class = GroupCoordinatorRequest[version]
        if version == 0:
            request = req_class(consumer_group=group_id)
        else:
            request = req_class(group_id, coordinator_type)
        response = await self._send_request(request)
        if response.error_code:
            err = for_code(response.error_code)
            raise err(f"Unable to get coordinator id for {group_id}")
        return response.coordinator_id

    async def list_consumer_group_offsets(
        self,
        group_id: str,
        group_coordinator_id: Optional[int] = None,
        partitions: Optional[list[TopicPartition]] = None,
    ) -> dict[TopicPartition, OffsetAndMetadata]:
        """Fetch Consumer Offsets for a single consumer group.

        Note:
        This does not verify that the group_id or partitions actually exist
        in the cluster.

        As soon as any error is encountered, it is immediately raised.

        :param group_id: The consumer group id name for which to fetch offsets.
        :param group_coordinator_id: The node_id of the group's coordinator
            broker. If set to None, will query the cluster to find the group
            coordinator. Explicitly specifying this can be useful to prevent
            that extra network round trip if you already know the group
            coordinator. Default: None.
        :param partitions: A list of TopicPartitions for which to fetch
            offsets. On brokers >= 0.10.2, this can be set to None to fetch all
            known offsets for the consumer group. Default: None.
        :return dictionary: A dictionary with TopicPartition keys and
            OffsetAndMetada values. Partitions that are not specified and for
            which the group_id does not have a recorded offset are omitted. An
            offset value of `-1` indicates the group_id has no offset for that
            TopicPartition. A `-1` can only happen for partitions that are
            explicitly specified.
        """
        version = self._matching_api_version(OffsetFetchRequest)
        if version <= 1 and partitions is None:
            raise ValueError(
                f"""OffsetFetchRequest_v{version} requires specifying the
                partitions for which to fetch offsets. Omitting the
                partitions is only supported on brokers >= 0.10.2"""
            )
        if partitions:
            topics_partitions_dict = defaultdict(set)
            for topic, partition in partitions:
                topics_partitions_dict[topic].add(partition)
            partitions = [
                (topic, list(partitions))
                for topic, partitions in topics_partitions_dict.items()
            ]
        request = OffsetFetchRequest[version](group_id, partitions)
        if group_coordinator_id is None:
            group_coordinator_id = await self.find_coordinator(group_id)
        response = await self._send_request(request, group_coordinator_id)
        response_dict = {}
        for topic, partitions in response.topics:
            for partition, offset, metadata, error_code in partitions:
                if error_code:
                    err = for_code(response.error_code)
                    raise err(f"Unable to get offset info for {topic} and {partition}")
                tp = TopicPartition(topic, partition)
                offset_plus_meta = OffsetAndMetadata(offset, metadata)
                response_dict[tp] = offset_plus_meta
        return response_dict

    async def delete_records(
        self,
        records_to_delete: dict[TopicPartition, RecordsToDelete],
        timeout_ms: Optional[int] = None,
    ) -> dict[TopicPartition, int]:
        """Delete records from partitions.

        :param records_to_delete: A map of RecordsToDelete for each TopicPartition
        :param timeout_ms: Milliseconds to wait for the deletion to complete.
        :return: Appropriate version of DeleteRecordsResponse class.
        """
        version = self._matching_api_version(DeleteRecordsRequest)

        metadata = await self._get_cluster_metadata()

        self._client.cluster.update_metadata(metadata)

        requests = defaultdict(lambda: defaultdict(list))
        responses = {}

        for tp, records in records_to_delete.items():
            leader = self._client.cluster.leader_for_partition(tp)
            if leader is None:
                raise NotLeaderForPartitionError()
            elif leader == -1:
                raise LeaderNotAvailableError()
            requests[leader][tp.topic].append((tp.partition, records))

        req_cls = DeleteRecordsRequest[version]

        for leader, delete_request in requests.items():
            request = req_cls(
                self._convert_records_to_delete(delete_request),
                timeout_ms or self._request_timeout_ms,
            )
            response = await self._client.send(leader, request)
            # Starting with v2, DeleteRecordsResponse contains extra field with tags
            for topic, partitions, *_ in response.topics:
                for partition_index, low_watermark, error_code, *_ in partitions:
                    if error_code:
                        err = for_code(error_code)
                        raise err
                    responses[TopicPartition(topic, partition_index)] = low_watermark
        return responses

    @staticmethod
    def _convert_records_to_delete(
        records_to_delete: dict[str, list[tuple[int, RecordsToDelete]]],
    ):
        return [
            (topic, [(partition, rec.before_offset) for partition, rec in records])
            for topic, records in records_to_delete.items()
        ]
