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

# pylint: disable=line-too-long

"""
Dispatcher module for processing messages from AMQP 1.0 broker.

This module provides the core message dispatching functionality for AMQP 1.0 consumers.
It handles both CloudEvents-formatted messages and plain AMQP messages, automatically
routing them to registered handler functions based on message type or subject.

The dispatcher supports:
- CloudEvents structured and binary content modes
- Plain AMQP messages with subject-based routing
- Message settlement (accept, reject, release)
- Automatic deserialization of message payloads
- Type-safe handler registration

Example:
    >>> consumer = FabrikamInkjetEventsEventConsumer()
    >>> consumer.printer_ready_handler = handle_printer_ready
    >>> await consumer.process_message(message, delivery)
"""

import typing
import json
import logging
from proton import Message
from cloudevents.http import CloudEvent, from_json

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

logger = logging.getLogger(__name__)


class MessageContext:
    """
    Context object for AMQP message processing.
    
    Provides methods to accept, reject, or release a message back to the broker.
    """
    
    def __init__(self, delivery):
        """
        Initialize the message context.
        
        Args:
            delivery: The proton delivery object for message settlement
        """
        self._delivery = delivery
        self._settled = False
    
    def accept(self) -> None:
        """Accept the message (successful processing)"""
        if not self._settled and self._delivery:
            self._delivery.update(self._delivery.ACCEPTED)
            self._delivery.settle()
            self._settled = True
    
    def reject(self) -> None:
        """Reject the message (permanent failure, don't retry)"""
        if not self._settled and self._delivery:
            self._delivery.update(self._delivery.REJECTED)
            self._delivery.settle()
            self._settled = True
    
    def release(self) -> None:
        """Release the message (temporary failure, allow retry)"""
        if not self._settled and self._delivery:
            self._delivery.update(self._delivery.RELEASED)
            self._delivery.settle()
            self._settled = True
    
    @property
    def is_settled(self) -> bool:
        """Check if the message has been settled"""
        return self._settled


{%- for messagegroupid, messagegroup in messagegroups.items() -%}
{%- set uses_cloudevents_message = (messagegroup | exists("envelope","CloudEvents/1.0")) %}
{%- set uses_amqp_message = true %}
{%- set groupname = messagegroupid | pascal %}
{%- set class_name = ( groupname | strip_dots )+"EventConsumer" %}

