{%- import "util.include.jinja" as util -%}
# 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
import datetime
from typing import Optional

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, Consumer, KafkaException, Message
from confluent_kafka.admin import AdminClient, NewTopic
from cloudevents.abstract import CloudEvent
from cloudevents.kafka import from_binary, from_structured, KafkaMessage
from testcontainers.kafka import KafkaContainer

{%- set imports = [] %}
{%- set test_imports = [] %}
{%- for messagegroupid, messagegroup in root.messagegroups.items() %}
{%- set groupname = messagegroupid | pascal %}
{%- set class_name = ( groupname | strip_dots ) + "EventProducer" %}
from {{main_project_name}}.producer 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) %}
{%- if message.dataschemaformat and "JSONStructure" in message.dataschemaformat %}
{%- set test_module_name = "test_" + (type_name | pascal | strip_namespace | lower) %}
{%- else %}
{%- set flat_package = (type_name | dotunderscore | lower) %}
{%- if flat_package.startswith('test_') %}
{%- set test_module_name = flat_package %}
{%- else %}
{%- set test_module_name = "test_" + flat_package %}
{%- endif %}
{%- endif %}
{%- 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:
        admin_client = AdminClient({'bootstrap.servers': kafka.get_bootstrap_server()})
        topic_list = [
            NewTopic("test_topic", num_partitions=1, replication_factor=1)
        ]
        admin_client.create_topics(topic_list)

        yield {
            "bootstrap_servers": kafka.get_bootstrap_server(),
            "topic": "test_topic",
        }

def parse_cloudevent(msg: Message) -> CloudEvent:
    headers_dict: Dict[str, bytes] = {header[0]: header[1] for header in msg.headers()}
    message = KafkaMessage(headers=headers_dict, key=msg.key(), value=msg.value())
    if message.headers and 'content-type' in message.headers:
        content_type = message.headers['content-type'].decode()
        if content_type.startswith('application/cloudevents'):
            ce = from_structured(message)
            if 'datacontenttype' not in ce:
                ce['datacontenttype'] = 'application/json'
        else:
            ce = from_binary(message)
            ce['datacontenttype'] = message.headers['content-type'].decode()
    else:
        ce = from_binary(message)
        ce['datacontenttype'] = 'application/json'
    return ce

{%- for messagegroupid, messagegroup in root.messagegroups.items() %}
{%- set groupname = messagegroupid | pascal %}
{%- set class_name = ( groupname | strip_dots ) + "EventProducer" %}
{%- 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(" ", "_") %}

def {{ test_function_name | dotunderscore }}(kafka_emulator):
    """Test the {{ messagename }} event from the {{ groupname }} message group"""

    bootstrap_servers = kafka_emulator["bootstrap_servers"]
    topic = kafka_emulator["topic"]

    producer = Producer({'bootstrap.servers': bootstrap_servers})
    consumer = Consumer({
        'bootstrap.servers': bootstrap_servers,
        'group.id': '{{ test_function_name | dotunderscore }}',  # Unique group per test
        'auto.offset.reset': 'earliest'
    })
    consumer.subscribe([topic])
    
    # Wait for partition assignment before producing messages
    import time
    assignment_timeout = time.time() + 10
    while not consumer.assignment() and time.time() < assignment_timeout:
        consumer.poll(0.1)
    
    # Verify partition assignment succeeded
    if not consumer.assignment():
        pytest.fail(f"Consumer failed to get partition assignment within 10 seconds. Topic: {topic}")
    
    # Give consumer time to stabilize and seek to beginning
    time.sleep(1)

    def on_event():
        import time
        timeout = time.time() + 20  # 20 second timeout for CI robustness
        while True:
            if time.time() > timeout:
                return False
            msg = consumer.poll(1.0)
            if msg is None:
                continue
            if msg.error():
                continue
            cloudevent = parse_cloudevent(msg)
            if cloudevent['type'] == "{{ messageid }}":
                return True

    kafka_producer = Producer({'bootstrap.servers': bootstrap_servers})
    producer_instance = {{ class_name }}(kafka_producer, topic, 'binary')
    {%- set type_name = util.DeclareDataType( data_project_name, root, message ) %}
    {%- if type_name != "object" %}
    # Create valid test data using the test helper
    event_data = Test_{{ type_name | pascal | strip_namespace }}.create_instance()
    {%- else %}
    event_data = {}
    {%- endif %}
    
    # Send 5 messages to test message settlement and ordering
    for i in range(5):
        producer_instance.send_{{ messageid | dotunderscore | snake }}(
            {%- for attrname, attribute in message.envelopemetadata.items() if attribute.required and attribute.value is not defined -%}
                _{{ attrname }} = {{ "f'test_{i}'" if not attrname == 'time' else 'datetime.datetime.now().isoformat()' }},
            {%- endfor -%}
            {%- for attrname, attribute in message.envelopemetadata.items() if attribute.type == "uritemplate" -%}
                {%- for placeholder in attribute.value | regex_search('\\{([A-Za-z0-9_]+)\\}') %}_{{ placeholder | snake }} = f'test_{i}', {% endfor -%}
            {%- endfor -%}
            data = event_data
            {%- for attrname, attribute in message.envelopemetadata.items() if not attribute.required and attribute.value is not defined -%}
                , _{{ attrname }} = ''
            {%- endfor -%}
        )
    
    # Flush producer to ensure messages are sent before consumer polling
    kafka_producer.flush(timeout=5.0)

    # Verify all 5 messages received
    for i in range(5):
        assert on_event(), f"Failed to receive message {i+1} of 5"
    consumer.close()

{%- endfor %}
{%- endfor %}
