# -*- coding: utf-8 -*-
# Copyright 2016 Yelp Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from collections import defaultdict
from collections import namedtuple

from kafka.common import UnknownTopicOrPartitionError
from kafka.structs import BrokerResponseError
from kafka.structs import check_error
from kafka.structs import OffsetCommitRequestPayload
from kafka.structs import OffsetFetchRequestPayload
from kafka.structs import OffsetFetchResponsePayload
from kafka.structs import OffsetRequestPayload
from kafka.structs import OffsetResponsePayload

from kafka_utils.util.error import InvalidOffsetStorageError
from kafka_utils.util.error import OffsetCommitError
from kafka_utils.util.error import UnknownPartitions
from kafka_utils.util.error import UnknownTopic


PartitionOffsets = namedtuple(
    'PartitionOffsets',
    ['topic', 'partition', 'highmark', 'lowmark']
)
"""Tuple representing the offsets for a topic partition.

* **topic**\(``str``): Name of the topic
* **partition**\(``int``): Partition number
* **highmark**\(``int``): high watermark
* **lowmark**\(``int``): low watermark
"""

HIGH_WATERMARK = "high"
LOW_WATERMARK = "low"


def pluck_topic_offset_or_zero_on_unknown(resp):
    try:
        check_error(resp)
    except UnknownTopicOrPartitionError:
        # If the server doesn't have any commited offsets by this group for
        # this topic, assume it's zero.
        pass
    # The API spec says server wont set an error, but 0.8.1.1 does. The actual
    # check is if the offset is -1.
    if resp.offset == -1:
        return OffsetFetchResponsePayload(
            resp.topic,
            resp.partition,
            0,
            resp.metadata,
            0,
        )
    return resp


def _check_fetch_response_error(resp):
    try:
        check_error(resp)
    except BrokerResponseError:
        # In case of error we set the offset to (-1,)
        return OffsetResponsePayload(
            resp.topic,
            resp.partition,
            resp.error,
            (-1,),
        )
    return resp


def _check_commit_response_error(resp):
    try:
        check_error(resp)
    except BrokerResponseError as e:
        exception = OffsetCommitError(
            resp.topic,
            resp.partition,
            e.message
        )
        return exception
    return resp


def _validate_topics_list_or_dict(topics):
    if isinstance(topics, dict):
        return topics
    elif isinstance(topics, (list, set, tuple)):
        return dict([(topic, []) for topic in topics])
    else:
        raise TypeError("Invalid topics: {topics}. It must be either a "
                        "list of topics or a dict "
                        "topic: [partitions]".format(topics=topics))


def _verify_topics_and_partitions(kafka_client, topics, raise_on_error):
    topics = _validate_topics_list_or_dict(topics)
    valid_topics = {}
    for topic, partitions in topics.iteritems():
        # Check topic exists
        if not kafka_client.has_metadata_for_topic(topic):
            if raise_on_error:
                raise UnknownTopic("Topic {topic!r} does not exist in "
                                   "kafka".format(topic=topic))
            else:
                continue
        if partitions:
            # Check the partitions really exist
            unknown_partitions = set(partitions) - \
                set(kafka_client.get_partition_ids_for_topic(topic))
            if unknown_partitions:
                if raise_on_error:
                    raise UnknownPartitions(
                        "Partitions {partitions!r} for topic {topic!r} do not"
                        "exist in kafka".format(
                            partitions=unknown_partitions,
                            topic=topic,
                        )
                    )
                else:
                    # We only use the available partitions in this case
                    partitions = set(partitions) - unknown_partitions
        else:
            # Default get all partitions metadata
            partitions = kafka_client.get_partition_ids_for_topic(topic)
        valid_topics[topic] = partitions
    return valid_topics


