# Copyright 2016 IBM Corp. All Rights Reserved.
#
# 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
#
# https://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 json
import logging
import os
import requests
import time

# HEADS UP! I'm importing confluent_kafka.Consumer as KafkaConsumer to avoid a
# naming conflict with my own Consumer class
from confluent_kafka import Consumer as KafkaConsumer, KafkaError
from database import Database
from datetime import datetime
from threading import Thread, Lock

local_dev = os.getenv('LOCAL_DEV', 'False')
check_ssl = (local_dev == 'False')

class Consumer:
    class State:
        Initializing = 'Initializing'
        Running = 'Running'
        Stopping = 'Stopping'
        Restart = 'Restart'
        Dead = 'Dead'

    def __init__(self, trigger, params):
        self.trigger = trigger
        self.params = params
        self.thread = ConsumerThread(trigger, params)

        # the following fields can be accessed from multiple threads
        # access needs to be protected with this Lock
        self.__lock = Lock()
        self.__restartCount = 0

        # this is weird.
        # The app needs this tho...
        self.triggerURL = params["triggerURL"]

    def currentState(self):
        return self.thread.currentState()

    def desiredState(self):
        return self.thread.desiredState()

    def shutdown(self):
        self.thread.shutdown()

    def start(self):
        self.thread.start()

    # should only be called by the Doctor thread
    def restart(self):
        if self.thread.desiredState() is Consumer.State.Dead:
            logging.info('[{}] Request to restart a consumer that is already slated for deletion.'.format(self.trigger))
            return

        with self.__lock:
            self.__restartCount += 1

        logging.info('[{}] Quietly shutting down consumer for restart'.format(self.trigger))
        self.thread.setDesiredState(Consumer.State.Restart)
        self.thread.join()
        logging.info('Consumer has shut down')

        # user may have interleaved a request to delete the trigger, check again
        if self.thread.desiredState() is not Consumer.State.Dead:
            logging.info('[{}] Starting new consumer thread'.format(self.trigger))
            self.thread = ConsumerThread(self.trigger, self.params)
            self.thread.start()

    def restartCount(self):
        with self.__lock:
            restartCount = self.__restartCount

        return restartCount

    def lastPoll(self):
        return self.thread.lastPoll()

    def secondsSinceLastPoll(self):
        return self.thread.secondsSinceLastPoll()


