#!/usr/bin/env python

import logging
from pythonjsonlogger import jsonlogger

import datetime

class CustomJsonFormatter(jsonlogger.JsonFormatter):
    def add_fields(self, log_record, record, message_dict):
        super(CustomJsonFormatter, self).add_fields(log_record, record, message_dict)
        if not log_record.get('@timestamp'):
            # this doesn't use record.created, so it is slightly off
            now = datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%S.%fZ')
            log_record['@timestamp'] = now
        if log_record.get('level'):
            log_record['level'] = log_record['level'].upper()
        else:
            log_record['level'] = record.levelname

def setup_logging(log_level='INFO'):
    logger = logging.getLogger(__name__)
    logger.propagate = 0
    logger.setLevel(log_level)
    logHandler = logging.StreamHandler()

    formatter = CustomJsonFormatter('%(@timestamp)s %(level)s %(name)s %(message)s')

    logHandler.setFormatter(formatter)
    logger.addHandler(logHandler)

def log(message, level):
    logger = logging.getLogger(__name__)
    if level == 'INFO':
        logger.info(message)
    elif level == 'WARN':
        logger.warn(message)
    elif level == 'ERR':
        logger.error(message)
    elif level == 'DEBUG':
        logger.debug(message)