import copy
import functools
import inspect
import json
from collections import OrderedDict
from typing import Callable, Dict, List, _GenericAlias

import carb

from . import type_definitions
from .util import GeneratedCode

EXEC_ATTRIBUTE_PREFIX = "exec"
NODE_DATA_LITERAL = "node_data"

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


class AutoFunctionWrapper(type_definitions.AutoNodeDefinitionWrapper):
    """Class to wrap python functions and classmethods. Contains methods for creating an Ogn annotation from a python
    function, and methods for applying values to the function.
    """

    def __init__(
        self, func: Callable, *, pure: bool, ui_name: str, unique_name: str, tags: List[str], annotation: Dict = None
    ) -> None:
        super().__init__()
        self._func_proto = func
        self._func = func
        self._is_bound = hasattr(func, "__self__")
        self._returns_tuple = False
        self._argument_cache: Dict = {}
        self._has_node_data_signature = False
        self._node_impl = None

        self.python_name = func.__name__ or func.__qualname__
        self.name = ui_name or self.python_name
        self.pure = pure or False
        if unique_name is None:
            raise NameError(f"{func.__qualname__} does not have a specified unique name")
        self.unique_name = unique_name

        self.descriptor = {}
        self.descriptor["uiName"] = self.name
        self.descriptor["version"] = 1
        self.descriptor["language"] = "Python"
        self.descriptor["tags"] = tags or []

        docstring = func.__doc__ or "[no documentation]"
        self.descriptor["description"] = f"{self.name}: {docstring}. Autogenerated"

        self.descriptor["inputs"] = OrderedDict(
            {
                f"_unique_name_{self.unique_name}": {
                    "uiName": "Function Name",
                    "type": "string",
                    "description": "autogenerated qualified name",
                    "metadata": {"hidden": 1},
                    "default": self.unique_name,
                }
            }
        )

        if not self.pure:
            self.descriptor["tags"].append("action")
            self.descriptor["inputs"][EXEC_ATTRIBUTE_PREFIX] = {
                "uiName": "Exec",
                "description": "exec input",
                "type": "execution",
            }

        if annotation is not None:
            # If the annotation shim exists, use it:
            for key in annotation:
                if key == "return":
                    self._set_output_type(annotation[key])
                elif key == "self":
                    # FIXME (OS): func.__class__ will not contain the whole type
                    # in the case pybind is used
                    self._add_input(key, func.__class__)
                else:
                    self._add_input(key, annotation[key])
        else:
            # fall back to inspection otherwise
            sig = inspect.signature(func)
            params = sig.parameters
            for name in params:
                item = params[name]
                if name == NODE_DATA_LITERAL:  # This node holds state
                    self._has_node_data_signature = True
                elif name == "self" and item.annotation is inspect._empty:
                    self._add_input(item.name, func.__class__)
                else:
                    self._add_input(item.name, item.annotation)
            self._set_output_type(sig.return_annotation)

    # --------------------------------------------------------------------------------
    @staticmethod
    def get_type_signature(type_desc: type) -> Dict:
        """Constructs an Ogn type signature from the incoming type. If the type can be converted to a native Ogn type,
        the method attempts to do this automatically.

        Args:
            type_desc: a type object to describe an input type

        Returns:
            A dict object with Ogn properties
        """
        if type_desc is None:
            return None
        conversion = type_definitions.AutoNodeTypeConversion.from_type(type_desc)
        ret = {"description": "autogenerated python type"}
        if conversion is not None:
            ret["type"] = conversion.og_type
            ret["default"] = conversion.default
        else:
            try:
                type_name = type_desc.__name__
            except AttributeError:
                type_name = str(type_desc)
            ret["type"] = "objectId"
            ret["metadata"] = {"python_type_desc": type_name}
        return ret

    # --------------------------------------------------------------------------------
    def _add_input(self, name: str, type_desc: type):
        """Adds an input to the function wrpper

        Args:
            name: a string to name the input in the node's public interface
            type_desc: a type for the function's public interface
        """
        self.descriptor["inputs"][name] = AutoFunctionWrapper.get_type_signature(type_desc)
        self.descriptor["inputs"].move_to_end(name)

    # --------------------------------------------------------------------------------
    def _set_output_type(self, type_desc: type):
        """Sets the function's output signature. Adds an "exec" port if the function isn't marked as "pure"
        If the function returns a 'Tuple' type, the Tuple is unpacked to individual node outputs,
        labeled 'out_0', 'out_1', etc.

        Args:
            type_desc: a return type

        """
        outputs = OrderedDict()
        if not self.pure:
            outputs[EXEC_ATTRIBUTE_PREFIX] = {"type": "execution", "description": "autogenerated output"}
        if type_desc is None:
            self.descriptor["outputs"] = outputs
            return

        # Handle case where tuples are returned
        if isinstance(type_desc, _GenericAlias) and type_desc.__origin__ == tuple:
            self._returns_tuple = True

        if self._returns_tuple:
            for num, tuple_type_desc in enumerate(reversed(type_desc.__args__)):
                name = f"out_{num}"
                outputs[name] = AutoFunctionWrapper.get_type_signature(tuple_type_desc)
                outputs.move_to_end(name)
        else:
            outputs["out_0"] = AutoFunctionWrapper.get_type_signature(type_desc)
            outputs.move_to_end("out_0")

        self.descriptor["outputs"] = outputs

    # --------------------------------------------------------------------------------
    def clear(self):
        """Resets all bindings in the function argument cache."""
        self._argument_cache = {}
        self._func = self._func_proto

    # --------------------------------------------------------------------------------
    @property
    def func(self) -> Callable:
        """Access to the bound function.

        Returns:
            A partially or fully bound function.
        """
        if not self._func:
            self._func = self.class_ref.__getattribute__(self.python_name)  # noqa: PLC2801,PLE1101
        return self._func

    # --------------------------------------------------------------------------------
    @property
    def is_bound(self) -> bool:
        """Whether the underlying function dependes on a specific self parameter"""
        return self._is_bound

    # --------------------------------------------------------------------------------
    @property
    def has_node_data(self) -> bool:
        """Whether this function's signature has a "node data" component"""
        return self._has_node_data_signature

    # --------------------------------------------------------------------------------
    def assign_input(self, name: str, value):
        """Assigns one of the function inputs into the function argument cache.

        Args:
            name: argument name
            value: unwrapped value to be applied.
        """
        self._argument_cache[name] = value
        return self

    # --------------------------------------------------------------------------------
    def assign_input_with_wrapper(self, name: str, value: type_definitions.AutographDataWrapper):
        """Unwraps an object wrapper and applies it to the function argument cache.

        Args:
            name: argument name
            value: object wrapper
        """
        return self.assign_input(name, value.value)

    # --------------------------------------------------------------------------------
    def assign_input_by_id(self, name: str, obj_id: int):
        """Retrieves an input from the object store by its id and applies it to the function argument cache.

        Args:
            name: argument name
            obj_id: the object ID for the argument wrapper, used to retrieve it from the object store

        Raises:
            KeyError in case the argument isn't found.

        """
        try:
            wrapper = type_definitions.TypeRegistry.instance().refs.pop(obj_id)
        except KeyError as e:
            carb.log_error(
                f"""Input assignment error: Referenced input [{self.name}:{name}] not found in object store.
                It may have been consumed: {e}"""
            )
            raise e

        return self.assign_input_with_wrapper(name, wrapper)

    # --------------------------------------------------------------------------------
    def execute(self):
        """Executes the wrapped function and clears the argument cache.

        Returns:
            The function's return value, else None.

        Raises:
            ValueError if not all arguments in the argument cache are proerly filled
        """
        args = {}
        if self.is_bound and "self" in self.descriptor["inputs"]:
            # method needs to provide its own 'self'
            func_name = self._func_proto.__name__
            func_source = self._argument_cache["self"]
            # find the specific method bound to the 'self' we're after
            self._func = getattr(func_source, func_name)
            args.update({k: v for k, v in self._argument_cache.items() if k != "self"})
        else:
            args.update(self._argument_cache)

        if len(args) > 0:
            self._func = functools.partial(self._func, **args)

        ret = self.func.__call__()  # noqa: PLR1705,PLC2801
        self.clear()
        return ret

    # --------------------------------------------------------------------------------
    def execute_and_add_to_graph(self) -> int:
        """Executes the wrapped function and adds it to the object store if it has a return value.
        Does not attempt to convert to an existing Ogn type.
        Clears the underlying function argument cache.

        Returns:
            The object ID if added to the graph, 0 otherwise.
        """
        ret = self.execute()
        if ret is not None:
            return type_definitions.TypeRegistry.add_to_graph(ret)

        return 0

    # --------------------------------------------------------------------------------
    def get_ogn(self) -> str:
        """Gets the Ogn function representation

        Returns:
            A JSON String with the representation
        """
        d = {self.unique_name: self.descriptor}
        return json.dumps(d)

    def get_module_name(self) -> str:
        return self._func.__module__

    def get_node_impl(self) -> type:
        if self._node_impl is None:

            class NodeImpl(
                AutographGenericFunction,
                unique_name=self.get_unique_name(),
                annotation=self.get_ogn(),
                # FIXME (OS): Change to actual type
                autograph_data={} if self.has_node_data else None,
                returns_tuple=self._returns_tuple,
            ):
                pass

            self._node_impl = NodeImpl

        return self._node_impl

    def get_unique_name(self) -> str:
        return self.unique_name


