{%- 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 json
import re
import typing
from typing import Callable, Awaitable, Optional, Dict, List
import asyncio
import paho.mqtt.client as mqtt
from cloudevents.conversion import to_binary, to_structured
from cloudevents.http import CloudEvent
import {{ data_project_name | dotunderscore | lower }}


{%- 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 %}
{%- 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 %}


# URI template regex pattern
_URI_TEMPLATE_PATTERN = re.compile(r'\{([A-Za-z0-9_]+)\}')


def _topic_to_mqtt_wildcard(topic: str) -> str:
    """Convert URI template placeholders to MQTT wildcards (+)."""
    return _URI_TEMPLATE_PATTERN.sub('+', topic)


def _apply_topic_template(topic_pattern: str, template_values: Optional[Dict[str, str]] = None) -> str:
    """Apply template values to a topic pattern."""
    if not template_values:
        return topic_pattern
    result = topic_pattern
    for key, value in template_values.items():
        result = result.replace('{' + key + '}', value)
    return result


def _build_topic_regex(topic_pattern: str) -> re.Pattern:
    """Build a regex pattern for matching topics and extracting template values."""
    # Escape regex special chars in the topic pattern
    escaped = re.escape(topic_pattern)
    # Restore placeholders (which were escaped to \{...\}) and convert to named groups
    pattern = re.sub(r'\\{([A-Za-z0-9_]+)\\}', r'(?P<\1>[^/]+)', escaped)
    return re.compile('^' + pattern + '$')


def _extract_topic_parameters(topic: str, pattern: re.Pattern) -> Dict[str, str]:
    """Extract URI template placeholder values from a topic."""
    match = pattern.match(topic)
    if match:
        return {k: v for k, v in match.groupdict().items() if v is not None}
    return {}


class _ClientBase:
    """Base class for MQTT client with CloudEvent detection."""
    
    @staticmethod
    def _is_cloud_event(message: mqtt.MQTTMessage) -> bool:
        """Check if the MQTT message contains a CloudEvent."""
        try:
            if message.payload:
                # Check for structured mode (JSON with CloudEvent fields)
                if isinstance(message.payload, bytes):
                    payload_str = message.payload.decode('utf-8')
                    data = json.loads(payload_str)
                    return all(k in data for k in ['specversion', 'type', 'source', 'id'])
                # Check for binary mode would require inspecting user properties (MQTT 5.0)
                # For now, we focus on structured mode detection
            return False
        except (json.JSONDecodeError, UnicodeDecodeError, AttributeError):
            return False
    
    @staticmethod
    def _cloud_event_from_message(message: mqtt.MQTTMessage) -> Optional[CloudEvent]:
        """Extract CloudEvent from MQTT message (structured mode)."""
        try:
            if isinstance(message.payload, bytes):
                payload_str = message.payload.decode('utf-8')
                # Use CloudEvents library's from_json to properly deserialize
                from cloudevents.http import from_json
                event = from_json(payload_str)
                return event
            return None
        except (json.JSONDecodeError, UnicodeDecodeError, KeyError, Exception):
            return None

{% 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 )+"MqttClient" %}