def _verify_commit_offsets_requests(kafka_client, new_offsets, raise_on_error):
    type_error_str = (
        "Invalid new_offsets: {new_offsets}. It must be a "
        "dict of the format: "
        "{{<topic>: {{<partition>: <offset>}}}}"
    ).format(new_offsets=new_offsets)

    if not isinstance(new_offsets, dict):
        raise TypeError(type_error_str)

    for topic, partitions in new_offsets.iteritems():
        if not isinstance(partitions, dict):
            raise TypeError(type_error_str)

    topics = dict(
        (topic, partitions.keys())
        for topic, partitions in new_offsets.iteritems()
    )

    valid_topics = _verify_topics_and_partitions(kafka_client, topics, raise_on_error)

    return dict(
        (topic, dict(
            (partition, new_offsets[topic][partition])
            for partition in partitions
        ))
        for topic, partitions in valid_topics.iteritems()
        if partitions
    )


def get_current_consumer_offsets(
    kafka_client,
    group,
    topics,
    raise_on_error=True,
    offset_storage='kafka',
):
    """ Get current consumer offsets.

    NOTE: This method does not refresh client metadata. It is up to the caller
    to avoid using stale metadata.

    If any partition leader is not available, the request fails for all the
    other topics. This is the tradeoff of sending all topic requests in batch
    and save both in performance and Kafka load.

    :param kafka_client: a connected KafkaToolClient
    :param group: kafka group_id
    :param topics: topic list or dict {<topic>: [partitions]}
    :param raise_on_error: if False the method ignores missing topics and
      missing partitions. It still may fail on the request send.
    :param offset_storage: String, one of {zookeeper, kafka}.
    :returns: a dict topic: partition: offset
    :raises:
      :py:class:`kafka_utils.util.error.UnknownTopic`: upon missing
      topics and raise_on_error=True

      :py:class:`kafka_utils.util.error.UnknownPartition`: upon missing
      partitions and raise_on_error=True

      :py:class:`kafka_utils.util.error.InvalidOffsetStorageError: upon unknown
      offset_storage choice.

      FailedPayloadsError: upon send request error.
    """

    topics = _verify_topics_and_partitions(kafka_client, topics, raise_on_error)

    group_offset_reqs = [
        OffsetFetchRequestPayload(topic, partition)
        for topic, partitions in topics.iteritems()
        for partition in partitions
    ]

    group_offsets = {}

    if offset_storage == 'zookeeper':
        send_api = kafka_client.send_offset_fetch_request
    elif offset_storage == 'kafka':
        send_api = kafka_client.send_offset_fetch_request_kafka
    else:
        raise InvalidOffsetStorageError(offset_storage)

    if group_offset_reqs:
        # fail_on_error = False does not prevent network errors
        group_resps = send_api(
            group=group,
            payloads=group_offset_reqs,
            fail_on_error=False,
            callback=pluck_topic_offset_or_zero_on_unknown,
        )
        for resp in group_resps:
            group_offsets.setdefault(
                resp.topic,
                {},
            )[resp.partition] = resp.offset

    return group_offsets


def get_topics_watermarks(kafka_client, topics, raise_on_error=True):
    """ Get current topic watermarks.

    NOTE: This method does not refresh client metadata. It is up to the caller
    to use avoid using stale metadata.

    If any partition leader is not available, the request fails for all the
    other topics. This is the tradeoff of sending all topic requests in batch
    and save both in performance and Kafka load.

    :param kafka_client: a connected KafkaToolClient
    :param topics: topic list or dict {<topic>: [partitions]}
    :param raise_on_error: if False the method ignores missing topics
      and missing partitions. It still may fail on the request send.
    :returns: a dict topic: partition: Part
    :raises:
      :py:class:`~kafka_utils.util.error.UnknownTopic`: upon missing
      topics and raise_on_error=True

      :py:class:`~kafka_utils.util.error.UnknownPartition`: upon missing
      partitions and raise_on_error=True

      FailedPayloadsError: upon send request error.
    """
    topics = _verify_topics_and_partitions(
        kafka_client,
        topics,
        raise_on_error,
    )
    highmark_offset_reqs = []
    lowmark_offset_reqs = []

    for topic, partitions in topics.iteritems():
        # Batch watermark requests
        for partition in partitions:
            # Request the the latest offset
            highmark_offset_reqs.append(
                OffsetRequestPayload(
                    topic, partition, -1, max_offsets=1
                )
            )
            # Request the earliest offset
            lowmark_offset_reqs.append(
                OffsetRequestPayload(
                    topic, partition, -2, max_offsets=1
                )
            )

    watermark_offsets = {}

    if not (len(highmark_offset_reqs) + len(lowmark_offset_reqs)):
        return watermark_offsets

    # fail_on_error = False does not prevent network errors
    highmark_resps = kafka_client.send_offset_request(
        highmark_offset_reqs,
        fail_on_error=False,
        callback=_check_fetch_response_error,
    )
    lowmark_resps = kafka_client.send_offset_request(
        lowmark_offset_reqs,
        fail_on_error=False,
        callback=_check_fetch_response_error,
    )

    # At this point highmark and lowmark should ideally have the same length.
    assert len(highmark_resps) == len(lowmark_resps)
    aggregated_offsets = defaultdict(lambda: defaultdict(dict))
    for resp in highmark_resps:
        aggregated_offsets[resp.topic][resp.partition]['highmark'] = \
            resp.offsets[0]
    for resp in lowmark_resps:
        aggregated_offsets[resp.topic][resp.partition]['lowmark'] = \
            resp.offsets[0]

    for topic, partition_watermarks in aggregated_offsets.iteritems():
        for partition, watermarks in partition_watermarks.iteritems():
            watermark_offsets.setdefault(
                topic,
                {},
            )[partition] = PartitionOffsets(
                topic,
                partition,
                watermarks['highmark'],
                watermarks['lowmark'],
            )
    return watermark_offsets