class ConsumerThread (Thread):

    retry_timeout = 1   # Timeout in seconds
    max_retries = 10    # Maximum number of times to retry firing trigger

    database = Database()

    def __init__(self, trigger, params):
        Thread.__init__(self)

        self.lock = Lock()

        # the following params may be set/read from other threads
        # only access through helper methods which handle thread safety!
        self.__currentState = Consumer.State.Initializing
        self.__desiredState = Consumer.State.Running
        self.__lastPoll = datetime.max

        self.daemon = True
        self.trigger = trigger
        self.isMessageHub = params["isMessageHub"]
        self.triggerURL = params["triggerURL"]
        self.brokers = params["brokers"]
        self.topic = params["topic"]

        if self.isMessageHub:
            self.username = params["username"]
            self.password = params["password"]

        # handle the case where there may be existing triggers that do not
        # have the isJSONData field set
        if "isJSONData" in params:
            self.parseAsJson = params["isJSONData"]
        else:
            self.parseAsJson = False

        # always init consumer to None in case the consumer needs to shut down
        # before the KafkaConsumer is fully initialized/assigned
        self.consumer = None

    # this only records the current state, and does not affect a state transition
    def __recordState(self, newState):
        with self.lock:
            self.__currentState = newState

    def currentState(self):
        with self.lock:
            state = self.__currentState

        return state

    def setDesiredState(self, newState):
        logging.info('[{}] Request to set desiredState to {}'.format(self.trigger, newState))

        with self.lock:
            if self.__desiredState is Consumer.State.Dead and newState is not Consumer.State.Dead:
                logging.info('[{}] Asking to kill a consumer that is already marked for death. Doing nothing.'.format(self.trigger))
                return
            else:
                logging.info('[{}] Setting desiredState to: {}'.format(self.trigger, newState))
                self.__desiredState = newState

    def desiredState(self):
        with self.lock:
            state = self.__desiredState

        return state

    # convenience method for checking if desiredState is Running
    def __shouldRun(self):
        return self.desiredState() is Consumer.State.Running

    def lastPoll(self):
        with self.lock:
            lastPoll = self.__lastPoll

        return lastPoll

    def updateLastPoll(self):
        with self.lock:
            self.__lastPoll = datetime.now()

    def secondsSinceLastPoll(self):
        lastPollDelta = datetime.now() - self.lastPoll()
        return lastPollDelta.total_seconds()

    def run(self):
        try:
            self.consumer = self.__createConsumer()

            while self.__shouldRun():
                messages = self.__pollForMessages()

                if len(messages) > 0:
                    self.__fireTrigger(messages)

            logging.info("[{}] Consumer exiting main loop".format(self.trigger))

            if self.desiredState() == Consumer.State.Dead:
                logging.info('[{}] Permanently killing consumer because desired state is Dead'.format(self.trigger))
                self.database.deleteTrigger(self.trigger)
            elif self.desiredState() == Consumer.State.Restart:
                logging.info('[{}] Quietly letting the consumer thread stop in order to allow restart.'.format(self.trigger))
                # nothing else to do because this Thread is about to go away
            else:
                # uh-oh... this really shouldn't happen
                logging.error('[{}] Consumer stopped without being asked'.format(self.trigger))
        except Exception as e:
            logging.error('[{}] Uncaught exception: {}'.format(self.trigger, e))

        try:
            if self.consumer is not None:
                logging.info('[{}] Cleaning up consumer'.format(self.trigger))
                logging.debug('[{}] Closing KafkaConsumer'.format(self.trigger))
                self.consumer.unsubscribe()
                self.consumer.close()
                logging.debug('[{}] Successfully closed KafkaConsumer'.format(self.trigger))

                logging.debug('[{}] Dellocating KafkaConsumer'.format(self.trigger))
                self.consumer = None
        except Exception as e:
            logging.error('[{}] Uncaught exception while shutting down consumer: {}'.format(self.trigger, e))
        finally:
            logging.info('[{}] Recording consumer as Dead. Bye bye!'.format(self.trigger))
            self.__recordState(Consumer.State.Dead)

    def __createConsumer(self):
        if self.__shouldRun():
            config = {'metadata.broker.list': ','.join(self.brokers),
                        'group.id': self.trigger,
                        'default.topic.config': {'auto.offset.reset': 'latest'},
                        'enable.auto.commit': False
                    }

            if self.isMessageHub:
                # append Message Hub specific config
                config.update({'ssl.ca.location': '/etc/ssl/certs/',
                                'sasl.mechanisms': 'PLAIN',
                                'sasl.username': self.username,
                                'sasl.password': self.password,
                                'security.protocol': 'sasl_ssl'
                             })

            consumer = KafkaConsumer(config)
            consumer.subscribe([self.topic])
            logging.info("[{}] Now listening in order to fire trigger".format(self.trigger))
            return consumer

    def __pollForMessages(self):
        messages = []
        messageSize = 0
        batchMessages = True

        if self.__shouldRun():
            while batchMessages and (self.secondsSinceLastPoll() < 2):
                message = self.consumer.poll(1.0)

                if self.secondsSinceLastPoll() < 0:
                    self.__recordState(Consumer.State.Running)

                if (message is not None):
                    if not message.error():
                        logging.debug("Consumed message: {}".format(str(message)))
                        messageSize += len(message.value())
                        messages.append(message)
                    elif message.error().code() != KafkaError._PARTITION_EOF:
                        logging.error('[{}] Error polling: {}'.format(self.trigger, message.error()))
                        batchMessages = False
                    else:
                        logging.debug('[{}] No more messages. Stopping batch op.'.format(self.trigger))
                        batchMessages = False
                else:
                    logging.debug('[{}] message was None. Stopping batch op.'.format(self.trigger))
                    batchMessages = False

        logging.debug('[{}] Completed poll'.format(self.trigger))

        if len(messages) > 0:
            logging.info("[{}] Found {} messages with a total size of {} bytes".format(self.trigger, len(messages), messageSize))

        self.updateLastPoll()
        return messages

    def __fireTrigger(self, messages):
        if self.__shouldRun():
            lastMessage = messages[len(messages) - 1]

            # I'm sure there is a much more clever way to do this ;)
            mappedMessages = []
            for message in messages:
                fieldsToSend = {
                    'value': self.__parseMessageIfNeeded(message.value()),
                    'topic': message.topic(),
                    'partition': message.partition(),
                    'offset': message.offset(),
                    'key': message.key()
                }
                mappedMessages.append(fieldsToSend)

            payload = {}
            payload['messages'] = mappedMessages
            retry = True
            retry_count = 0

            logging.info("[{}] Firing trigger with {} messages".format(self.trigger,len(messages)))

            while retry:
                try:
                    response = requests.post(self.triggerURL, json=payload, timeout=10.0, verify=check_ssl)
                    status_code = response.status_code
                    logging.info("[{}] Repsonse status code {}".format(self.trigger, status_code))

                    # Manually commit offset if the trigger was fired successfully. Retry firing the trigger if the
                    # service was unavailable (503), an internal server error occurred (500), the request timed out
                    # (408), or the request was throttled (429).
                    if status_code == 200:
                        self.consumer.commit()
                        retry = False
                    elif status_code not in [503, 500, 408, 429]:
                        logging.error('[{}] Error talking to OpenWhisk, status code {}'.format(self.trigger, status_code))

                        # abandon all hope?
                        self.setDesiredState(Consumer.State.Dead)
                        retry = False
                except requests.exceptions.RequestException as e:
                    logging.error('[{}] Error talking to OpenWhisk: {}'.format(self.trigger, e))

                if retry:
                    retry_count += 1

                    if retry_count < self.max_retries:
                        logging.info("[{}] Retrying in {} second(s)".format(
                            self.trigger, self.retry_timeout))
                        time.sleep(self.retry_timeout)
                    else:
                        logging.warn("[{}] Skipping {} messages to offset {} of partition {}".format(self.trigger, len(messages), lastMessage.offset, lastMessage.partition))
                        self.consumer.commit()
                        retry = False

    def shutdown(self):
        logging.info("[{}] Shutting down consumer for trigger".format(self.trigger))
        self.setDesiredState(Consumer.State.Dead)
        self.database.deleteTrigger(self.trigger)

    def __parseMessageIfNeeded(self, value):
        if self.parseAsJson:
            try:
                parsed = json.loads(value)
                logging.debug(
                    '[{}] Successfully parsed a message as JSON.'.format(self.trigger))
                return parsed
            except ValueError:
                # no big deal, just return the original value
                logging.warn(
                    '[{}] I was asked to parse a message as JSON, but I failed.'.format(self.trigger))
                pass

        logging.debug('[{}] Returning un-parsed message'.format(self.trigger))
        return value