"""
DO NOT MODIFY THIS FILE. {# Unless you see this, which means you are editing the template. #} It is automatically generated and changes will be over-written
on the next generation.

This is the Client for the {{stinger.name}} interface.
"""

from typing import Dict, Callable, List, Any, Optional
from uuid import uuid4
from functools import partial
import json
import logging
from datetime import datetime, timedelta, UTC
from isodate import parse_duration
{%if stinger.methods | length > 0 %}
import asyncio
import concurrent.futures as futures
from .method_codes import *
from .interface_types import *
import threading
{%endif%}
from .connection import IBrokerConnection
{%if stinger.uses_enums()%}from . import {{stinger.get_enum_module_name()}} as {{stinger.get_enum_module_alias()}}{%endif%}

logging.basicConfig(level=logging.DEBUG)

{%for sig_name, sig in stinger.signals.items()-%}
{{sig_name | UpperCamelCase}}SignalCallbackType = Callable[[{%for arg in sig.arg_list%}{{arg.python_annotation}}{%if not loop.last%}, {%endif%}{%endfor%}], None]
{%endfor%}
{%-for method_name, method in stinger.methods.items()-%}
{{method_name | UpperCamelCase}}MethodResponseCallbackType = Callable[[{%if method.return_value%}{{method.return_value_python_annotation}}{%endif%}], None]
{%endfor%}
{%for prop_name, prop in stinger.properties.items()-%}
{{prop_name | UpperCamelCase}}PropertyUpdatedCallbackType = Callable[[{{prop.python_annotation}}], None]
{%endfor%}