def _commit_offsets_to_watermark(
    kafka_client,
    group,
    topics,
    watermark,
    raise_on_error,
    offset_storage,
):
    topics = _verify_topics_and_partitions(kafka_client, topics, raise_on_error)

    watermark_offsets = get_topics_watermarks(kafka_client, topics, raise_on_error)

    if watermark == HIGH_WATERMARK:
        group_offset_reqs = [
            OffsetCommitRequestPayload(
                topic, partition,
                watermark_offsets[topic][partition].highmark,
                metadata=''
            )
            for topic, partitions in topics.iteritems()
            for partition in partitions
        ]
    elif watermark == LOW_WATERMARK:
        group_offset_reqs = [
            OffsetCommitRequestPayload(
                topic, partition,
                watermark_offsets[topic][partition].lowmark,
                metadata=''
            )
            for topic, partitions in topics.iteritems()
            for partition in partitions
        ]
    else:
        raise ValueError(
            "Unknown watermark: {watermark}".format(watermark=watermark)
        )

    if offset_storage == 'kafka' or not offset_storage:
        send_api = kafka_client.send_offset_commit_request_kafka
    elif offset_storage == 'zookeeper':
        send_api = kafka_client.send_offset_commit_request
    else:
        raise InvalidOffsetStorageError(offset_storage)

    status = []
    if group_offset_reqs:
        status = send_api(
            group,
            group_offset_reqs,
            raise_on_error,
            callback=_check_commit_response_error
        )

    return filter(None, status)


def advance_consumer_offsets(
    kafka_client,
    group,
    topics,
    raise_on_error=True,
    offset_storage='kafka',
):
    """Advance consumer offsets to the latest message in the topic
    partition (the high watermark).

    This method shall refresh the client metadata prior to updating
    the offsets.

    If any partition leader is not available, the request fails for all the
    other topics. This is the tradeoff of sending all topic requests in batch
    and save both in performance and Kafka load.

    :param kafka_client: a connected KafkaToolClient
    :param group: kafka group_id
    :param topics: topic list or dict {<topic>: [partitions]}
    :param raise_on_error: if False the method does not raise exceptions
      on missing topics/partitions. It may still fail on the request send.
    :param offset_storage: String, one of {zookeeper, kafka}.
    :returns: a list of errors for each partition offset update that failed.
    :rtype: list [OffsetCommitError]
    :raises:
      :py:class:`kafka_utils.util.error.UnknownTopic`: upon missing
      topics and raise_on_error=True

      :py:class:`kafka_utils.util.error.UnknownPartition`: upon missing
      partitions and raise_on_error=True

      FailedPayloadsError: upon send request error.
    """
    kafka_client.load_metadata_for_topics()

    return _commit_offsets_to_watermark(
        kafka_client, group, topics,
        HIGH_WATERMARK, raise_on_error,
        offset_storage,
    )