{#- Helper macro to get the default topic for a message -#}
{%- macro get_default_topic(message, messageid) -%}
{%- if message.protocoloptions and message.protocoloptions.properties and message.protocoloptions.properties.topic -%}
{{ message.protocoloptions.properties.topic.value }}
{%- else -%}
{{ messageid }}
{%- endif -%}
{%- endmacro %}

def get_default_topic_mappings_{{ groupName }}() -> Dict[str, str]:
    """
    Get the default topic mappings for the {{ messagegroupid }} message group.
    
    Returns:
        Dictionary mapping message identifiers to their default topic patterns.
    """
    return {
        {%- for messageid, message in messagegroup.messages.items() %}
        "{{ messageid }}": "{{ get_default_topic(message, messageid) }}",
        {%- endfor %}
    }


def get_subscription_topics_{{ groupName }}(topic_mappings: Optional[Dict[str, str]] = None) -> List[str]:
    """
    Get subscription topics with URI template placeholders replaced by MQTT wildcards (+).
    
    Args:
        topic_mappings: Optional topic mappings. If None, uses default mappings.
        
    Returns:
        List of topic patterns suitable for MQTT subscription.
    """
    mappings = topic_mappings or get_default_topic_mappings_{{ groupName }}()
    topics = []
    for topic in mappings.values():
        wildcard_topic = _topic_to_mqtt_wildcard(topic)
        if wildcard_topic not in topics:
            topics.append(wildcard_topic)
    return topics


class {{ class_name }}(_ClientBase):
    """MQTT Client for producing and consuming messages in the {{ messagegroupid }} message group."""
    
    def __init__(
        self, 
        client: mqtt.Client, 
        topic_mappings: Optional[Dict[str, str]] = None,
        content_mode: typing.Literal['structured', 'binary'] = 'structured', 
        loop: Optional[asyncio.AbstractEventLoop] = None
    ):
        """
        Initialize the MQTT client.

        Args:
            client: Paho MQTT client instance
            topic_mappings: Optional dictionary mapping message identifiers to topic patterns.
                Topic patterns may be URI templates with placeholders like {placeholder}.
                For consumers, placeholders are converted to MQTT wildcards (+) for subscription.
                If not provided, uses default 1:1 mapping.
            content_mode: The content mode for CloudEvents ('structured' or 'binary')
            loop: Optional event loop to use for async operations. If None, will try to get the running loop.
        """
        self.client = client
        self.content_mode = content_mode
        self.loop = loop
        self._topic_mappings = topic_mappings or get_default_topic_mappings_{{ groupName }}()
        self._topic_patterns = {k: _build_topic_regex(v) for k, v in self._topic_mappings.items()}
        
        # Message handler callbacks (Dispatcher pattern)
        {% for messageid, message in messagegroup.messages.items() if (message | exists("envelope","CloudEvents/1.0")) -%}
        {%- set messagename = messageid | dotunderscore | snake %}
        {%- set type_name = util.DeclareDataType( data_project_name, root, message ) %}
        self.{{messagename.split('.')[-1]}}_async: Optional[Callable[[mqtt.MQTTMessage, CloudEvent, {{type_name}}, Dict[str, str]], Awaitable[None]]] = None
        {% endfor %}
        
        # Attach message callback
        self.client.on_message = self._on_message

    @property
    def topic_mappings(self) -> Dict[str, str]:
        """Get the current topic mappings."""
        return self._topic_mappings.copy()
    
    def _on_message(self, client, userdata, message: mqtt.MQTTMessage):
        """Internal MQTT message callback that dispatches to async handlers."""
        loop = self.loop
        if loop is None:
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = asyncio.get_event_loop()
        asyncio.run_coroutine_threadsafe(self._process_message(message), loop)
    
    async def _process_message(self, message: mqtt.MQTTMessage):
        """Process incoming MQTT message and dispatch to appropriate handler."""
        try:
            if self._is_cloud_event(message):
                cloud_event = self._cloud_event_from_message(message)
                if cloud_event:
                    await self._dispatch_cloud_event(message, cloud_event)
            else:
                # Handle plain MQTT messages (if needed)
                pass
        except Exception as e:
            print(f"Error processing message: {e}")
    
    def _extract_topic_params(self, topic: str, message_id: str) -> Dict[str, str]:
        """Extract URI template placeholder values from a topic."""
        if message_id in self._topic_patterns:
            return _extract_topic_parameters(topic, self._topic_patterns[message_id])
        return {}
    
    async def _dispatch_cloud_event(self, mqtt_message: mqtt.MQTTMessage, cloud_event: CloudEvent):
        """Dispatch CloudEvent to the appropriate handler based on type."""
        event_type = cloud_event['type']
        
        {% for messageid, message in messagegroup.messages.items() if (message | exists("envelope","CloudEvents/1.0")) -%}
        {%- set messagename = messageid | dotunderscore | snake %}
        {%- set type_name = util.DeclareDataType( data_project_name, root, message ) %}
        {%- if "type" in message.envelopemetadata and "value" in message.envelopemetadata["type"] %}
        if event_type == "{{ message.envelopemetadata["type"]["value"] }}":
        {%- else %}
        if event_type == "{{ messageid }}":
        {%- endif %}
            if self.{{messagename.split('.')[-1]}}_async:
                try:
                    content_type = cloud_event.get_attributes().get('datacontenttype', 'application/json')
                    {%- if type_name == "object" %}
                    data = cloud_event.data
                    {%- else %}
                    # CloudEvent.data is now a dict or string, not bytes
                    data = {{type_name}}.from_data(cloud_event.data, content_type)
                    {%- endif %}
                    topic_params = self._extract_topic_params(mqtt_message.topic, "{{ messageid }}")
                    await self.{{messagename.split('.')[-1]}}_async(mqtt_message, cloud_event, data, topic_params)
                except Exception as e:
                    print(f"Error in {{messagename.split('.')[-1]}} handler: {e}")
            return
        {% endfor %}
    
    async def subscribe(self, topics: Optional[typing.List[str]] = None, qos: int = 0):
        """
        Subscribe to MQTT topics.

        Args:
            topics: List of topic patterns to subscribe to. If None, subscribes to all 
                default topics with URI template placeholders replaced by MQTT wildcards.
            qos: Quality of Service level (0, 1, or 2)
        """
        if topics is None:
            topics = get_subscription_topics_{{ groupName }}(self._topic_mappings)
        for topic in topics:
            self.client.subscribe(topic, qos)
    
    async def unsubscribe(self, topics: typing.List[str]):
        """
        Unsubscribe from MQTT topics.

        Args:
            topics: List of topics to unsubscribe from
        """
        for topic in topics:
            self.client.unsubscribe(topic)
    
    async def connect(self, broker: str, port: int = 1883, keepalive: int = 60):
        """
        Connect to MQTT broker.

        Args:
            broker: Broker hostname or IP
            port: Broker port
            keepalive: Keepalive interval in seconds
        """
        self.client.connect(broker, port, keepalive)
        self.client.loop_start()
    
    async def disconnect(self):
        """Disconnect from MQTT broker."""
        self.client.loop_stop()
        self.client.disconnect()

    # Producer methods
    {% for messageid, message in messagegroup.messages.items() if (message | exists("envelope","CloudEvents/1.0")) -%}
    {%- set messagename = messageid | dotunderscore | snake %}
    {%- set type_name = util.DeclareDataType( data_project_name, root, message ) %}
    {%- set uri_template_vars = [] %}
    {%- for attrname, attribute in message.envelopemetadata.items() %}
        {%- if attribute.value and attribute.type == "uritemplate" %}
            {%- set phs = attribute.value | regex_search('\\{([A-Za-z0-9_]+)\\}') %}
            {%- if phs %}
                {%- for placeholder in phs %}
                    {%- if placeholder not in uri_template_vars %}
                        {%- set _ = uri_template_vars.append(placeholder) %}
                    {%- endif %}
                {%- endfor %}
            {%- endif %}
        {%- endif %}
    {%- endfor %}
    async def publish_{{messagename.split('.')[-1]}}(self,
        {%- for attrname, attribute in message.envelopemetadata.items() if not attribute.value -%}
        {%- if attribute.type != "uritemplate" -%}
        _{{ attrname | snake }}: str,
        {%- endif -%}
        {%- endfor %}
        {%- for var in uri_template_vars %}
        {{ var | snake }}: str,
        {%- endfor %}
        data: {{type_name}},
        topic: Optional[str] = None,
        content_type: str = "application/json") -> None:
        """
        Publish the '{{ messageid }}' event to an MQTT topic.

        Args:
        {% for attrname, attribute in message.envelopemetadata.items() if not attribute.value -%}
        {%- if attribute.type != "uritemplate" -%}
            _{{ attrname | snake }}: CloudEvents attribute '{{attrname}}'
        {%- endif -%}
        {%- endfor %}
        {%- for var in uri_template_vars %}
            {{ var | snake }}: URI template variable for '{{var}}'
        {%- endfor %}
            data: The event data to be published
            topic: Optional topic override. If not provided, uses default topic '{{ get_default_topic(message, messageid) }}'.
            content_type: The content type for the event data
        """
        target_topic = topic if topic is not None else "{{ get_default_topic(message, messageid) }}"
        
        attributes = {
        {%- 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
        byte_data = data.to_byte_array(content_type) if data is not None else b''
        event = CloudEvent(attributes, byte_data)
        
        if self.content_mode == "structured":
            headers, body = to_structured(event)
            payload = body
        else:
            headers, body = to_binary(event)
            payload = body
            # Binary mode headers would go into MQTT 5.0 user properties
            # For MQTT 3.1.1, we'll use structured mode
        
        self.client.publish(target_topic, payload, qos=1)

    {% endfor %}
{% endfor %}
