"""Improved Color Kafka"""

import argparse
import json
from kafka import KafkaConsumer
from pygments import highlight
from pygments.formatters import TerminalFormatter
from pygments.lexers import JsonLexer

# parse the args
parser = argparse.ArgumentParser(description='Consume from kafka')
parser.add_argument(
    '-s', '--server', help='Kafka broker ip or hostname', default="localhost")
parser.add_argument(
    '-g', '--group', help='kafka consumer group', default="defaultgroup")
parser.add_argument('-t', '--topic', help='Kafka topic name', required=True)
args = parser.parse_args()

# start consuming them bytes
consumer = KafkaConsumer(
    args.topic,
    bootstrap_servers=[args.server],
    enable_auto_commit=True,
    group_id=args.group,
)

# print each message that is consumed
for count, message in enumerate(consumer):
    try:
        message = message.value.decode('utf-8')
        message = json.loads(message)
        message = json.dumps(message, indent=4, sort_keys=True)
        print(highlight(message, JsonLexer(), TerminalFormatter()))
    except Exception:
        print(message)
    print("message count: {}".format(count+1))
