{%- 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" %}
{#- Look up Kafka key template from endpoint for this message group -#}
{%- set endpoint_kafka_key = [] %}
{%- if root.endpoints %}
{%- for endpointid, ep in root.endpoints.items() %}
{%- if "protocol" in ep and ep.protocol | lower == "kafka" and "messagegroups" in ep and ep.messagegroups %}
{%- for mgref in ep.messagegroups %}
{%- if mgref == "#/messagegroups/" + messagegroupid %}
{%- if "protocoloptions" in ep and "options" in ep.protocoloptions and "key" in ep.protocoloptions.options %}
{%- set _ = endpoint_kafka_key.append(ep.protocoloptions.options.key) %}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- endif %}
{%- endfor %}
{%- endif %}
{%- 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(" ", "_") %}
{#- Determine Kafka key for this message -#}
{%- set msg_key_holder = [] %}
{%- if "protocolmetadata" in message and "key" in message.protocolmetadata and "value" in message.protocolmetadata.key %}
{%- set _ = msg_key_holder.append(message.protocolmetadata.key.value) %}
{%- elif endpoint_kafka_key %}
{%- set _ = msg_key_holder.append(endpoint_kafka_key[-1]) %}
{%- endif %}
{%- set kafka_key_value = msg_key_holder[0] if msg_key_holder else None %}
{%- set declared_params = [] %}
{%- for attrname, attribute in message.envelopemetadata.items() if attribute.type == "uritemplate" -%}
{%- for ph in attribute.value | regex_search('\\{([A-Za-z0-9_]+)\\}') %}
{%- if (ph | snake) not in declared_params %}{%- set _ = declared_params.append(ph | snake) %}{%- endif %}
{%- endfor -%}
{%- endfor %}

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 None
            msg = consumer.poll(1.0)
            if msg is None:
                continue
            if msg.error():
                continue
            cloudevent = parse_cloudevent(msg)
            if cloudevent['type'] == "{{ messageid }}":
                return msg.key().decode('utf-8') if msg.key() else None

    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 -%}
            {%- if kafka_key_value -%}
                {%- for placeholder in kafka_key_value | regex_search('\\{([A-Za-z0-9_]+)\\}') -%}
                    {%- if (placeholder | snake) not in declared_params %}_{{ placeholder | snake }} = f'test_{i}', {% endif -%}
                {%- endfor -%}
            {%- endif -%}
            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 and assert Kafka key
    for i in range(5):
        received_key = on_event()
        assert received_key is not None, f"Failed to receive message {i+1} of 5"
        {%- if kafka_key_value %}
        {%- set kphs = kafka_key_value | regex_search('\\{([A-Za-z0-9_]+)\\}') %}
        {%- if kphs %}
        expected_key = "{{ kafka_key_value }}".format({% for ph in kphs %}{{ ph }}=f'test_{i}'{% if not loop.last %}, {% endif %}{% endfor %})
        assert received_key == expected_key, f"Expected Kafka key '{expected_key}' but got '{received_key}'"
        {%- else %}
        assert received_key == "{{ kafka_key_value }}", f"Expected Kafka key '{{ kafka_key_value }}' but got '{received_key}'"
        {%- endif %}
        {%- endif %}
    consumer.close()
{% endfor %}
{%- endfor %}

{#- Cross-event-type Kafka key test: verifies that different event types with the same
    key template in a message group produce the same Kafka key for the same placeholder values -#}
{%- for messagegroupid, messagegroup in root.messagegroups.items() %}
{%- set groupname = messagegroupid | pascal %}
{%- set class_name = ( groupname | strip_dots ) + "EventProducer" %}
{#- Look up Kafka key template from endpoint for this message group -#}
{%- set ep_key_holder = [] %}
{%- if root.endpoints %}
{%- for endpointid, ep in root.endpoints.items() %}
{%- if "protocol" in ep and ep.protocol | lower == "kafka" and "messagegroups" in ep and ep.messagegroups %}
{%- for mgref in ep.messagegroups %}
{%- if mgref == "#/messagegroups/" + messagegroupid %}
{%- if "protocoloptions" in ep and "options" in ep.protocoloptions and "key" in ep.protocoloptions.options %}
{%- set _ = ep_key_holder.append(ep.protocoloptions.options.key) %}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- endif %}
{%- endfor %}
{%- endif %}
{%- set group_kafka_key = ep_key_holder[0] if ep_key_holder else None %}
{#- Only generate cross-event-type test when group has >=2 messages sharing same key template -#}
{%- if group_kafka_key and messagegroup.messages | length >= 2 %}
{%- set msg_list = messagegroup.messages.items() | list %}

def test_{{ groupname | lower | replace(" ", "_") | dotunderscore }}_cross_event_type_kafka_key(kafka_emulator):
    """Test that different event types in {{ groupname }} produce the same Kafka key for the same placeholder values"""

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

    consumer = Consumer({
        'bootstrap.servers': bootstrap_servers,
        'group.id': 'test_{{ groupname | lower | dotunderscore }}_cross_key',
        'auto.offset.reset': 'latest'
    })
    consumer.subscribe([topic])

    import time
    assignment_timeout = time.time() + 10
    while not consumer.assignment() and time.time() < assignment_timeout:
        consumer.poll(0.1)
    if not consumer.assignment():
        pytest.fail(f"Consumer failed to get partition assignment within 10 seconds. Topic: {topic}")
    # Drain any pre-existing messages before producing our test messages
    drain_timeout = time.time() + 3
    while time.time() < drain_timeout:
        msg = consumer.poll(0.5)
    time.sleep(1)

    kafka_producer = Producer({'bootstrap.servers': bootstrap_servers})
    producer_instance = {{ class_name }}(kafka_producer, topic, 'binary')

    shared_key_value = "shared_entity_42"

    {#- Send one message of each of the first two message types with the same key placeholder -#}
    {%- set first_msg_id = msg_list[0][0] %}
    {%- set first_msg = msg_list[0][1] %}
    {%- set second_msg_id = msg_list[1][0] %}
    {%- set second_msg = msg_list[1][1] %}
    {%- set first_type_name = util.DeclareDataType( data_project_name, root, first_msg ) %}
    {%- set second_type_name = util.DeclareDataType( data_project_name, root, second_msg ) %}

    {%- if first_type_name != "object" %}
    data1 = Test_{{ first_type_name | pascal | strip_namespace }}.create_instance()
    {%- else %}
    data1 = {}
    {%- endif %}
    {%- if second_type_name != "object" %}
    data2 = Test_{{ second_type_name | pascal | strip_namespace }}.create_instance()
    {%- else %}
    data2 = {}
    {%- endif %}

    producer_instance.send_{{ first_msg_id | dotunderscore | snake }}(
        {%- for attrname, attribute in first_msg.envelopemetadata.items() if attribute.required and attribute.value is not defined -%}
            _{{ attrname }} = 'test',
        {%- endfor -%}
        {%- for attrname, attribute in first_msg.envelopemetadata.items() if attribute.type == "uritemplate" -%}
            {%- for placeholder in attribute.value | regex_search('\\{([A-Za-z0-9_]+)\\}') %}_{{ placeholder | snake }} = shared_key_value, {% endfor -%}
        {%- endfor -%}
        {%- set first_msg_declared_params = [] %}
        {%- for attrname, attribute in first_msg.envelopemetadata.items() if attribute.type == "uritemplate" -%}
        {%- for ph in attribute.value | regex_search('\\{([A-Za-z0-9_]+)\\}') %}
        {%- if (ph | snake) not in first_msg_declared_params %}{%- set _ = first_msg_declared_params.append(ph | snake) %}{%- endif %}
        {%- endfor -%}
        {%- endfor -%}
        {%- if group_kafka_key -%}
            {%- for placeholder in group_kafka_key | regex_search('\\{([A-Za-z0-9_]+)\\}') -%}
                {%- if (placeholder | snake) not in first_msg_declared_params %}_{{ placeholder | snake }} = shared_key_value, {% endif -%}
            {%- endfor -%}
        {%- endif -%}
        data = data1
        {%- for attrname, attribute in first_msg.envelopemetadata.items() if not attribute.required and attribute.value is not defined -%}
            , _{{ attrname }} = ''
        {%- endfor -%}
    )
    producer_instance.send_{{ second_msg_id | dotunderscore | snake }}(
        {%- for attrname, attribute in second_msg.envelopemetadata.items() if attribute.required and attribute.value is not defined -%}
            _{{ attrname }} = 'test',
        {%- endfor -%}
        {%- for attrname, attribute in second_msg.envelopemetadata.items() if attribute.type == "uritemplate" -%}
            {%- for placeholder in attribute.value | regex_search('\\{([A-Za-z0-9_]+)\\}') %}_{{ placeholder | snake }} = shared_key_value, {% endfor -%}
        {%- endfor -%}
        {%- set second_msg_declared_params = [] %}
        {%- for attrname, attribute in second_msg.envelopemetadata.items() if attribute.type == "uritemplate" -%}
        {%- for ph in attribute.value | regex_search('\\{([A-Za-z0-9_]+)\\}') %}
        {%- if (ph | snake) not in second_msg_declared_params %}{%- set _ = second_msg_declared_params.append(ph | snake) %}{%- endif %}
        {%- endfor -%}
        {%- endfor -%}
        {%- if group_kafka_key -%}
            {%- for placeholder in group_kafka_key | regex_search('\\{([A-Za-z0-9_]+)\\}') -%}
                {%- if (placeholder | snake) not in second_msg_declared_params %}_{{ placeholder | snake }} = shared_key_value, {% endif -%}
            {%- endfor -%}
        {%- endif -%}
        data = data2
        {%- for attrname, attribute in second_msg.envelopemetadata.items() if not attribute.required and attribute.value is not defined -%}
            , _{{ attrname }} = ''
        {%- endfor -%}
    )
    kafka_producer.flush(timeout=5.0)

    # Collect keys from both messages
    collected_keys = []
    timeout = time.time() + 20
    while len(collected_keys) < 2 and time.time() < timeout:
        msg = consumer.poll(1.0)
        if msg is None or msg.error():
            continue
        cloudevent = parse_cloudevent(msg)
        if cloudevent['type'] in ["{{ first_msg_id }}", "{{ second_msg_id }}"]:
            key = msg.key().decode('utf-8') if msg.key() else None
            collected_keys.append(key)

    assert len(collected_keys) == 2, f"Expected 2 messages but received {len(collected_keys)}"
    assert collected_keys[0] == collected_keys[1], \
        f"Expected same Kafka key for different event types but got '{collected_keys[0]}' and '{collected_keys[1]}'"
    {%- set gkphs = group_kafka_key | regex_search('\\{([A-Za-z0-9_]+)\\}') %}
    {%- if gkphs %}
    expected_key = "{{ group_kafka_key }}".format({% for ph in gkphs %}{{ ph }}=shared_key_value{% if not loop.last %}, {% endif %}{% endfor %})
    {%- else %}
    expected_key = "{{ group_kafka_key }}"
    {%- endif %}
    assert collected_keys[0] == expected_key, \
        f"Expected Kafka key '{expected_key}' but got '{collected_keys[0]}'"
    consumer.close()
{%- endif %}
{%- endfor %}
