{%- import "util.include.jinja" as util -%}

# pylint: disable=unused-import, line-too-long, missing-module-docstring, missing-function-docstring, missing-class-docstring, consider-using-f-string, trailing-whitespace, trailing-newlines
import sys
import json
import uuid
import typing
from datetime import datetime
from confluent_kafka import Producer, KafkaException, Message
from cloudevents.kafka import to_binary, to_structured, KafkaMessage
from cloudevents.http import CloudEvent

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

{%- for messagegroupid, messagegroup in messagegroups.items() -%}
{%- set parts = messagegroupid.split('.') -%}
{%- set groupName = parts | join('_') | lower %}
{%- set groupname = messagegroupid | pascal %}
{%- set class_name = ( groupname | strip_dots ) + "EventProducer" %}

class {{ class_name }}:
    def __init__(self, producer: Producer, topic: str, content_mode:typing.Literal['structured','binary']='structured'):
        """
        Initializes the Kafka producer

        Args:
            producer (Producer): The Kafka producer client
            topic (str): The Kafka topic to send events to
            content_mode (typing.Literal['structured','binary']): The content mode to use for sending events
        """
        self.producer = producer
        self.topic = topic
        self.content_mode = content_mode

    def __key_mapper(self, x: CloudEvent, m: typing.Any, key_mapper: typing.Callable[[CloudEvent, typing.Any], str]) -> str:
        """
        Maps a CloudEvent to a Kafka key

        Args:
            x (CloudEvent): The CloudEvent to map
            m (Any): The event data
            key_mapper (Callable[[CloudEvent, Any], str]): The user's key mapper function
        """
        if key_mapper:
            return key_mapper(x, m)
        else:
            return f'{str(x.get("type"))}:{str(x.get("source"))}{("-"+str(x.get("subject"))) if x.get("subject") else ""}'
    {%- for messageid, message in messagegroup.messages.items() -%}
    {%- set message_snake = messageid | dotunderscore | snake %}
    {%- set data_type = util.DeclareDataType( data_project_name, root, message ) %}

    def send_{{message_snake}}(self,
        {%- for attrname in ['source', 'type'] if attrname not in message.envelopemetadata -%}
        _{{ attrname }}: str, {{ '' }}
        {%- endfor -%}
        {%- for attrname, attribute in message.envelopemetadata.items() if attribute.required and attribute.value is not defined -%}
        _{{ attrname }}: str, {{ '' }}
        {%- 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 }} : str, {% endfor -%}
        {%- endfor -%}
        data: {{ data_type | strip_namespace -}}, content_type: str = "application/json"
        {%- for attrname, attribute in message.envelopemetadata.items() if not attribute.required and attribute.value is not defined -%}
            , _{{ attrname }}: typing.Optional[str] = None
        {%- endfor -%}, flush_producer=True, key_mapper: typing.Callable[[CloudEvent, {{ data_type | strip_namespace -}}], str]=None) -> None:
        """
        Sends the '{{ messageid }}' event to the Kafka topic

        Args:
        {%- for attrname in ['source', 'type'] if attrname not in message.envelopemetadata %}
            _{{ attrname }}(str) : CloudEvents required attribute '{{ attrname }}'
        {%- endfor -%}
        {%- for attrname, attribute in message.envelopemetadata.items() if attribute.required and attribute.value is not defined %}
            _{{ attrname }}(str) : {{ attribute.description if attribute.description else "CloudEvents attribute '"+attrname+"'" }}
        {%- 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 }}(str):  Value for placeholder {{ placeholder }} in attribute {{ attrname }}
            {%- endfor -%}
        {%- endfor %}
            data: ({{ data_type | strip_namespace -}}): The event data to be sent
            content_type (str): The content type that the event data shall be sent with
        {%- for attrname, attribute in message.envelopemetadata.items() if not attribute.required and attribute.value is not defined %}
            _{{ attrname }}(typing.Optional[str]): {{ attribute.description if attribute.description else "CloudEvents attribute '"+attrname+"'" }}
        {%- endfor %}
            flush_producer(bool): Whether to flush the producer after sending the event (default: True)
            key_mapper(Callable[[CloudEvent, {{ data_type | strip_namespace -}}], str]): A function to map the CloudEvent contents to a Kafka key (default: None). 
                The default key mapper maps the CloudEvent type, source, and subject to the Kafka key
        """
        attributes = {
        {%- for attrname in ['source', 'type'] if attrname not in message.envelopemetadata %}
             "{{ attrname }}": _{{ attrname }},
        {%- endfor -%}
        {%- for attrname, attribute in message.envelopemetadata.items() %}
             "{{ attrname }}":
            {%- if attribute.value -%}
            "{{ attribute.value }}"
            {%- if attribute.type == "uritemplate" -%}
                {%- set phs = attribute.value | regex_search('\\{([A-Za-z0-9_]+)\\}') %}
                {%- if phs -%}
                .format(
                {%- for placeholder in phs -%}
                    {{ placeholder }} = _{{ placeholder | snake }}
                {%- if not loop.last -%},{%- endif%}
                {%- endfor -%}
                )
                {%- endif %}
            {%- endif %}
            {%- else -%}
                _{{ attrname }}
            {%- endif-%}
            {%- if not loop.last -%},{%- endif %}
        {%- endfor %}
        }
        attributes["datacontenttype"] = content_type
        event = CloudEvent.create(attributes, data)
        if self.content_mode == "structured":
            message = to_structured(event, data_marshaller=lambda x: {% if data_type != "object"%}json.loads(x.to_json()){%else%}json.dumps(x){%endif%}, key_mapper=lambda x: self.__key_mapper(x, data, key_mapper))
            message.headers[b"content-type"] = b"application/cloudevents+json"
        else:
            # For binary mode, datacontenttype is already set in attributes above
            # The to_binary() function will create the ce_datacontenttype header
            message = to_binary(event, data_marshaller=lambda x: {% if data_type != "object"%}x.to_byte_array("application/json"){%else%}json.dumps(x){%endif%}, key_mapper=lambda x: self.__key_mapper(x, data, key_mapper))
        self.producer.produce(self.topic, key=message.key, value=message.value, headers=message.headers)
        if flush_producer:
            self.producer.flush()
{% endfor %}

    @classmethod
    def parse_connection_string(cls, connection_string: str) -> typing.Tuple[typing.Dict[str, str], str]:
        """
        Parse the connection string and extract bootstrap server, topic name, username, and password.

        Args:
            connection_string (str): The connection string.

        Returns:
            Tuple[Dict[str, str], str]: Kafka config, topic name
        """
        config_dict = {
            'security.protocol': 'SASL_SSL',
            'sasl.mechanisms': 'PLAIN',
            'sasl.username': '$ConnectionString',
            'sasl.password': connection_string.strip()
        }
        kafka_topic = None
        try:
            for part in connection_string.split(';'):
                if 'Endpoint' in part:
                    config_dict['bootstrap.servers'] = part.split('=')[1].strip(
                        '"').replace('sb://', '').replace('/', '')+':9093'
                elif 'EntityPath' in part:
                    kafka_topic = part.split('=')[1].strip('"')
        except IndexError as e:
            raise ValueError("Invalid connection string format") from e
        return config_dict, kafka_topic

    @classmethod
    def from_connection_string(cls, connection_string: str, topic: typing.Optional[str]=None, content_mode: typing.Literal['structured','binary']='structured') -> '{{ class_name }}':
        """
        Create a Kafka producer from a connection string and a topic name.

        Args:
            connection_string (str): The connection string.
            topic (Optional[str]): The Kafka topic.
            content_mode (typing.Literal['structured','binary']): The content mode to use for sending events

        Returns:
            Producer: The Kafka producer
        """
        config, topic_name = cls.parse_connection_string(connection_string)
        if topic:
            topic_name = topic
        if not topic_name:
            raise ValueError("Topic name not found in connection string")
        return cls(Producer(config), topic_name, content_mode)

{% endfor %}
