import json
from uuid import uuid4
from aiokafka import AIOKafkaConsumer
from aiokafka.helpers import create_ssl_context
import asyncio

GROUP_ID = 'streams_plugin_3to4_test'
SECURITY_PROTOCOL='SSL'
AUTO_OFFSET_RESET='latest'
ENABLE_AUTO_COMMIT=True

LSR_TOPIC = 'cdc.musicGraph.labelSoundRecording'
GSR_TOPIC = 'cdc.musicGraph.globalSoundRecording'
REL_TOPIC = 'cdc.musicGraph.represents'

TOPICS = [LSR_TOPIC, GSR_TOPIC, REL_TOPIC]

RECORDS_TO_COMPARE = {
    LSR_TOPIC: {},
    GSR_TOPIC: {},
    REL_TOPIC: {}
}

# KAFKA CONSUMER CONFIGURATION
BOOTSTRAP_SERVERS_44 = [
    "b-1.devmanagedkafkacdcdes.wwhlmg.c23.kafka.us-east-1.amazonaws.com:9094",
    "b-2.devmanagedkafkacdcdes.wwhlmg.c23.kafka.us-east-1.amazonaws.com:9094",
    "b-3.devmanagedkafkacdcdes.wwhlmg.c23.kafka.us-east-1.amazonaws.com:9094"
]
BOOTSTRAP_SERVERS_35 = [
    "b-1.dev-managed-kafka-cdc.rk4es0.c11.kafka.us-east-1.amazonaws.com:9094",
    "b-2.dev-managed-kafka-cdc.rk4es0.c11.kafka.us-east-1.amazonaws.com:9094",
    "b-2.dev-managed-kafka-cdc.rk4es0.c11.kafka.us-east-1.amazonaws.com:9094"
]

def get_kafka_consumer(topics, brokers) -> AIOKafkaConsumer:
    """Factory which returns configured kafka consumer."""
    return AIOKafkaConsumer(
        *topics,
        value_deserializer=lambda v: json.loads(str(v.decode())),
        bootstrap_servers=brokers,
        security_protocol=SECURITY_PROTOCOL,
        ssl_context=create_ssl_context(),
        enable_auto_commit=ENABLE_AUTO_COMMIT,
        auto_offset_reset=AUTO_OFFSET_RESET,
        group_id=GROUP_ID,
        client_id=str(uuid4())
    )

async def consume(brokers, handler=None):
    """Method to consume the messages."""
    consumer = get_kafka_consumer(TOPICS, brokers)

    await consumer.start()
    try:
        async for msg in consumer:
            if handler is not None:
                await handler(msg)
    finally:
        await consumer.stop()
        await asyncio.sleep(20)

async def handler_44(msg):
    id = ''
    if msg.topic in (LSR_TOPIC, GSR_TOPIC):
        id = msg.value['payload']['before']['properties']['id']
    if msg.topic == REL_TOPIC:
        # since relations do not have uuids let's use combined key
        id = msg.value['payload']['start']['ids']['id'] + msg.value['payload']['end']['ids']['id']

    # use node uuid as record key in pushed records
    RECORDS_TO_COMPARE[msg.topic][id] = msg.value
    record_len = len(RECORDS_TO_COMPARE[msg.topic])

    # log only last 10 pushes. 999 msgs to be pushed
    if record_len > 990:
        print('pushing records:', msg.topic, record_len)
        
async def handler_35(msg):
    id = ''
    if msg.topic in (LSR_TOPIC, GSR_TOPIC):
        id = msg.value['payload']['before']['properties']['id']
    if msg.topic == REL_TOPIC:
        id = msg.value['payload']['start']['ids']['id'] + msg.value['payload']['end']['ids']['id']

    try:
        # get record to compare with from the pushed ones
        stack_msg = RECORDS_TO_COMPARE[msg.topic][id]

        # we only want to compare message structure. 
        # data contained in the properties cannot be the same 
        # especially timestamps
        
        # metadata key names
        meta = msg.value['meta'].keys()

        # payload key names
        payload = msg.value['payload'].keys()

        # schema key names
        schema = msg.value['schema']

        # there is a small prop name drift because we do not refresh 
        # neo4j and thus on some nodes in active 3.5 dev cluster 
        # there are additional properties not present in 4.4 
        schema_props = set(schema['properties'])
        stack_schema_props = set(stack_msg['schema']['properties'])

        # unset schema_props to compare node/relation labels
        schema['properties'] = {}
        stack_msg['schema']['properties'] = {}

        # if everything is fine the RECORDS_TO_COMPARE should 
        # be popped from the dict where keys are uuids of nodes
        if (stack_msg['meta'].keys() == meta and 
                stack_msg['payload'].keys() == payload and 
                stack_msg['schema'] == schema and 
                stack_schema_props.issubset(schema_props) or 
                schema_props.issubset(stack_schema_props)):
            RECORDS_TO_COMPARE[msg.topic].pop(id, None)

            record_len = len(RECORDS_TO_COMPARE[msg.topic])

            # log only last 10 pushes, no need to print all 999
            if record_len < 10:
                print('records left to compare:', msg.topic, record_len)
        else:
            # always print if there is structure mismatch between 3.5 and 4.4
            print('structure match failed.', msg.topic, id, schema_props, stack_schema_props)

    except KeyError:
        print('msg not found. id:', id)


loop = asyncio.get_event_loop()
loop.create_task(consume(BOOTSTRAP_SERVERS_44, handler_44))
loop.create_task(consume(BOOTSTRAP_SERVERS_35, handler_35))
loop.run_forever()