import enum
from abc import ABC, abstractmethod
from collections import namedtuple
from typing import Callable, Dict, Iterable, NewType, Optional, Tuple

import numpy as np

Vector3d = NewType("Vector3d", np.ndarray)
Vector3f = NewType("Vector3f", np.ndarray)
Vector3h = NewType("Vector3h", np.ndarray)

Float = NewType("Float", np.float)
Float2 = NewType("Float2", np.ndarray)
Float3 = NewType("Float3", np.ndarray)
Float4 = NewType("Float4", np.ndarray)

Half = NewType("Half", np.half)
Half2 = NewType("Half2", np.ndarray)
Half3 = NewType("Half3", np.ndarray)
Half4 = NewType("Half4", np.ndarray)

Double = NewType("Double", np.double)
Double2 = NewType("Double2", np.ndarray)
Double3 = NewType("Double3", np.ndarray)
Double4 = NewType("Double4", np.ndarray)

Int = NewType("Int", int)
Int2 = NewType("Int2", np.ndarray)
Int3 = NewType("Int3", np.ndarray)
Int4 = NewType("Int4", np.ndarray)


def int_to_int32(input: int) -> Int:
    return input & 0xFFFFFFFF


Timecode = NewType("Timecode", float)
Token = NewType("Token", str)

UInt = NewType("UInt", np.uint)
UChar = NewType("UChar", np.ubyte)


def int_to_uchar(input: int) -> UChar:
    return input & 0xFF


Matrix2d = NewType("Matrix2d", np.ndarray)
Matrix3d = NewType("Matrix3d", np.ndarray)
Matrix4d = NewType("Matrix4d", np.ndarray)

Normal3f = NewType("Normal3f", np.ndarray)
Normal3d = NewType("Normal3d", np.ndarray)
Normal3h = NewType("Normal3h", np.ndarray)

Point3f = NewType("Point3f", np.ndarray)
Point3d = NewType("Point3d", np.ndarray)
Point3h = NewType("Point3h", np.ndarray)

Quatd = NewType("Quatd", np.ndarray)
Quatf = NewType("Quatf", np.ndarray)
Quath = NewType("Point3h", np.ndarray)

TexCoord2d = NewType("TexCoord2d", np.ndarray)
TexCoord2f = NewType("TexCoord2f", np.ndarray)
TexCoord2h = NewType("TexCoord2h", np.ndarray)
TexCoord3d = NewType("TexCoord3d", np.ndarray)
TexCoord3f = NewType("TexCoord3f", np.ndarray)
TexCoord3h = NewType("TexCoord3h", np.ndarray)

Color3f = NewType("Color3f", np.ndarray)
Color4f = NewType("Color4f", np.ndarray)
Color3d = NewType("Color3d", np.ndarray)
Color4d = NewType("Color4d", np.ndarray)
Color3h = NewType("Color3h", np.ndarray)
Color4h = NewType("Color4h", np.ndarray)

all_types = [
    Vector3d,
    Vector3f,
    Vector3h,
    Float,
    Float2,
    Float3,
    Float4,
    Half,
    Half2,
    Half3,
    Half4,
    Double,
    Double2,
    Double3,
    Double4,
    Int,
    Int2,
    Int3,
    Int4,
    Timecode,
    Token,
    UInt,
    UChar,
    Matrix2d,
    Matrix3d,
    Matrix4d,
    Normal3f,
    Normal3d,
    Normal3h,
    Point3f,
    Point3d,
    Point3h,
    Quatd,
    Quatf,
    Quath,
    TexCoord2d,
    TexCoord2f,
    TexCoord2h,
    TexCoord3d,
    TexCoord3f,
    TexCoord3h,
    Color3f,
    Color4f,
    Color3d,
    Color4d,
    Color3h,
    Color4h,
]

TypeDesc = namedtuple(
    "TypeDesc",
    [
        "type",
        "og_type",
        "type_to_og",
        "type_to_og_conversion_method",
        "og_to_type",
        "og_to_type_conversion_method",
        "default",
    ],
)