class {{stinger.python.client_class_name}}:

    def __init__(self, connection: IBrokerConnection, service_instance_id: str):
        """ Constructor for a `{{stinger.python.client_class_name}}` object.
        """
        self._logger = logging.getLogger('{{stinger.python.client_class_name}}')
        self._logger.setLevel(logging.DEBUG)
        self._logger.debug("Initializing {{stinger.python.client_class_name}}")
        self._conn = connection
        self._conn.add_message_callback(self._receive_message)
        self._service_id = service_instance_id
        {%if stinger.methods | length > 0 %}
        self._pending_method_responses: dict[str, Callable[..., None]] = {}
        {%endif%}
        {%for prop_name, prop_spec in stinger.properties.items()-%}
        self._property_{{prop_name|snake_case}} = None # type: Optional[{{prop_spec.python_class}}]
        self._property_{{prop_name|snake_case}}_mutex = threading.Lock()
        self._property_{{prop_name|snake_case}}_version = -1
        self._conn.subscribe("{{prop_spec.value_topic}}".format(self._service_id), self._receive_{{prop_name | snake_case}}_property_update_message)
        self._changed_value_callbacks_for_{{prop_name}}: list[{{prop_name | UpperCamelCase}}PropertyUpdatedCallbackType] = []
        {%endfor-%}
        {%-for sig_name in stinger.signals.keys()-%}
        self._signal_recv_callbacks_for_{{sig_name | snake_case }}: list[{{sig_name | UpperCamelCase}}SignalCallbackType] = []
        {%endfor-%}
        {%-for method_name, method in stinger.methods.items()-%}
        self._conn.subscribe(f"{{method.response_topic('{self._conn.client_id}')}}", self._receive_{{method_name | snake_case}}_response_message)
        {%endfor%}

    {%for prop_name, prop_spec in stinger.properties.items()-%}
    @property
    def {{prop_name}}(self) -> Optional[{{prop_spec.python_annotation}}]:
        """ Property '{{prop_name}}' getter.
        """
        return self._property_{{prop_name}}
    {%if not prop_spec.read_only %}
    @{{prop_name}}.setter
    def {{prop_name}}(self, value: {{prop_spec.python_annotation}}):
        """ Serializes and publishes the '{{prop_name}}' property.
        """
        if not isinstance(value, {{prop_spec.python_local_type}}):
            raise ValueError("The '{{prop_name}}' property must be a {{prop_spec.python_class}}")
        {%if prop_spec.arg_list | length > 1 -%}
        serialized = value.model_dump_json(exclude_none=True)
        {%else-%}
        serialized = json.dumps({ "{{prop_spec.arg_list[0].name}}": value.{{prop_spec.arg_list[0].name}} })
        {%endif-%}
        self._logger.debug("Setting '{{prop_name}}' property to %s", serialized)
        self._conn.publish("{{prop_spec.update_topic}}".format(self._service_id), serialized, qos=1)
    {%endif%}
    def {{prop_name | snake_case}}_changed(self, handler: {{prop_name | UpperCamelCase}}PropertyUpdatedCallbackType, call_immediately: bool=False):
        """ Sets a callback to be called when the '{{prop_name}}' property changes.
        Can be used as a decorator.
        """
        with self._property_{{prop_name|snake_case}}_mutex:
            self._changed_value_callbacks_for_{{prop_name|snake_case}}.append(handler)
            if call_immediately and self._property_{{prop_name|snake_case}} is not None:
                handler(self._property_{{prop_name|snake_case}})
        return handler
    {%endfor%}

    def _do_callbacks_for(self, callbacks: List[Callable[..., None]], **kwargs):
        """ Call each callback in the callback dictionary with the provided args.
        """
        for cb in callbacks:
            cb(**kwargs)

    @staticmethod
    def _filter_for_args(args: Dict[str, Any], allowed_args: List[str]) -> Dict[str, Any]:
        """ Given a dictionary, reduce the dictionary so that it only has keys in the allowed list.
        """
        filtered_args = {}
        for k, v in args.items():
            if k in allowed_args:
                filtered_args[k] = v
        return filtered_args

    {%for sig_name, sig in stinger.signals.items()-%}
    def _receive_{{sig_name|snake_case}}_signal_message(self, topic: str, payload: str, properties: Dict[str, Any]):
        if 'ContentType' not in properties or properties['ContentType'] != 'application/json':
            self._logger.warning("Received '{{sig_name}}' signal with non-JSON content type")
            return

        model = {{sig_name|UpperCamelCase}}SignalPayload.model_validate_json(payload)
        kwargs = model.model_dump()

        self._do_callbacks_for(self._signal_recv_callbacks_for_{{sig_name | snake_case}}, **kwargs)
    {%endfor%}{# end signal loop #}

    {%-for method_name, method in stinger.methods.items()%}
    def _receive_{{method_name | snake_case}}_response_message(self, topic: str, payload: str, properties: Dict[str, Any]):
        # Handle '{{method_name}}' method response.
        return_code = MethodReturnCode.SUCCESS
        debug_message = None
        if "UserProperty" in properties:
            user_properties = properties["UserProperty"]
            if "DebugInfo" in user_properties:
                self._logger.info("Received Debug Info to '%s': %s", topic, user_properties["DebugInfo"])
                debug_message = user_properties["DebugInfo"]
            if "ReturnCode" in user_properties:
                return_code = MethodReturnCode(int(user_properties["ReturnCode"]))
        if "CorrelationData" in properties:
            correlation_id = properties["CorrelationData"].decode()
            if correlation_id in self._pending_method_responses:
                cb = self._pending_method_responses[correlation_id]
                del self._pending_method_responses[correlation_id]
                cb(payload, return_code, debug_message)
            else:
                self._logger.warning("Correlation id %s was not in the list of pending method responses... %s", correlation_id, [k for k in self._pending_method_responses.keys()])
        else:
            self._logger.warning("No correlation data in properties sent to %s... %s", topic, [s for s in properties.keys()])
    {%endfor%}{# end method loop #}

    {%-for prop_name, prop_spec in stinger.properties.items()%}
    def _receive_{{prop_name | snake_case}}_property_update_message(self, topic: str, payload: str, properties: Dict[str, Any]):
        # Handle '{{prop_name}}' property change.
        if 'ContentType' not in properties or properties['ContentType'] != 'application/json':
            self._logger.warning("Received '{{prop_name}}' property change with non-JSON content type")
            return
        try:
            prop_obj = {{prop_name|UpperCamelCase}}Property.model_validate_json(payload)
            with self._property_{{prop_name|snake_case}}_mutex:
                self._property_{{prop_name|snake_case}} = prop_obj
                if ver := properties.get("PropertyVersion", False):
                    if int(ver) > self._property_{{prop_name|snake_case}}_version:
                        self._property_{{prop_name|snake_case}}_version = int(ver)
                {% if prop_spec.arg_list | length == 1 %}
                self._do_callbacks_for(self._changed_value_callbacks_for_{{prop_name|snake_case}}, value=prop_obj.{{prop_spec.arg_list[0].name|snake_case}})
                {%else%}
                self._do_callbacks_for(self._changed_value_callbacks_for_{{prop_name|snake_case}}, value=prop_obj)
                {%endif%}
        except Exception as e:
            self._logger.exception("Error processing '{{prop_name}}' property change: %s", exc_info=e)
    {%endfor%}{# end property loop #}

    def _receive_message(self, topic: str, payload: str, properties: Dict[str, Any]):
        """ New MQTT messages are passed to this method, which, based on the topic,
        calls the appropriate handler method for the message.
        """
        self._logger.warning("Receiving message sent to %s, but without a handler", topic)

    {%for sig_name, sig in stinger.signals.items()%}
    def receive_{{sig_name | snake_case }}(self, handler: {{sig_name | UpperCamelCase}}SignalCallbackType):
        """ Used as a decorator for methods which handle particular signals.
        """
        self._signal_recv_callbacks_for_{{sig_name | snake_case }}.append(handler)
        if len(self._signal_recv_callbacks_for_{{sig_name | snake_case}}) == 1:
            self._conn.subscribe("{{sig.topic}}".format(self._service_id), self._receive_{{sig_name|snake_case}}_signal_message)
        return handler
    {%endfor%}

    {%for method_name, method in stinger.methods.items()%}
    def {{method_name | snake_case}}(self, {%for arg in method.arg_list%}{{arg.name|snake_case}}: {{arg.python_annotation}}{%if not loop.last%}, {%endif%}{%endfor%}) -> futures.Future:
        """ Calling this initiates a `{{method_name}}` IPC method call.
        """
        fut = futures.Future() # type: futures.Future
        correlation_id = str(uuid4())
        self._pending_method_responses[correlation_id] = partial(self._handle_{{method_name | snake_case}}_response, fut)
        payload = {{method_name|UpperCamelCase}}MethodRequest( {%-for arg in method.arg_list%}
            {{arg.name}}={{arg.name|snake_case}},{%endfor%}
        )
        json_payload = payload.model_dump_json()
        self._logger.debug("Calling '{{method_name}}' method with payload %s", json_payload)
        self._conn.publish("{{method.topic}}".format(self._service_id), payload.model_dump_json(), qos=2, retain=False,
                           correlation_id=correlation_id, response_topic=f"{{method.response_topic('{self._conn.client_id}')}}")
        return fut

    def _handle_{{method_name | snake_case}}_response(self, fut: futures.Future, response_json_text: str, return_value: MethodReturnCode, debug_message: Optional[str]=None):
        """ This called with the response to a `{{method_name}}` IPC method call.
        """
        self._logger.debug("Handling {{method_name | snake_case}} response message %s", fut)

        if return_value != MethodReturnCode.SUCCESS.value:
            self._logger.warning("Received error return value %s from '{{method_name}}' method: %s", return_value, debug_message)
            fut.set_exception(stinger_exception_factory(return_value, debug_message))
            return

        try:
            resp_model = {{method_name|UpperCamelCase}}MethodResponse.model_validate_json(response_json_text)
        except Exception as e:
            fut.set_exception(ClientDeserializationErrorStingerMethodException(f"Failed to deserialize response to '{{method_name}}' method: {e}"))

        if not fut.done():
        {%-if method.return_arg_list | length == 0 %}
            fut.set_result(None)
            return
        {%elif method.return_arg_list | length == 1 %}
            fut.set_result(resp_model.{{method.return_arg_list[0].name|snake_case}})
        {%else%}
            fut.set_result(resp_model)
        {%endif-%}
        else:
            self._logger.warning("Future for '{{method_name}}' method was already done!")
    {%endfor%}

class {{stinger.python.client_class_name}}Builder:
    """ Using decorators from {{stinger.python.client_class_name}} doesn't work if you are trying to create multiple instances of {{stinger.python.client_class_name}}.
    Instead, use this builder to create a registry of callbacks, and then build clients using the registry.

    When ready to create a {{stinger.python.client_class_name}} instance, call the `build(broker, service_instance_id)` method.
    """

    def __init__(self):
        """ Creates a new {{stinger.python.client_class_name}}Builder.
        """
        self._logger = logging.getLogger('{{stinger.python.client_class_name}}Builder')
        {%for sig_name in stinger.signals.keys()-%}
        self._signal_recv_callbacks_for_{{sig_name | snake_case}} = [] # type: List[{{sig_name | UpperCamelCase}}SignalCallbackType]
        {%endfor-%}
        {%for prop_name in stinger.properties.keys()-%}
        self._property_updated_callbacks_for_{{prop_name | snake_case}}: list[{{prop_name | UpperCamelCase}}PropertyUpdatedCallbackType] = []
        {%endfor-%}

    {%for sig_name, sig in stinger.signals.items()%}
    def receive_{{sig_name | snake_case }}(self, handler):
        """ Used as a decorator for methods which handle particular signals.
        """
        self._signal_recv_callbacks_for_{{sig_name | snake_case }}.append(handler)
    {%endfor%}

    {%for prop_name, prop in stinger.properties.items()%}
    def {{prop_name|snake_case}}_updated(self, handler: {{prop_name | UpperCamelCase}}PropertyUpdatedCallbackType):
        """ Used as a decorator for methods which handle updates to properties.
        """
        self._property_updated_callbacks_for_{{prop_name | snake_case}}.append(handler)
    {%endfor%}

    def build(self, broker: IBrokerConnection, service_instance_id: str) -> {{stinger.python.client_class_name}}:
        """ Builds a new {{stinger.python.client_class_name}}.
        """
        self._logger.debug("Building {{stinger.python.client_class_name}} for service instance %s", service_instance_id)
        client = {{stinger.python.client_class_name}}(broker, service_instance_id)
        {%for sig_name, sig in stinger.signals.items()%}
        for cb in self._signal_recv_callbacks_for_{{sig_name | snake_case }}:
            client.receive_{{sig_name | snake_case}}(cb)
        {%endfor%}
        {%for prop_name, prop in stinger.properties.items()%}
        for cb in self._property_updated_callbacks_for_{{prop_name | snake_case}}:
            client.{{prop_name | snake_case}}_changed(cb)
        {%endfor%}
        return client


class {{stinger.python.client_class_name}}Discoverer:
    
    def __init__(self, connection: IBrokerConnection, builder: Optional[{{stinger.python.client_class_name}}Builder]=None):
        """ Creates a new {{stinger.python.client_class_name}}Discoverer.
        """
        self._conn = connection
        self._builder = builder
        self._logger = logging.getLogger('{{stinger.python.client_class_name}}Discoverer')
        self._logger.setLevel(logging.DEBUG)
        service_discovery_topic = "{{stinger.interface_info_topic}}".format('+')
        self._conn.subscribe(service_discovery_topic, self._process_service_discovery_message)
        self._mutex = threading.Lock()
        self._discovered_services: Dict[str, InterfaceInfo] = {}
        self._discovered_service_callbacks: List[Callable[[InterfaceInfo], None]] = []
        self._pending_futures : List[futures.Future] = []
        self._removed_service_callbacks: List[Callable[[str], None]] = []

    def add_discovered_service_callback(self, callback: Callable[[InterfaceInfo], None]):
        """ Adds a callback to be called when a new service is discovered.
        """
        with self._mutex:
            self._discovered_service_callbacks.append(callback)

    def add_removed_service_callback(self, callback: Callable[[str], None]):
        """ Adds a callback to be called when a service is removed.
        """
        with self._mutex:
            self._removed_service_callbacks.append(callback)

    def get_service_instance_ids(self) -> List[str]:
        """ Returns a list of currently discovered service instance IDs.
        """
        with self._mutex:
            return list(self._discovered_services.keys())

    def get_singleton_client(self) -> futures.Future[{{stinger.python.client_class_name}}]:
        """ Returns a {{stinger.python.client_class_name}} for the single discovered service.
        Raises an exception if there is not exactly one discovered service.
        """
        fut = futures.Future()
        with self._mutex:
            if len(self._discovered_services) > 0:
                service_instance_id = next(iter(self._discovered_services))
                if self._builder is None:
                    fut.set_result({{stinger.python.client_class_name}}(self._conn, service_instance_id))
                else:
                    new_client = self._builder.build(self._conn, service_instance_id)
                    fut.set_result(new_client)
            else:
                self._pending_futures.append(fut)
        return fut

    def _process_service_discovery_message(self, topic: str, payload: str, properties: Dict[str, Any]):
        """ Processes a service discovery message.
        """
        self._logger.debug("Processing service discovery message on topic %s", topic)
        if len(payload) > 0:
            try:
                service_info = InterfaceInfo.model_validate_json(payload)
            except Exception as e:
                self._logger.warning("Failed to process service discovery message: %s", e)
            with self._mutex:
                self._discovered_services[service_info.instance] = service_info
                while self._pending_futures:
                    fut = self._pending_futures.pop(0)
                    if not fut.done():
                        if self._builder is not None:
                            fut.set_result(self._builder.build(self._conn, service_info.instance))
                        else:
                            fut.set_result({{stinger.python.client_class_name}}(self._conn, service_info.instance))
                if not service_info.instance in self._discovered_services:
                    self._logger.info("Discovered service: %s.instance", service_info.instance)
                    for cb in self._discovered_service_callbacks:
                        cb(service_info)
                else:
                    self._logger.debug("Updated info for service: %s", service_info.instance)
        else: # Empty payload means the service is going away
            instance_id = topic.split('/')[-2]
            with self._mutex:
                if instance_id in self._discovered_services:
                    self._logger.info("Service %s is going away", instance_id)
                    del self._discovered_services[instance_id]
                    for cb in self._removed_service_callbacks:
                            cb(instance_id)

