{%- import "util.include.jinja" as util -%}
{%- set messagegroups = root.messagegroups %}
"""

{%- filter wordwrap(80) %}
This is sample code to consume events from Apache Kafka with the
dispatcher client{%- if messagegroups|length > 1 -%}s{% endif %} contained in this project.

There is a handler for each defined event type. The handler is a function that takes the following parameters:
- record: The Kafka record.
- cloud_event: The CloudEvent data.
- event_data: The event data.

Handlers are registered with the dispatcher client{%- if messagegroups|length > 1 -%} associated with their respective message group{% endif %}. A "dispatcher" is a class that inspects incoming events and routes them to the appropriate handler based on the event type.

{%- if messagegroups|length == 1 -%}
The main function creates a dispatcher client and registers the handlers for the events in the message group. Then it creates an event processor and starts it.

The event processor will call the appropriate handler for each event received.
{%- else -%}
The main function creates a dispatcher client for each message group and registers the handlers for the events in the message group. Then it creates an event processor and adds the dispatcher clients to it.

The event processor will call the appropriate handler for each event received.
{%- endif %}

This sample uses a Kafka cluster with the necessary topics and configurations. The script reads the configuration from the command line or uses the environment variables. The following environment variables are recognized:

- BOOTSTRAP_SERVERS: The Kafka bootstrap servers.
- GROUP_ID: The Kafka consumer group ID.
- TOPICS: The Kafka topics to subscribe to.

Alternatively, you can pass the configuration as command-line arguments.

python sample.py --bootstrap-servers <bootstrap_servers> --group-id <group_id> --topics <topics>

The main function waits for a signal (Press Ctrl+C) to stop the event processor.
{%- endfilter %}
"""

import argparse
import asyncio
import os
from confluent_kafka import Consumer, Message
from cloudevents.abstract import CloudEvent

{% set messagegroups = root.messagegroups %}
{%- set imports = [] %}
{%- for messagegroupid, messagegroup in messagegroups.items() -%}
{%- for _, message in messagegroup.messages.items() -%}
{%- set type_name = util.DeclareDataType(data_project_name, root, message) %}
{%- if type_name != "object" %}
{%- set import_statement = "from " + (type_name | lower) + " import " + type_name | strip_namespace %}
{%- if import_statement not in imports %}
{%- set _ = imports.append(import_statement) %}
{{ import_statement }}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- endfor %}

{% for messagegroupid, messagegroup in messagegroups.items() %}
{%- set groupname = messagegroupid | pascal %}
{%- set class_name = ( groupname | strip_dots ) + "EventDispatcher" %}
from {{main_project_name}}.dispatcher import {{class_name}}

{% for messageid, message in messagegroup.messages.items() %}
{%- set type_name = util.DeclareDataType( data_project_name, root, message ) %}
async def handle_{{ messageid | dotunderscore | snake }}(dispatcher: {{class_name}}, record: Message, cloud_event: CloudEvent, {{ messageid | dotunderscore | snake }}_event_data: {{ type_name | strip_namespace }}):
    """ Handles the {{ messageid }} event """
    print(f"{{ messageid }}: { {{- messageid | dotunderscore | snake -}}_event_data.to_json()}")
    # await some_processing_function(record, cloud_event, {{ messageid | dotunderscore | snake }}_event_data)

{% endfor %}
{%- endfor %}

async def main(bootstrap_servers, group_id, topics):
    """ Main function to consume events from Kafka """
    {%- if messagegroups|length == 1 -%}
    {%- set messagegroups = root.messagegroups %}
    {%- for messagegroupid, messagegroup in messagegroups.items() %}
    {%- set groupname = messagegroupid | pascal %}
    {%- set class_name = ( groupname | strip_dots ) + "EventDispatcher" %}
    dispatcher = {{ class_name }}()
    {%- for messageid, message in messagegroup.messages.items() %}
    dispatcher.{{ messageid | dotunderscore | snake }}_async = handle_{{ messageid | dotunderscore | snake }}
    {%- endfor %}
    {%- endfor %}
    consumer = Consumer()
    async with dispatcher.create_consumer(
                kafka_config={
                    'bootstrap.servers': bootstrap_servers,
                    'group.id': group_id,
                    'auto.offset.reset': 'earliest'
                },
                topics=topics.split(',')):
        # Terminate after 60 secs
        await asyncio.sleep(60)
    {%- else -%}
    consumers = []
    for topic in topics.split(','):
        consumer = Consumer({
            'bootstrap.servers': bootstrap_servers,
            'group.id': group_id,
            'auto.offset.reset': 'earliest'
        })
        consumer.subscribe([topic])
        consumers.append(consumer)

    try:
        while True:
            for consumer in consumers:
                msg = consumer.poll(timeout=1.0)
                if msg is None:
                    continue
                if msg.error():
                    if msg.error().code() == KafkaError._PARTITION_EOF:
                        continue
                    else:
                        raise KafkaException(msg.error())
                {%- for messagegroupid, messagegroup in messagegroups.items() %}
                {%- set groupname = messagegroupid | pascal %}
                {%- set class_name = ( groupname | strip_dots ) + "EventDispatcher" %}
                {{messagegroupid | dotunderscore | snake}}_dispatcher = {{ class_name }}()
                await {{messagegroupid | dotunderscore | snake}}_dispatcher._process_event(msg)
                {%- endfor %}
    finally:
        for consumer in consumers:
            consumer.close()
    {%- endif %}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Kafka Consumer")
    parser.add_argument('--bootstrap-servers', default=os.getenv('BOOTSTRAP_SERVERS', 'your_bootstrap_servers'), help='Kafka bootstrap servers', required=True)
    parser.add_argument('--group-id', default=os.getenv('GROUP_ID', 'your_group_id'), help='Kafka consumer group ID', required=True)
    parser.add_argument('--topics', default=os.getenv('TOPICS', 'your_topics'), help='Kafka topics to subscribe to', required=True)

    args = parser.parse_args()

    asyncio.run(main(
        args.bootstrap_servers,
        args.group_id,
        args.topics
    ))