def rewind_consumer_offsets(
    kafka_client,
    group,
    topics,
    raise_on_error=True,
    offset_storage='kafka',
):
    """Rewind consumer offsets to the earliest message in the topic
    partition (the low watermark).

    This method shall refresh the client metadata prior to updating
    the offsets.

    If any partition leader is not available, the request fails for all the
    other topics. This is the tradeoff of sending all topic requests in batch
    and save both in performance and Kafka load.

    :param kafka_client: a connected KafkaToolClient
    :param group: kafka group_id
    :param topics: topic list or dict {<topic>: [partitions]}
    :param raise_on_error: if False the method does not raise exceptions
      on missing topics/partitions. It may still fail on the request send.
    :param offset_storage: String, one of {zookeeper, kafka}.
    :returns: a list of errors for each partition offset update that failed.
    :rtype: list [OffsetCommitError]
    :raises:
      :py:class:`kafka_utils.util.error.UnknownTopic`: upon missing
      topics and raise_on_error=True

      :py:class:`kafka_utils.util.error.UnknownPartition`: upon missing
      partitions and raise_on_error=True

      FailedPayloadsError: upon send request error.
    """
    kafka_client.load_metadata_for_topics()

    return _commit_offsets_to_watermark(
        kafka_client, group, topics,
        LOW_WATERMARK, raise_on_error,
        offset_storage,
    )


def set_consumer_offsets(
    kafka_client,
    group,
    new_offsets,
    raise_on_error=True,
    offset_storage='kafka',
):
    """Set consumer offsets to the specified offsets.

    This method does not validate the specified offsets, it is up to
    the caller to specify valid offsets within a topic partition.

    If any partition leader is not available, the request fails for all the
    other topics. This is the tradeoff of sending all topic requests in batch
    and save both in performance and Kafka load.

    :param kafka_client: a connected KafkaToolClient
    :param group: kafka group_id
    :param topics: dict {<topic>: {<partition>: <offset>}}
    :param raise_on_error: if False the method does not raise exceptions
      on errors encountered. It may still fail on the request send.
    :param offset_storage: String, one of {zookeeper, kafka}.
    :returns: a list of errors for each partition offset update that failed.
    :rtype: list [OffsetCommitError]
    :raises:
      :py:class:`kafka_utils.util.error.UnknownTopic`: upon missing
      topics and raise_on_error=True

      :py:class:`kafka_utils.util.error.UnknownPartition`: upon missing
      partitions and raise_on_error=True

      :py:class:`exceptions.TypeError`: upon badly formatted input
      new_offsets

      :py:class:`kafka_utils.util.error.InvalidOffsetStorageError: upon unknown
      offset_storage choice.

      FailedPayloadsError: upon send request error.
    """
    valid_new_offsets = _verify_commit_offsets_requests(
        kafka_client,
        new_offsets,
        raise_on_error
    )

    group_offset_reqs = [
        OffsetCommitRequestPayload(
            topic,
            partition,
            offset,
            metadata='',
        )
        for topic, new_partition_offsets in valid_new_offsets.iteritems()
        for partition, offset in new_partition_offsets.iteritems()
    ]

    if offset_storage == 'kafka' or not offset_storage:
        send_api = kafka_client.send_offset_commit_request_kafka
    elif offset_storage == 'zookeeper':
        send_api = kafka_client.send_offset_commit_request
    else:
        raise InvalidOffsetStorageError(offset_storage)

    status = []
    if group_offset_reqs:
        status = send_api(
            group,
            group_offset_reqs,
            raise_on_error,
            callback=_check_commit_response_error
        )

    return filter(None, status)


def _nullify_partition_offsets(partition_offsets):
    result = {}
    for partition in partition_offsets:
        result[partition] = 0
    return result


def nullify_offsets(offsets):
    """Modify offsets metadata so that the partition offsets
    have null payloads.

    :param offsets: dict {<topic>: {<partition>: <offset>}}
    :returns: a dict topic: partition: offset
    """
    result = {}
    for topic, partition_offsets in offsets.iteritems():
        result[topic] = _nullify_partition_offsets(partition_offsets)
    return result
