# (C) Copyright 2014-2016 Hewlett Packard Enterprise Development Company LP
#
# 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.

import datetime
import logging
import threading
import time

import kafka.client
import kafka.common
import kafka.consumer

from kazoo.client import KazooClient
from kazoo.recipe.partitioner import SetPartitioner

log = logging.getLogger(__name__)

"""Kafka consumer interface

Kafka consumer class that will automatically share partitions between processes
using the same zookeeper path.

For performance it is often required that data from the kafka queue be
batched before being processed.  There are two important concerns to
keep in mind when dealing with batched data.

    1. Negotiating partitions takes a rather long amount of time so when
    the negotiation process begins a defined repartition_callback will be
    called.  This is a good time to process whatever has been batched.
    2. If the traffic across the kafka topic is low enough it will take a long
    time to build a batch of data.  A commit_callback is available that will
    fire when the commit_timeout duration has elapsed since the last commit.
"""


class KafkaConsumer(object):
    def __init__(self, kafka_url,
                 zookeeper_url, zookeeper_path,
                 group, topic,
                 fetch_size=1048576,
                 repartition_callback=None,
                 commit_callback=None,
                 commit_timeout=30):
        """Init

             kafka_url            - Kafka location
             zookeeper_url        - Zookeeper location
             zookeeper_path       - Zookeeper path used for partition
                                    negotiation
             group                - Kafka consumer group
             topic                - Kafka topic
             repartition_callback - Callback to run when the Kafka consumer
                                    group changes.  Repartitioning takes a
                                    relatively long time so this is a good
                                    time to flush and commit any data.
             commit_callback      - Callback to run when the commit_timeout
                                    has elapsed between commits.
             commit_timeout       - Timeout between commits.
        """

        self._kazoo_client = None
        self._set_partitioner = None

        self._repartition_callback = repartition_callback

        self._commit_callback = commit_callback
        self._commit_timeout = commit_timeout

        self._last_commit = 0

        self._partitions = []

        self._kafka_group = group
        self._kafka_topic = topic
        self._kafka_fetch_size = fetch_size

        self._zookeeper_url = zookeeper_url
        self._zookeeper_path = zookeeper_path

        self._kafka = kafka.client.KafkaClient(kafka_url)

        self._consumer = self._create_kafka_consumer()

    def _create_kafka_consumer(self, partitions=None):
        # No auto-commit so that commits only happen after the message is processed.
        consumer = kafka.consumer.SimpleConsumer(
            self._kafka,
            self._kafka_group,
            self._kafka_topic,
            auto_commit=False,
            partitions=partitions,
            iter_timeout=5,
            fetch_size_bytes=self._kafka_fetch_size,
            buffer_size=self._kafka_fetch_size,
            max_buffer_size=None)

        consumer.provide_partition_info()
        consumer.fetch_last_known_offsets()
        return consumer

    def __iter__(self):
        self._partition()

        self._last_commit = datetime.datetime.now()

        while 1:
            if self._repartition():
                if self._repartition_callback:
                    self._repartition_callback()
                self._partition()

            # When Kafka resizes the partitions it's possible that it
            # will remove data at our current offset.  When this
            # happens the next attempt to read from Kafka will generate
            # an OffsetOutOfRangeError.  We trap this error and seek to
            # the head of the current Kafka data.  Because this error
            # only happens when Kafka removes data we're currently
            # pointing at we're gauranteed that we won't read any
            # duplicate data however we will lose any information
            # between our current offset and the new Kafka head.

            try:
                message = self._consumer.get_message()
                if message:
                    yield message
                else:
                    time.sleep(0.01)

                if self._commit_callback:
                    time_now = datetime.datetime.now()
                    time_delta = time_now - self._last_commit
                    if time_delta.total_seconds() > self._commit_timeout:
                        self._commit_callback()

            except kafka.common.OffsetOutOfRangeError:
                log.error("Kafka OffsetOutOfRange.  Jumping to head.")
                self._consumer.seek(0, 0)

    def _repartition(self):
        return not self._set_partitioner.acquired

    def _partition(self):
        """Consume messages from kafka

           Consume messages from kafka using the Kazoo SetPartitioner to
           allow multiple consumer processes to negotiate access to the kafka
           partitions
        """

        # KazooClient and SetPartitioner objects need to be instantiated after
        # the consumer process has forked.  Instantiating prior to forking
        # gives the appearance that things are working but after forking the
        # connection to zookeeper is lost and no state changes are visible

        if not self._kazoo_client:
            self._kazoo_client = KazooClient(hosts=self._zookeeper_url)
            self._kazoo_client.start()

            state_change_event = threading.Event()

            self._set_partitioner = (
                SetPartitioner(self._kazoo_client,
                               path=self._zookeeper_path,
                               set=self._consumer.fetch_offsets.keys(),
                               state_change_event=state_change_event,
                               identifier=str(datetime.datetime.now())))

        try:
            while 1:
                if self._set_partitioner.failed:
                    raise Exception("Failed to acquire partition")

                elif self._set_partitioner.release:
                    log.info("Releasing locks on partition set {} "
                             "for topic {}".format(self._partitions,
                                                   self._kafka_topic))
                    self._set_partitioner.release_set()

                    self._partitions = []

                elif self._set_partitioner.acquired:
                    if not self._partitions:
                        self._partitions = [p for p in self._set_partitioner]

                        if not self._partitions:
                            log.info("Not assigned any partitions on topic {},"
                                     " waiting for a Partitioner state change"
                                     .format(self._kafka_topic))
                            state_change_event.wait()
                            state_change_event.clear()
                            continue

                        log.info("Acquired locks on partition set {} "
                                 "for topic {}".format(self._partitions, self._kafka_topic))

                        # Reconstruct the kafka consumer object because the
                        # consumer has no API that allows the set of partitons
                        # to be updated outside of construction.
                        self._consumer.stop()
                        self._consumer = self._create_kafka_consumer(self._partitions)
                        return

                elif self._set_partitioner.allocating:
                    log.info("Waiting to acquire locks on partition set")
                    self._set_partitioner.wait_for_acquire()

        except Exception:
            log.exception('KafkaConsumer encountered fatal exception '
                          'processing messages.')
            raise

    def commit(self):
        self._last_commit = datetime.datetime.now()
        self._consumer.commit(partitions=self._partitions)