# ================================================================================
class AutoNodeTypeConversion:
    """Static class for storing conversion methods between python types and Omnigraph types"""

    class Method(enum.Enum):
        ASSIGN = (0,)
        MODIFY = 1

    # flake8: noqa: E241
    types = [
        TypeDesc(bool, "bool", bool, Method.ASSIGN, None, Method.ASSIGN, False),
        TypeDesc(Color3d, "colord[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Color3f, "colorf[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Color3h, "colorh[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Color4d, "colord[4]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0, 0)),
        TypeDesc(Color4f, "colorf[4]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0, 0)),
        TypeDesc(Color4h, "colorh[4]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0, 0)),
        TypeDesc(Double, "double", None, Method.ASSIGN, None, Method.ASSIGN, 0),
        TypeDesc(Double2, "double[2]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0)),
        TypeDesc(Double3, "double[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Double4, "double[4]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0, 0)),
        TypeDesc(float, "double", None, Method.ASSIGN, None, Method.ASSIGN, 0),
        TypeDesc(Float, "float", np.float, Method.ASSIGN, None, Method.ASSIGN, 0),
        TypeDesc(Float2, "float[2]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0)),
        TypeDesc(Float3, "float[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Float4, "float[4]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0, 0)),
        TypeDesc(Half, "half", np.half, Method.ASSIGN, None, Method.ASSIGN, 0),
        TypeDesc(Half2, "half[2]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0)),
        TypeDesc(Half3, "half[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Half4, "half[4]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0, 0)),
        TypeDesc(int, "int64", None, Method.ASSIGN, None, Method.ASSIGN, 0),
        TypeDesc(Int, "int", int_to_int32, Method.ASSIGN, None, Method.ASSIGN, 0),
        TypeDesc(Int2, "int[2]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0)),
        TypeDesc(Int3, "int[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Int4, "int[4]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0, 0)),
        TypeDesc(Matrix2d, "matrixd[2]", None, Method.ASSIGN, None, Method.ASSIGN, ((1, 0), (0, 1))),
        TypeDesc(Matrix3d, "matrixd[3]", None, Method.ASSIGN, None, Method.ASSIGN, ((1, 0, 0), (0, 1, 0), (0, 0, 1))),
        TypeDesc(
            Matrix4d,
            "matrixd[4]",
            None,
            Method.ASSIGN,
            None,
            Method.ASSIGN,
            ((1, 0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1)),
        ),
        TypeDesc(Normal3d, "normald[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Normal3f, "normalf[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Normal3h, "normalh[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Point3d, "pointd[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Point3f, "pointf[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Point3h, "pointh[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Quatd, "quatd[4]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0, 0)),
        TypeDesc(Quatf, "quatf[4]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0, 0)),
        TypeDesc(Quath, "quath[4]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0, 0)),
        TypeDesc(str, "string", None, Method.ASSIGN, None, Method.ASSIGN, ""),
        TypeDesc(TexCoord2d, "texcoordd[2]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0)),
        TypeDesc(TexCoord2f, "texcoordf[2]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0)),
        TypeDesc(TexCoord2h, "texcoordh[2]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0)),
        TypeDesc(TexCoord3d, "texcoordd[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(TexCoord3f, "texcoordf[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(TexCoord3h, "texcoordh[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Timecode, "timecode", None, Method.ASSIGN, None, Method.ASSIGN, 0),
        TypeDesc(Token, "token", None, Method.ASSIGN, None, Method.ASSIGN, 0),
        TypeDesc(UInt, "uint", int_to_uchar, Method.ASSIGN, None, Method.ASSIGN, 0),
        TypeDesc(UChar, "uchar", int_to_uchar, Method.ASSIGN, None, Method.ASSIGN, 0),
        TypeDesc(Vector3d, "vectord[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Vector3f, "vectorf[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
        TypeDesc(Vector3h, "vectorh[3]", None, Method.ASSIGN, None, Method.ASSIGN, (0, 0, 0)),
    ]

    user_types = []

    # --------------------------------------------------------------------------------
    def __init__(self):
        raise RuntimeError("AutoNodeTypeConversion is a static class, do not instantiate")

    # --------------------------------------------------------------------------------
    @classmethod
    def from_type(cls, type_desc: type) -> Optional[TypeDesc]:
        """Searches the conversion registry using a python type.

        Args:
            type_desc: python type to convert

        Returns:
            TypeDesc for the found python type if it's found, None otherwise
        """
        for tb in cls.user_types:
            if tb.type == type_desc:
                return tb
        for tb in cls.types:
            if tb.type == type_desc:
                return tb
        return None

    # --------------------------------------------------------------------------------
    @classmethod
    def from_ogn_type(cls, og_type: str) -> Optional[TypeDesc]:
        """Searches the conversion registry using an Omnigraph type.
        Searches the user types first, defaults to the system types.

        Args:
            og_type: string representing the incoming ogn type

        Returns:
            TypeDesc for the found python type if it's found, None otherwise
        """

        # TODO (OS): Solve for arrays
        for tb in cls.user_types:
            og_coded_type = tb.og_type
            if og_coded_type == og_type:
                return tb
        for tb in cls.types:
            og_coded_type = tb.og_type
            if og_coded_type == og_type:
                return tb
        return None

    # --------------------------------------------------------------------------------
    @classmethod
    def register_type_conversion(
        cls,
        python_type: type,
        ogn_typename: str,
        python_to_ogn: Optional[Callable] = None,
        python_to_ogn_method: Method = Method.ASSIGN,
        ogn_to_python: Optional[Callable] = None,
        ogn_to_python_method: Method = Method.ASSIGN,
        default=None,
    ):
        """Registers a type conversion between a python type and an ogn type.
        Masks any existing system setting. If a previous user-submitted type conversion is registered, it will be overridden.

        Args:
            python_type: the type representation in python.
            ogn_typename: string representation of the ogn type. Node generation will fail if the type isn't recognized by ogn.
            python_to_ogn: [optional] function to convert a python return value to an OGN struct. Signature is Callable[[[python_type], object]. Defaults to None
            ogn_to_python: [optional] function to convert an OGN struct to a python return value. Signature is Callable[[[object], python_type]. Defaults to None
        """

        desc = TypeDesc(
            python_type, ogn_typename, python_to_ogn, python_to_ogn_method, ogn_to_python, ogn_to_python_method, default
        )
        unregistered = cls.unregister_type_conversion(python_type)
        if unregistered:
            carb.log_warn(
                "Registering an autograph type conversion for {type}->{ogn_typename} replaces anohter conversion for: {desc.type}->{og_type}"
            )

        cls.user_types.append(desc)

    # --------------------------------------------------------------------------------
    @classmethod
    def unregister_type_conversion(cls, python_type: type = None, ogn_type_name: str = None) -> Optional[TypeDesc]:
        """Unregisters a type conversion from python to ogn.
        Doesn't unregister system types.

        Args:
            python_type: the python type to be removed from support
            ogn_type_name: the ogn type name type to be removed from support

        Returns:
            The TypeDesc tuple just unregistered, or none.
        """
        if python_type is not None:
            for desc in cls.user_types:
                if desc.type == python_type:
                    cls.user_types.remove(desc)
                    return desc
        if ogn_type_name is not None:
            for desc in cls.user_types:
                if desc.og_type == ogn_type_name:
                    cls.user_types.remove(desc)
                    return desc
        return None


# ================================================================================
class AutoNodeDefinitionWrapper(ABC):
    """Container for a single node representation consumed by the Ogn code generator.
    Class is abstract and meant to be overridden. A sufficient implementation overrides these methods:
    * get_ogn(self) -> Dict
    * get_node_impl(self)
    * get_unique_name(self) -> str
    * get_module_name(self) -> str
    """

    def __init__(self):
        super().__init__()

    # --------------------------------------------------------------------------------
    @abstractmethod
    def get_ogn(self) -> Dict:
        """Get the Ogn dictionary representation of the node interface."""
        return {}

    # --------------------------------------------------------------------------------
    @abstractmethod
    def get_node_impl(self):
        """Returns the Ogn class implementing the node behavior. See omnigraph docuemntation on how to implement.
        A sufficient implementation contains a staticmethod with the function: compute(db)
        """
        return None

    # --------------------------------------------------------------------------------
    @abstractmethod
    def get_unique_name(self) -> str:
        """Get nodes unique name, to be saved as an accessor in the node database.

        Returns:
            the non-manlged unique name
        """
        return ""

    # --------------------------------------------------------------------------------
    @abstractmethod
    def get_module_name(self) -> str:
        """Get the module this autograph method was defined in.

        Returns:
            the module name
        """
        return ""


# ================================================================================
class AutoNodeDefinitionGenerator(ABC):
    """Defines an interface for generating a node definition"""

    _name = ""

    # --------------------------------------------------------------------------------
    @classmethod
    def name(cls):
        return cls._name

    # --------------------------------------------------------------------------------
    @classmethod
    @abstractmethod
    def generate_from_definitions(cls, new_type: type) -> Tuple[Iterable[AutoNodeDefinitionWrapper], Iterable[str]]:
        """This method scans the type new_type and outputs an AutoNodeDefinitionWrapper from it, representing the type,
        as well as a list of members it wishes to hide from the rest of the node extraction process.

        Args:
            new_type: the type to analyze by attribute

        Returns: a tuple of:
            Iterable[AutoNodeDefinitionWrapper] - an iterable of AutoNodeDefinitionWrapper - every node wrapper that is generated from this type.
            Iterable[str] - an iterable of all members covered by this handler that other handlers should ignore.
        """
        pass


class AutographDataWrapper:
    """Class to wrap around data being passed in the graph, to ensure a standard interface for passing data in the
    graph.  Accepts both reference and value types.
    """

    # --------------------------------------------------------------------------------
    def __init__(self, value) -> None:
        self._value = value
        self._type = type(value)
        if isinstance(self._type, AutographDataWrapper):
            raise KeyError("Attempted to wrap a datawrapper")

    # --------------------------------------------------------------------------------
    @property
    def value(self):
        """Returns the stored value without side effects"""
        return self._value

    # --------------------------------------------------------------------------------
    @property
    def type(self) -> type:
        """Returns the value type stored at initialization time without side effects"""
        return self._type


class AutographObjectStore:
    """Obejct store with simple put-pop interface."""

    # --------------------------------------------------------------------------------
    def __init__(self):
        self._store = {}

    # --------------------------------------------------------------------------------
    def pop(self, obj_id: int) -> AutographDataWrapper:
        """Attempts to return a value from the object store, and deletes it from the store.

        Args:
            obj_id: the object ID to be retrieved

        Returns:
            an AutographDataWrapper if successful

        Raises:
            KeyError if the object isn't there.
        """
        return self._store.pop(obj_id)

    # --------------------------------------------------------------------------------
    def put(self, obj_id: int, value: AutographDataWrapper) -> Optional[AutographDataWrapper]:
        """Places an object in the object store according to an object ID. If an object already exists in the data store
        with the same ID, it is popped.

        Args:
            obj_id: integer object ID for input object
            value: the actual object ID. Does not check for types.

        Returns:
            The previous object with the same ID stored in the object store, None otherwise.

        """
        swapped = self._store.get(obj_id, None)
        self._store[obj_id] = value
        return swapped


# ================================================================================
class TypeRegistry:

    """Main singleton for storing graph objects and generating and registering functions."""

    def __init__(self):
        raise RuntimeError("Singleton class, use instance() instead")

    # --------------------------------------------------------------------------------
    def _initialize(self):

        self.class_to_methods: Dict[type, Dict[str, str]] = {}

        # Meant to make accessing functions easier
        self.func_name_to_func = {}

        # Meant to keep a pointer to an object, so it doesn't get garbage collected
        self.refs = AutographObjectStore()

        # Meant to handle types whene generating nodes
        self.type_handlers = set()

        # Stores omnigraph implementations
        self._impl_modules = {}

    # --------------------------------------------------------------------------------
    @classmethod
    def _reset(cls):
        """Resets all internal data strucutres. Does not unregister nodes."""
        cls.instance()._initialize()

    # --------------------------------------------------------------------------------
    @classmethod
    def instance(cls):
        """Retrieves the class instance for this singleton.

        Returns:
            Class instance
        """
        if not hasattr(cls, "_instance"):
            cls._instance = cls.__new__(cls)
            cls._instance._initialize()
        return cls._instance

    # --------------------------------------------------------------------------------
    @classmethod
    def get_func(cls, unique_name: str) -> Optional[AutoNodeDefinitionWrapper]:
        """Retrieves a function from the object store

        Attributes
            unique_name: function's qualified name. Name mangling is handled by this class

        Returns:
            Function Wrapper.

        """
        return cls.instance().func_name_to_func.get(unique_name, None)

    # --------------------------------------------------------------------------------
    @classmethod
    def add_to_graph(cls, obj) -> int:
        """Adds an object `obj` to the data store without checking for uniqueness of held object.
        Does check uniqueness of reference object.

        Attributes
            obj: Object to add to the data store.

        Returns:
            the object id.
        """
        wrapper = AutographDataWrapper(obj)
        obj_id = id(wrapper)
        cls.instance().refs.put(obj_id, wrapper)
        return obj_id

    # --------------------------------------------------------------------------------
    @classmethod
    def remove_from_graph(cls, obj_id: int) -> AutographDataWrapper:
        """Attempts to remove an object ID from the object store.

        Attributes
            obj_id: the object ID to be removed from the data store.

        Returns:
            the object stored if it was found, None otherwise.
        """
        return cls.instance().refs.pop(obj_id)