# ================================================================================
class AutographGenericFunction:
    """Ogn Python function runner for single output"""

    def __init_subclass__(cls, unique_name: str, annotation: Dict, autograph_data: Dict, returns_tuple: bool) -> None:
        cls.unique_name = unique_name
        cls.annotation = json.loads(annotation)[unique_name]
        cls.autograph_data = autograph_data
        cls.returns_tuple = returns_tuple
        cls.generate_compute()

    # --------------------------------------------------------------------------------
    @classmethod
    def generate_compute(cls):
        """Generates the compute function for an instntiated class.
        This code generator aims to minimize the size of code that runs at every frame
        """
        code = GeneratedCode()

        code.line("@classmethod")
        with code.indent("def compute(*args) -> bool:"):
            code.line("cls = args[0]")
            code.line("db = args[1]")
            code.line(f"func_obj = type_definitions.TypeRegistry.get_func('{cls.unique_name}')")

            def key_filter(x: str) -> bool:
                ret = not x.startswith("_")  # no private attributes
                ret &= not x.startswith(EXEC_ATTRIBUTE_PREFIX)  # no exec attributes
                return not ret

            for key in cls.annotation["inputs"]:
                if key_filter(key):
                    continue
                typename = cls.annotation["inputs"][key]["type"]
                conversion = type_definitions.AutoNodeTypeConversion.from_ogn_type(typename)
                if conversion:
                    if conversion.og_to_type is not None:
                        code.line(
                            f"value = type_definitions.AutoNodeTypeConversion.from_ogn_type('{typename}')"
                            f".og_to_type(db.inputs.{key})"
                        )
                    else:
                        code.line(f"value = db.inputs.{key}")
                    code.line(f"func_obj.assign_input('{key}',value)")
                else:
                    with code.indent("try:"):
                        code.line(f"func_obj.assign_input_by_id('{key}', db.inputs.{key})")
                    with code.indent("except KeyError as e:"):
                        code.line(
                            f"""carb.log_error("Failed to assign input at [{cls.unique_name} : {key}] : " + str(e))"""
                        )
                        code.line("return False")

            if cls.autograph_data is not None:
                code.line(f"func_obj.assign_input('{NODE_DATA_LITERAL}', db.internal_state)")

            output_list = [
                output
                for output in cls.annotation["outputs"]
                if output.startswith("out_") and output[len("out_") :].isnumeric()
            ]

            if len(output_list) == 0:
                code.line("func_obj.execute()")

            else:  # this class returns something, unclear if native Ogn
                # extract the output; if it's a tuple iterate over values, if not continue for single value
                code.line("ret = func_obj.execute()")
                by_keys = {int(output[len("out_") :]): output for output in output_list}

                for key, output in by_keys.items():

                    if not cls.returns_tuple:
                        code.line("out_value = ret")
                    else:
                        code.line(f"out_value = ret[{key}]")

                    typename = cls.annotation["outputs"][output]["type"]
                    conversion = type_definitions.AutoNodeTypeConversion.from_ogn_type(typename)
                    if conversion:  # convertible

                        if conversion.type_to_og is not None:
                            # converter = f"{conversion.type_to_og.__qualname__}"
                            if (
                                conversion.type_to_og_conversion_method
                                == type_definitions.AutoNodeTypeConversion.Method.ASSIGN
                            ):  # noqa
                                code.line(
                                    f"db.outputs.out_{key} = type_definitions.AutoNodeTypeConversion"
                                    f".from_ogn_type('{typename}').type_to_og(out_value)"
                                )
                            elif (
                                conversion.type_to_og_conversion_method
                                == type_definitions.AutoNodeTypeConversion.Method.MODIFY
                            ):  # noqa
                                code.line(
                                    f"type_definitions.AutoNodeTypeConversion.from_ogn_type('{typename}')"
                                    f".type_to_og(out_value, db.outputs.out_{key})"
                                )
                        else:
                            code.line(f"db.outputs.out_{key} = out_value")
                    else:  # non-convertible
                        code.line("out_id = type_definitions.TypeRegistry.add_to_graph(out_value)")
                        code.line(f"db.outputs.out_{key} = out_id")

            code.line("db.outputs.exec = True")

            code.line("return True")

        carb.log_verbose(f"Generated code for {cls.unique_name}:\n{str(code)}")
        exec(str(code), globals())  # noqa: PLW0122

        # makes this generated function a class function
        cls.compute = compute  # noqa - compute is defined by the exec

    # ================================================================================
    @classmethod
    def internal_state(cls):
        return None if cls.autograph_data is None else copy.copy(cls.autograph_data)
