import logging
from .producer_utils import confirm_produce
from .consumer_utils import commit_message
from .confluent_runtime_vars import env_vars

LOGGER = logging.getLogger(__name__)


def success_headers(headers):
    """
    Updates headers for a message when a process succeeds

    Resets the kafka_retry_count to 0,
    or adds if it didn't originally exist
    """

    try:
        headers = dict(headers)
    except TypeError:
        headers = {}

    headers['kafka_retry_count'] = '0'
    return headers


def produce_message_callback(error, message):
    """
    Logs the returned message from the Broker after producing
    NOTE: Headers not supported on the message for callback for some reason.
    """
    LOGGER.debug('Producer Callback...')
    if error:
        LOGGER.critical(error)
    else:
        LOGGER.debug('Producer Callback - Message produced successfully')


def consume_message_callback(error, partitions):
    """
    Logs the info returned when a successful commit is performed
    """
    LOGGER.debug('Consumer Callback...')
    if error:
        LOGGER.critical(error)
    else:
        LOGGER.debug('Consumer Callback - Message consumption committed successfully')


def synchronous_message_handling(producer=None, consumer=None):
    if env_vars()['CONSUMER_ENABLE_AUTO_COMMIT'].lower() == 'true':
        LOGGER.warning(f'The environment variable "CONSUMER_ENABLE_AUTO_COMMIT" has been set to "true" when it is'
                       f' expected to be "false" using this method. Commiting may not behave as expected!')
    if producer:
        LOGGER.debug('Waiting for produce to finish...')
        confirm_produce(producer)
    if consumer:
        LOGGER.debug('Waiting for commits to finalize...')
        commit_message(consumer)


def shutdown_cleanup(producer=None, consumer=None):
    LOGGER.info('Performing graceful teardown of producer and/or consumer...')
    if consumer:
        LOGGER.debug("Shutting down consumer; no further commits can be queued or finalized.")
        consumer.close()
    if producer:
        LOGGER.debug("Sending/confirming the leftover messages in producer message queue")
        producer.flush(timeout=30)
