"""
This is the implementation of the OGN node defined in OgnOnMessageBusEvent.ogn
"""
from contextlib import suppress
from typing import Optional

import carb
import carb.dictionary
import carb.events
import numpy as np
import omni.graph.core as og
import omni.kit.app
from omni.graph.action import get_interface
from omni.graph.action.ogn.OgnOnMessageBusEventDatabase import OgnOnMessageBusEventDatabase


class OgnOnMessageBusEventInternalState:
    """Convenience class for maintaining per-node state information"""

    def __init__(self):
        # This subscription object controls the lifetime of our callback, it will be
        # cleaned up automatically when our node is destroyed
        self.sub: carb.events.ISubscription = None
        # Set when the callback has triggered
        self.is_set = False
        # The last payload received
        self.payload: carb.dictionary.Item = None
        # The event_name we used to subscribe
        self.sub_event_name = ""
        # The node instance handle
        self.node: og.Node = None

    def on_event(self, custom_event: carb.events.IEvent):
        if custom_event is None:
            return
        self.is_set = True
        self.payload = custom_event.payload
        # Tell the evaluator we need to be computed
        if self.node.is_valid():
            self.node.request_compute()

    def first_time_subscribe(self, node: og.Node, event_name: str) -> bool:
        """Checked call to set up carb subscription
        Args:
            node: The node instance
            event_name: The name of the carb event
        Returns:
            True if we subscribed, False if we are already subscribed
        """
        if self.sub is not None and self.sub_event_name != event_name:
            # event name changed since we last subscribed, unsubscribe
            self.sub.unsubscribe()
            self.sub = None

        if self.sub is None:
            # Add a subscription for the given event name. This is a pop subscription, so we expect a 1-frame
            # lag between send and receive
            reg_event_name = carb.events.type_from_string(event_name)
            message_bus = omni.kit.app.get_app().get_message_bus_event_stream()
            self.sub = message_bus.create_subscription_to_pop_by_type(reg_event_name, self.on_event)
            self.sub_event_name = event_name
            self.node = node
            return True

        return False

    def try_pop_event(self) -> Optional[carb.dictionary.Item]:
        """Pop the payload of the last event received, or None if there is no event to pop"""
        if self.is_set:
            self.is_set = False
            payload = self.payload
            self.payload = None
            return payload
        return None


# ======================================================================


class OgnOnMessageBusEvent:
    """
    This node triggers when the specified message bus event is received
    """

    @staticmethod
    def internal_state():
        """Returns an object that will contain per-node state information"""
        return OgnOnMessageBusEventInternalState()

    @staticmethod
    def compute(db) -> bool:
        event_name = db.inputs.eventName
        if not event_name:
            return True

        state = db.internal_state

        if state.first_time_subscribe(db.node, event_name):
            return True

        # Drop events if we are disabled
        if db.inputs.onlyPlayback and (not db.node.get_graph().get_default_graph_context().get_is_playing()):
            state.try_pop_event()
            return True

        payload = state.try_pop_event()

        if payload is None:
            return True

        node: og.Node = db.node

        # Copy the event dict contents into dynamic attributes if they exist
        for key in payload.get_keys():
            attr_name = f"outputs:{key}"
            if not node.get_attribute_exists(attr_name):
                continue
            attrib = node.get_attribute(attr_name)
            if attrib:
                value = payload[key]
                og_type: og.Type = attrib.get_resolved_type()
                is_array = og_type.array_depth > 0
                is_tuple = og_type.tuple_count > 1
                is_matrix = is_tuple and (
                    og_type.role in (og.AttributeRole.FRAME, og.AttributeRole.MATRIX, og.AttributeRole.TRANSFORM)
                )
                if is_array:
                    if is_matrix:
                        dim = 2 if og_type.tuple_count == 4 else 3 if og_type.tuple_count == 9 else 4
                        value = np.array(value).reshape((-1, dim, dim))
                    elif is_tuple:
                        dim = og_type.tuple_count
                        value = np.array(value).reshape((-1, dim))
                elif is_matrix:
                    dim = 2 if og_type.tuple_count == 4 else 3 if og_type.tuple_count == 9 else 4
                    value = np.array(value).reshape((dim, dim))

                try:
                    og.Controller.set(attrib, value)
                except TypeError as exc:
                    db.log_error(f"Failed to copy payload data {key} to attribute {attr_name}:\n{exc}")
                    return False
        get_interface().set_execution_enabled("outputs:execOut")
        return True

    # ----------------------------------------------------------------------------
    @staticmethod
    def release(node):
        # Unsubscribe right away instead of waiting for GC cleanup, we don't want our callback firing
        # after the node has been released.
        with suppress(og.OmniGraphError):
            state = OgnOnMessageBusEventDatabase.per_node_internal_state(node)
            if state.sub:
                state.sub.unsubscribe()
                state.sub = None