class {{ class_name }}:
    """
    Event consumer for the {{ messagegroupid }} message group.
    
    This class receives and dispatches AMQP 1.0 messages, supporting both CloudEvents
    envelope format and plain AMQP messages. It provides type-safe message routing
    to registered handler functions.
    
    The consumer automatically detects message format (CloudEvents vs. plain AMQP),
    deserializes message payloads based on content type, and routes messages to
    the appropriate handler based on CloudEvent type or AMQP subject.
    
    Attributes:
        {%- for messageid, message in messagegroup.messages.items() %}
        {%- set messagename = messageid | strip_dots | pascal %}
        {{ messagename | snake }}_handler: Optional handler for {{ messageid }} messages
        {%- endfor %}
        {%- if uses_amqp_message %}
        unhandled_message_handler: Optional handler for unrecognized plain AMQP messages
        {%- endif %}
        {%- if uses_cloudevents_message %}
        unhandled_event_handler: Optional handler for unrecognized CloudEvents
        {%- endif %}
    
    Example:
        >>> consumer = {{ class_name }}()
        {%- set first_message = messagegroup.messages.items() | first %}
        {%- if first_message %}
        {%- set messageid, message = first_message %}
        {%- set messagename = messageid | strip_dots | pascal %}
        {%- set isCloudEvent = (message | exists("envelope","CloudEvents/1.0")) %}
        >>> async def handler({% if isCloudEvent %}cloud_event, {% endif %}data, context):
        ...     print(f"Received: {data}")
        ...     context.accept()
        >>> consumer.{{ messagename | snake }}_handler = handler
        >>> await consumer.process_message(message, delivery)
        {%- endif %}
    """
    
    def __init__(self):
        """
        Initialize the event consumer.
        
        Creates a new consumer instance with all message handlers initialized to None.
        Handlers must be registered before calling process_message().
        """
        {%- for messageid, message in messagegroup.messages.items() %}
        {%- set messagename = messageid | strip_dots | pascal %}
        {%- set message_body_type = util.DeclareDataType( data_project_name, root, message ) | strip_namespace %}
        {%- set isCloudEvent = (message | exists("envelope","CloudEvents/1.0")) %}
        self.{{ messagename | snake }}_handler: typing.Optional[typing.Callable[[{% if isCloudEvent %}CloudEvent, {% else %}Message, {% endif %}{{ message_body_type }}, MessageContext], typing.Awaitable[None]]] = None
        {%- endfor %}
        
        {%- if uses_amqp_message %}
        self.unhandled_message_handler: typing.Optional[typing.Callable[[Message, MessageContext], typing.Awaitable[None]]] = None
        {%- endif %}
        
        {%- if uses_cloudevents_message %}
        self.unhandled_event_handler: typing.Optional[typing.Callable[[CloudEvent, MessageContext], typing.Awaitable[None]]] = None
        {%- endif %}
    
    async def process_message(self, message: Message, delivery=None) -> None:
        """
        Process an AMQP message and route it to the appropriate handler.
        
        This method is the main entry point for message processing. It automatically
        detects the message format (CloudEvents vs. plain AMQP), creates a MessageContext
        for settlement, and dispatches to the appropriate handler based on message type.
        
        The method performs the following steps:
        1. Creates a MessageContext wrapping the delivery for settlement operations
        2. Checks if the message is a CloudEvent (structured or binary mode)
        3. Converts the message to the appropriate format (CloudEvent or Message)
        4. Dispatches to the registered handler for the message type/subject
        5. Calls unhandled_*_handler if no specific handler is registered
        
        Args:
            message: The AMQP Proton Message object to process
            delivery: The Proton Delivery object for message settlement. Optional for
                testing scenarios where settlement is not required. If provided, handlers
                can use context.accept(), context.reject(), or context.release() to
                settle the message.
        
        Raises:
            Exception: Any exception raised by the handler is propagated to the caller.
                The caller is responsible for catching exceptions and settling the
                message appropriately (e.g., reject on error, release for retry).
        
        Example:
            >>> async def handle_message(cloud_event, data, context):
            ...     try:
            ...         await process_data(data)
            ...         context.accept()
            ...     except TransientError:
            ...         context.release()  # Allow retry
            ...     except PermanentError:
            ...         context.reject()   # Don't retry
            >>> consumer.message_handler = handle_message
            >>> await consumer.process_message(proton_message, delivery)
        """
        context = MessageContext(delivery)
        
        {%- if uses_cloudevents_message %}
        if self._is_cloud_event(message):
            cloud_event = self._cloud_event_from_message(message)
            await self._dispatch_cloud_event(cloud_event, context)
            return
        {%- endif %}
        
        await self._dispatch_message(message, context)
    
    {%- if uses_cloudevents_message %}
    
    def _is_cloud_event(self, message: Message) -> bool:
        """
        Check if an AMQP message is a CloudEvent.
        
        Args:
            message (Message): The AMQP message to check
            
        Returns:
            bool: True if the message is a CloudEvent, False otherwise
        """
        # Check for structured mode CloudEvent
        content_type = message.content_type
        if content_type and (content_type.startswith('application/cloudevents+json') or 
                            content_type.startswith('application/cloudevents')):
            return True
        
        # Check for binary mode CloudEvent
        if message.properties:
            app_props = message.properties
            if 'cloudEvents:specversion' in app_props or 'ce_specversion' in app_props:
                return True
        
        return False
    
    def _cloud_event_from_message(self, message: Message) -> CloudEvent:
        """
        Convert an AMQP message to a CloudEvent object.
        
        Handles both CloudEvents content modes:
        - Structured mode: Entire CloudEvent as JSON in message body, identified by
          content type 'application/cloudevents+json'
        - Binary mode: CloudEvents attributes in AMQP application properties with
          'cloudEvents:', 'ce_', or 'ce-' prefixes, event data in message body
        
        Args:
            message: The AMQP Proton Message object to convert
            
        Returns:
            A CloudEvent object containing the event attributes and data
            
        Raises:
            ValueError: If the message cannot be parsed as a valid CloudEvent
            json.JSONDecodeError: If structured mode JSON is malformed
        
        Example:
            Structured mode message will have content type 'application/cloudevents+json'
            and body containing: {"specversion": "1.0", "type": "...", ...}
            
            Binary mode message will have application properties like:
            {'ce_specversion': '1.0', 'ce_type': '...', 'ce_source': '...'}
        """
        content_type = message.content_type
        
        # Structured mode
        if content_type and content_type.startswith('application/cloudevents'):
            body = message.body
            if isinstance(body, bytes):
                body = body.decode('utf-8')
            if isinstance(body, str):
                return from_json(body)
            elif isinstance(body, dict):
                return CloudEvent(body)
        
        # Binary mode
        app_props = message.properties or {}
        
        # Handle different property name formats
        def get_prop(name: str) -> typing.Optional[str]:
            """Get property value from application properties with multiple name formats"""
            for prefix in ['cloudEvents:', 'cloudEvents_', 'ce_', 'ce-']:
                key = prefix + name
                if key in app_props:
                    val = app_props[key]
                    return val.decode('utf-8') if isinstance(val, bytes) else str(val)
            return None
        
        attributes = {
            'specversion': get_prop('specversion') or '1.0',
            'type': get_prop('type'),
            'source': get_prop('source'),
            'id': get_prop('id'),
        }
        
        # Add optional attributes
        time_val = get_prop('time')
        if time_val:
            attributes['time'] = time_val
        
        subject_val = get_prop('subject')
        if subject_val:
            attributes['subject'] = subject_val
        
        # Get data and content type
        data = message.body
        attributes['datacontenttype'] = content_type or 'application/json'
        
        return CloudEvent(attributes, data)
    
    async def _dispatch_cloud_event(self, cloud_event: CloudEvent, context: MessageContext) -> None:
        """
        Dispatch a CloudEvent to the appropriate registered handler.
        
        Routes the CloudEvent to the handler registered for its 'type' attribute.
        If no specific handler is registered, calls unhandled_event_handler if available.
        Automatically deserializes the event data using the generated from_data() method
        of the message data class.
        
        Args:
            cloud_event: The CloudEvent object to dispatch
            context: The MessageContext for message settlement operations
            
        Note:
            The handler signature is:
            async def handler(cloud_event: CloudEvent, data: DataType, context: MessageContext)
            
            The handler receives:
            - The complete CloudEvent for accessing metadata (id, source, time, etc.)
            - The deserialized data object with strong typing
            - The MessageContext for settlement (accept/reject/release)
            
            If the handler does not settle the message, it remains unacknowledged.
        """
        cloud_event_type = cloud_event['type']
        
        {% for messageid, message in messagegroup.messages.items() -%}
        {%- set messagename = messageid | strip_dots | pascal %}
        {%- set isCloudEvent = (message | exists("envelope","CloudEvents/1.0")) %}
        {%- set message_body_type = util.DeclareDataType( data_project_name, root, message ) | strip_namespace %}
        {%- if isCloudEvent %}
        if cloud_event_type == '{{ messageid }}':
            if self.{{ messagename | snake }}_handler:
                {%- if message_body_type == 'object' %}
                # No data schema defined for this message
                await self.{{ messagename | snake }}_handler(cloud_event, None, context)
                {%- else %}
                body_data = cloud_event.data
                content_type = cloud_event.get('datacontenttype', 'application/json')
                
                # Convert memoryview to bytes if needed
                if isinstance(body_data, memoryview):
                    body_data = body_data.tobytes()
                
                # Deserialize using the generated from_data method
                data = {{ message_body_type }}.from_data(body_data, content_type)
                await self.{{ messagename | snake }}_handler(cloud_event, data, context)
                {%- endif %}
                return
        
        {% endif %}
        {%- endfor %}
        # No handler found
        if self.unhandled_event_handler:
            await self.unhandled_event_handler(cloud_event, context)
        else:
            logger.warning(f"No handler for CloudEvent type: {cloud_event_type}")
    {%- endif %}
    
    {%- if uses_amqp_message %}
    
    async def _dispatch_message(self, message: Message, context: MessageContext) -> None:
        """
        Dispatch a plain AMQP message to the appropriate registered handler.
        
        Routes the message to the handler registered for its 'subject' application property.
        The subject is extracted from application properties using either 'subject' or
        'Subject' keys. If no specific handler is registered, calls unhandled_message_handler
        if available. Automatically deserializes the message body using the generated
        from_data() method of the message data class.
        
        Args:
            message: The AMQP Proton Message object to dispatch
            context: The MessageContext for message settlement operations
            
        Note:
            The handler signature is:
            async def handler(message: Message, data: DataType, context: MessageContext)
            
            The handler receives:
            - The complete AMQP Message for accessing properties and metadata
            - The deserialized data object with strong typing
            - The MessageContext for settlement (accept/reject/release)
            
            If the handler does not settle the message, it remains unacknowledged.
        """
        # Get subject from application properties
        message_subject = None
        if message.properties:
            app_props = message.properties
            for key in ['subject', 'Subject']:
                if key in app_props:
                    val = app_props[key]
                    message_subject = val.decode('utf-8') if isinstance(val, bytes) else str(val)
                    break
        
        {% for messageid, message in messagegroup.messages.items() if not (message | exists("envelope","CloudEvents/1.0")) -%}
        {%- set messagename = messageid | strip_dots | pascal %}
        {%- set message_body_type = util.DeclareDataType( data_project_name, root, message ) | strip_namespace %}
        if message_subject == '{{ messageid }}':
            if self.{{ messagename | snake }}_handler:
                {%- if message_body_type == 'object' %}
                # No data schema defined for this message
                await self.{{ messagename | snake }}_handler(message, None, context)
                {%- else %}
                body_data = message.body
                content_type = message.content_type or 'application/json'
                
                # Convert memoryview to bytes if needed
                if isinstance(body_data, memoryview):
                    body_data = body_data.tobytes()
                
                # Deserialize using the generated from_data method
                data = {{ message_body_type }}.from_data(body_data, content_type)
                await self.{{ messagename | snake }}_handler(message, data, context)
                {%- endif %}
                return
        
        {% endfor %}
        # No handler found
        if self.unhandled_message_handler:
            await self.unhandled_message_handler(message, context)
        else:
            logger.warning(f"No handler for message subject: {message_subject}")
    {%- endif %}

{% endfor %}
