{% import "cloudevents.jinja.include" as cloudEvents -%}
{%- import "kafka.jinja.include" as kafka -%}
{%- import "util.include.jinja" as util -%}
{%- set messagegroups = root.messagegroups %}
{%- set uses_cloudevents_message = (root | exists("envelope","CloudEvents/1.0")) %}
{%- set uses_plain_kafka_message = (root | exists( "protocol", "kafka" )) %}
{%- set uses_kafka_endpoint = (root | exists( "protocol", "kafka" )) %}
{%- set function_name = project_name | pascal | strip_dots -%}
# pylint: disable=missing-function-docstring, wrong-import-position, import-error, no-name-in-module, import-outside-toplevel, no-member, redefined-outer-name, unused-argument, unused-variable, invalid-name, redefined-outer-name, missing-class-docstring

import asyncio
import logging
import os
import sys

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../{{data_project_name}}/src')))
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../{{data_project_name}}/tests')))
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../{{main_project_name}}/src')))

import tempfile
import pytest
from confluent_kafka import Producer
from confluent_kafka.admin import AdminClient, NewTopic
from testcontainers.kafka import KafkaContainer

{%- set imports = [] %}
{%- set test_imports = [] %}
{%- 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 ) %}
{%- if type_name != "object" %}
{%- set import_statement = "from " + (data_project_name | dotunderscore | lower) + " import " + (type_name | pascal | strip_namespace) %}
{%- if import_statement not in imports %}
{%- set _ = imports.append(import_statement) %}
{{ import_statement }}
{%- endif %}
{%- set test_class_name = "Test_" + (type_name | pascal | strip_namespace) %}
{%- set test_module_name = "test_" + (type_name | dotunderscore | lower) %}
{%- set test_import_statement = "from " + test_module_name + " import " + test_class_name %}
{%- if test_import_statement not in test_imports %}
{%- set _ = test_imports.append(test_import_statement) %}
{{ test_import_statement }}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- endfor %}

@pytest.fixture(scope="module")
def kafka_emulator():
    with KafkaContainer() as kafka:
        connection = kafka.get_bootstrap_server()
        admin_client = AdminClient({'bootstrap.servers': connection})
        topic_list = [NewTopic("test_topic", num_partitions=1, replication_factor=1)]
        admin_client.create_topics(topic_list)

        yield {
            "bootstrap_servers": connection
        }

{%- for messagegroupid, messagegroup in messagegroups.items() %}
{%- set groupname = messagegroupid | pascal %}
{%- set class_name = ( groupname | strip_dots ) + "EventDispatcher" %}
{%- set test_class_name = (project_name | pascal | strip_dots) + (groupname | strip_dots) + "Tests" %}

{%- for messageid, message in messagegroup.messages.items() %}
{%- set messagename = messageid | pascal | strip_dots | strip_namespace %}
{%- set test_function_name = "test_" + groupname | lower | replace(" ", "_") + "_" + messagename | lower | replace(" ", "_") %}

@pytest.mark.asyncio
async def {{ test_function_name | dotunderscore }}(kafka_emulator):
    """Test the {{ messagename }} message from the {{ groupname }} message group"""

    # Turn on logging if running in debugger
    if sys.gettrace() is not None:
        logging.basicConfig(level=logging.DEBUG)
        kafka_logger = logging.getLogger('kafka')
        kafka_logger.setLevel(logging.DEBUG)
        logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

    bootstrap_servers = kafka_emulator["bootstrap_servers"]

    dispatcher = {{ class_name }}()

    future = asyncio.Future()
    received_count = 0
    received_data_list = []
    
    async def set_future(consumer, record, cloud_event, data):
        nonlocal received_count
        received_count += 1
        received_data_list.append(data)
        if received_count >= 5:
            if not future.done():
                future.set_result(data)

    dispatcher.{{ messageid | dotunderscore | snake }}_async = set_future
    async with dispatcher.create_consumer({
                'bootstrap.servers': bootstrap_servers,
                'group.id': "test-group",
                'auto.offset.reset': 'earliest'
            },
            ["test_topic"]):
        producer = Producer({
            "bootstrap.servers":bootstrap_servers
        })
        {%- set type_name = util.DeclareDataType( data_project_name, root, message ) %}
        {%- if type_name != "object" %}
        # Create valid test data using the test helper
        test_instance = Test_{{ type_name | pascal | strip_namespace }}.create_instance()
        event_data = test_instance.to_json().encode('utf-8')
        {%- else %}
        event_data = b'{}'
        {%- endif %}
        
        # Send 5 messages to test proper message settlement and ordering
        for i in range(5):
            producer.produce(topic='test_topic', key=b'abc', value=event_data, headers=[
                ("ce_type", b"{{ messageid }}"),
                ("ce_specversion", b"1.0"),
                ("ce_source", b"/test"),
                ("ce_id", f"test-id-{i}".encode('utf-8')),
                ("content-type", b"application/json")
            ])
        producer.flush()
        
        # Wait for all 5 messages to be received
        await asyncio.wait_for(future, timeout=15)
        assert received_count == 5, f"Expected 5 messages, got {received_count}"
        assert len(received_data_list) == 5, f"Expected 5 data items, got {len(received_data_list)}"

{% endfor %}
{% endfor %}
