import ctypes
import platform
import re
from typing import TypeVar

from pytensor.graph.basic import Constant
from pytensor.graph.type import Type
from pytensor.link.c.interface import CLinkerType
from pytensor.utils import Singleton


D = TypeVar("D")
T = TypeVar("T", bound=Type)


class CType(Type[D], CLinkerType):
    """Convenience wrapper combining `Type` and `CLinkerType`.

    PyTensor comes with several subclasses of such as:

    - `Generic`: for any python type

    - `TensorType`: for numpy.ndarray

    - `SparseTensorType`: for scipy.sparse

    But you are encouraged to write your own, as described in WRITEME.

    The following code illustrates the use of a Type instance,
    here tensor.fvector:

    .. code-block:: python

        # Declare a symbolic floating-point vector using __call__
        b = tensor.fvector()

        # Create a second Variable with the same Type instance
        c = tensor.fvector()

    Whenever you create a symbolic variable in pytensor (technically,
    `Variable`) it will contain a reference to a Type instance. That
    reference is typically constant during the lifetime of the
    Variable.  Many variables can refer to a single Type instance, as
    do b and c above.  The Type instance defines the kind of value
    which might end up in that variable when executing a `Function`.
    In this sense, pytensor is like a strongly-typed language because
    the types are included in the graph before the values.  In our
    example above, b is a Variable which is guaranteed to correspond
    to a numpy.ndarray of rank 1 when we try to do some computations
    with it.

    Many `Op` instances will raise an exception if they are applied to
    inputs with incorrect types.  Type references are also useful to
    do type-checking in pattern-based optimizations.

    """


class Generic(CType, Singleton):
    r"""A type for a generic Python object exposed directly in C.

    This class implements the `CType` and `CLinkerType` interfaces
    for generic ``PyObject`` instances.

    It can be used to easily expose Python objects to `COp`\s.

    For example, ``obj_const = Constant(Generic(), my_obj)`` will construct
    an PyTensor constant that can be used to access the Python object ``my_obj``
    directly in a `COp` (e.g. as a ``PyObject *``).

    """

    def filter(self, data, strict=False, allow_downcast=None):
        return data

    def is_valid_value(self, a):
        return True

    def c_declare(self, name, sub, check_input=True):
        return f"""
        PyObject* {name};
        """

    def c_init(self, name, sub):
        return f"""
        {name} = NULL;
        """

    def c_extract(self, name, sub, check_input=True, **kwargs):
        return f"""
        Py_INCREF(py_{name});
        {name} = py_{name};
        """

    def c_cleanup(self, name, sub):
        return f"""
        Py_XDECREF({name});
        """

    def c_sync(self, name, sub):
        return f"""
        assert(py_{name}->ob_refcnt > 1);
        Py_DECREF(py_{name});
        py_{name} = {name} ? {name} : Py_None;
        Py_INCREF(py_{name});
        """

    def c_code_cache_version(self):
        return (1,)

    def __str__(self):
        return self.__class__.__name__


generic = Generic()

_cdata_type = None

if platform.python_implementation() != "PyPy":
    _cdata_type = ctypes.py_object.from_address(
        ctypes.addressof(ctypes.pythonapi.PyCapsule_Type)
    ).value


class CDataType(CType[D]):
    """
    Represents opaque C data to be passed around. The intent is to
    ease passing arbitrary data between ops C code.

    The constructor builds a type made to represent a C pointer in pytensor.

    Parameters
    ----------
    ctype
        The type of the pointer (complete with the `*`).

    freefunc
        A function to call to free the pointer. This function must
        have a `void` return and take a single pointer argument.

    version
        The version to use in PyTensor cache system.
    """

    __props__ = (
        "ctype",
        "freefunc",
        "headers",
        "header_dirs",
        "libraries",
        "lib_dirs",
        "extra_support_code",
        "compile_args",
        "version",
    )

    def __init__(
        self,
        ctype,
        freefunc=None,
        headers=(),
        header_dirs=(),
        libraries=(),
        lib_dirs=(),
        compile_args=(),
        extra_support_code="",
        version=None,
    ):
        assert isinstance(ctype, str)
        self.ctype = ctype
        if freefunc is not None:
            assert isinstance(freefunc, str)
        self.freefunc = freefunc
        self.headers = tuple(headers)
        self.header_dirs = tuple(header_dirs)
        self.libraries = tuple(libraries)
        self.lib_dirs = tuple(lib_dirs)
        self.compile_args = tuple(compile_args)
        self.extra_support_code = extra_support_code
        self._fn = None
        self.version = version

    def filter(self, data, strict=False, allow_downcast=None):
        # We ignore this type-check (_cdata_type is None) in PyPy
        # because this type is not exposed to us.
        if data is not None and _cdata_type is not None:
            if not isinstance(data, _cdata_type):
                raise TypeError("expected None or a PyCapsule")

        return data

    def c_declare(self, name, sub, check_input=True):
        return f"""
        {self.ctype} {name};
        """

    def c_init(self, name, sub):
        return f"{name} = NULL;"

    def c_extract(self, name, sub, check_input=True, **kwargs):
        fail = sub["fail"]
        return f"""
  {name} = ({self.ctype})PyCapsule_GetPointer(py_{name}, NULL);
  if ({name} == NULL) {fail}
        """

    def c_sync(self, name, sub):
        freefunc = self.freefunc
        if freefunc is None:
            freefunc = "NULL"
        s = f"""
Py_XDECREF(py_{name});
if ({name} == NULL) {{
  py_{name} = Py_None;
  Py_INCREF(py_{name});
}} else {{
  py_{name} = PyCapsule_New((void *){name}, NULL,
                              _capsule_destructor);
  if (py_{name} != NULL) {{
    if (PyCapsule_SetContext(py_{name}, (void *){freefunc}) != 0) {{
      /* This won't trigger a call to freefunc since it could not be
         set. The error case below will do it. */
      Py_DECREF(py_{name});
      /* Signal the error */
      py_{name} = NULL;
    }}
  }}
}}"""
        if self.freefunc is not None:
            s += f"""
if (py_{name} == NULL) {{ {freefunc}({name}); }}
"""
        return s

    def c_cleanup(self, name, sub):
        # No need to do anything here since the CObject/Capsule will
        # free the data for us when released.
        return ""

    def c_headers(self, **kwargs):
        return self.headers

    def c_header_dirs(self, **kwargs):
        return self.header_dirs

    def c_libraries(self, **kwargs):
        return self.libraries

    def c_lib_dirs(self, **kwargs):
        return self.lib_dirs

    def c_support_code(self, **kwargs):
        return (
            """
void _capsule_destructor(PyObject *o) {
    void *d = PyCapsule_GetContext(o);
    void *p = PyCapsule_GetPointer(o, NULL);
    void (*f)(void *) = (void (*)(void *))d;
    if (f != NULL) f(p);
}
"""
            + self.extra_support_code
        )

    def c_compile_args(self, **kwargs):
        return self.compile_args

    def c_code_cache_version(self):
        v = (3,)
        if self.version is not None:
            v = (*v, self.version)
        return v

    def __str__(self):
        return f"{self.__class__.__name__}{{{self.ctype}}}"

    def __setstate__(self, dct):
        self.__dict__.update(dct)
        if not hasattr(self, "headers"):
            self.headers = ()
            self.header_dirs = ()
            self.libraries = ()
            self.lib_dirs = ()
        if not hasattr(self, "version"):
            self.version = None


class CDataTypeConstant(Constant[T]):
    def merge_signature(self):
        # We don't want to merge constants that don't point to the
        # same object.
        return id(self.data)

    def signature(self):
        # There is no way to put the data in the signature, so we
        # don't even try
        return (self.type,)


CDataType.constant_type = CDataTypeConstant


class EnumType(CType, dict):
    """
    Main subclasses:
     - :class:`EnumList`
     - :class:`CEnumType`

    Op parameter class that allows to create enumerations of constant values.

     - Constants are available as object attributes in Python code and as macro-defined constants in C code.
     - Constants can be floating values, integers, or booleans (automatically converted to integers).
     - Constants name must start with a capital letter and contain capital letters, underscores or digits.
     - A constant can have an alias, and then be available through both constant name and constant alias.

    **Example**

    .. code-block:: python

        enum = EnumType(CONSTANT_1=1, CONSTANT_2=2.5, CONSTANT_3=False, CONSTANT_4=True)
        print(enum.CONSTANT_1, enum.CONSTANT_2, enum.CONSTANT_3, enum.CONSTANT_4)
        # will print 1 2.5 0 1

    In C code:

    .. code-block:: c

        int constant_1 = CONSTANT_1;
        double constant_2 = CONSTANT_2;
        int constant_3 = CONSTANT_3; // constant_3 == 0
        int constant_4 = CONSTANT_4; // constant_4 == 1

    You can also specify a C type for the op param. Default C type is ``double``.

    .. code-block:: python

        enum = EnumType(CONSTANT_1=0, CONSTANT_2=1, CONSTANT_3=2, ctype="size_t")
        # In C code, the Op param will then be a ``size_t``.

    .. note::

        You can also specify a C name (``cname``) or the current enumeration. This C name may be
        used to name functions related to that specific enumeration, e.g. for debugging
        purposes. Default C name is the C type (with any sequence of spaces replaced with
        an underscore). If you want to debug and your C type is quite generic (e.g.
        ``int`` or ``double``), we recommend you specify a C name.

        C name must be a valid C identifier.

        .. code-block:: python

            enum = EnumType(
                CONSTANT_1=0, CONSTANT_2=1, CONSTANT_3=2, ctype="size_t", cname="MyEnumName"
            )

    **Example with aliases**

    When creating an enum, you can give some aliases to specific constants while keeping other constants without aliases.
    An alias must be a string, and there is currently no string format constraints.

    To give an alias to a constant in the EnumType constructor, use the following key-value syntax::

        constant_name = (constant_alias, constant_value)

    You can then retrieve a constant from an alias with method ``EnumType.fromalias()``.

    Aliases are intended to be used in Python code only (only constants names are available in C code).
    Especially, an alias will be recognized by ``Enumtype.filter()`` method with non-strict filtering,
    allowing a maximum flexibility for converting strings to numeric constants available in Python and C code.

    .. code-block:: python

        from pytensor.link.c.type import EnumType

        # You can remark that constant 'C' does not have an alias.
        enum = EnumType(A=("alpha", 1), B=("beta", 2), C=3, D=("delta", 4))

        # Constants are all directly available by name.
        print(enum.A, enum.B, enum.C, enum.D)

        # But we can also now get some constants by alias.
        a = enum.fromalias("alpha")
        b = enum.fromalias("beta")
        d = enum.fromalias("delta")

        # If method fromalias() receives an unknown alias,
        # it will looks for a constant with this alias
        # as exact constant name.
        c = enum.fromalias("C")  # will get enum.C

        # An alias defined in an EnumType will be correctly converted with non-strict filtering.
        value = enum.filter("delta", strict=False)
        # value now contains enum.D, ie. 4.

    .. note::

        This Type (and subclasses) is not complete and should never be used for regular graph operations.

    """

    def __init_ctype(self, ctype):
        # C type may be a list of keywords, e.g. "unsigned long long".
        # We should check each part.
        ctype_parts = ctype.split()
        if not all(re.match("^[A-Za-z_][A-Za-z0-9_]*$", el) for el in ctype_parts):
            raise TypeError(f"{type(self).__name__}: invalid C type.")
        self.ctype = " ".join(ctype_parts)

    def __init_cname(self, cname):
        if not re.match("^[A-Za-z_][A-Za-z0-9_]*$", cname):
            raise TypeError(f"{type(self).__name__}: invalid C name.")
        self.cname = cname

    def __init__(self, **kwargs):
        self.__init_ctype(kwargs.pop("ctype", "double"))
        self.__init_cname(kwargs.pop("cname", self.ctype.replace(" ", "_")))
        self.aliases = dict()
        for k in kwargs:
            if re.match("^[A-Z][A-Z0-9_]*$", k) is None:
                raise AttributeError(
                    f'{type(self).__name__}: invalid enum name: "{k}". '
                    "Only capital letters, underscores and digits "
                    "are allowed."
                )
            if isinstance(kwargs[k], list | tuple):
                if len(kwargs[k]) != 2:
                    raise TypeError(
                        f"{type(self).__name__}: when using a tuple to define a constant, your tuple should contain 2 values: "
                        "constant alias followed by constant value."
                    )
                alias, value = kwargs[k]
                if not isinstance(alias, str):
                    raise TypeError(
                        f'{type(self).__name__}: constant alias should be a string, got "{alias}".'
                    )
                if alias == k:
                    raise TypeError(
                        f"{type(self).__name__}: it's useless to create an alias "
                        "with the same name as its associated constant."
                    )
                if alias in self.aliases:
                    raise TypeError(
                        f'{type(self).__name__}: consant alias "{alias}" already used.'
                    )
                self.aliases[alias] = k
                kwargs[k] = value
            if isinstance(kwargs[k], bool):
                kwargs[k] = int(kwargs[k])
            elif not isinstance(kwargs[k], int | float):
                raise TypeError(
                    f'{type(self).__name__}: constant "{k}": expected integer or floating value, got "{type(kwargs[k]).__name__}".'
                )
        if [a for a in self.aliases if a in self]:
            raise TypeError(
                f"{type(self).__name__}: some aliases have same names as constants."
            )
        super().__init__(**kwargs)

    def fromalias(self, alias):
        """
        Get a constant value by its alias.
        If there is not such alias in this enum, look for a constant
        with this alias as constant name.
        """
        return self[self.aliases[alias]] if alias in self.aliases else self[alias]

    def has_alias(self, alias):
        """
        return True if and only if this enum has this alias.
        """
        return alias in self.aliases

    def get_aliases(self):
        """
        Return the sorted tuple of all aliases in this enumeration.
        """
        return tuple(sorted(self.aliases))

    def __repr__(self):
        names_to_aliases = {constant_name: "" for constant_name in self}
        for alias in self.aliases:
            names_to_aliases[self.aliases[alias]] = f"({alias})"
        args = ", ".join(f"{k}{names_to_aliases[k]}:{self[k]}" for k in sorted(self))
        return f"{type(self).__name__}<{self.ctype}>({args})"

    def __getattr__(self, key):
        if key in self:
            return self[key]
        return CType.__getattr__(self, key)

    def __setattr__(self, key, value):
        if key in self:
            raise NotImplementedError("constant values are immutable.")
        CType.__setattr__(self, key, value)

    def __setitem__(self, key, value):
        raise NotImplementedError("constant values are immutable.")

    def __delitem__(self, key):
        raise NotImplementedError("constant values are immutable.")

    def __hash__(self):
        # All values are Python basic types, then easy to hash.
        return hash(
            (
                type(self),
                self.ctype,
                *sorted(self.items()),
                *sorted(self.aliases.items()),
            )
        )

    def __eq__(self, other):
        return (
            type(self) is type(other)
            and self.ctype == other.ctype
            and len(self) == len(other)
            and len(self.aliases) == len(other.aliases)
            and all(k in other for k in self)
            and all(a in other.aliases for a in self.aliases)
            and all(self[k] == other[k] for k in self)
            and all(self.aliases[a] == other.aliases[a] for a in self.aliases)
        )

    # EnumType should be used to create constants available in both Python and C code.
    # However, for convenience, we make sure EnumType can have a value, like other common types,
    # such that it could be used as-is as an op param.
    # C type of value is defined in self.ctype.

    def filter(self, data, strict=False, allow_downcast=None):
        if not strict:
            if isinstance(data, bool):
                data = int(data)
            elif isinstance(data, str):
                # We now accept strings as data values.
                # Strings should be a constant alias or a constant name.
                data = self.fromalias(data)
        assert data in self.values()
        return data

    def values_eq(self, a, b):
        return a == b

    def values_eq_approx(self, a, b):
        # For an enum, it does not have a meaning to be approx equal.
        return self.values_eq(a, b)

    pyint_compat_code = """
    #if PY_MAJOR_VERSION >= 3
        #ifndef PyInt_Check
            #define PyInt_Check PyLong_Check
        #endif
        #ifndef PyInt_AsLong
            #define PyInt_AsLong PyLong_AsLong
        #endif
    #endif
    """

    def c_to_string(self):
        """
        Return code for a C function that will convert an enumeration value
        to a string representation. The function prototype is:

        .. code-block:: c

            int pytensor_enum_to_string_<cname>(<ctype> value, char* output_string);

        Where ``ctype`` and ``cname`` are the C type and the C name of current PyTensor enumeration.

        ``output_string`` should be large enough to contain the longest name in this enumeration.

        If given value is unknown, the C function sets a Python ValueError exception and returns a non-zero.

        This C function may be useful to retrieve some runtime information.
        It is available in C code when pytensor flag ``config.cmodule__debug`` is set to ``True``.
        """
        cases = "".join(
            f"""
                   case {name}: sprintf(out, "{name}"); break;
                   """
            for name in self
        )
        return f"""
        #ifdef DEBUG
        int pytensor_enum_to_string_{self.cname}({self.ctype} in, char* out) {{
            int ret = 0;
            switch(in) {{
                {cases}
                default:
                    PyErr_SetString(PyExc_ValueError, "{type(self).__name__}:  unknown enum value.");
                    ret = -1;
                    break;
            }}
            return ret;
        }}
        #endif
        """

    def c_support_code(self, **kwargs):
        return (
            self.pyint_compat_code
            + "".join(
                f"""
            #define {k} {self[k]!s}
            """
                for k in sorted(self)
            )
            + self.c_to_string()
        )

    def c_declare(self, name, sub, check_input=True):
        return f"""{self.ctype} {name};"""

    def c_init(self, name, sub):
        return f"{name} = ({self.ctype})0;"

    def c_cleanup(self, name, sub):
        return ""

    def c_extract(self, name, sub, check_input=True, **kwargs):
        fail = sub["fail"]
        return f"""
        if (PyInt_Check(py_{name})) {{
            {name} = ({self.ctype})PyInt_AsLong(py_{name});
        }} else {{
            {name} = ({self.ctype})PyFloat_AsDouble(py_{name});
        }}
        if (PyErr_Occurred()) {{
            {fail}
        }}
        """

    def c_code_cache_version(self):
        return (2, self.ctype, self.cname, tuple(self.items()))

    def c_sync(self, name, sub):
        raise NotImplementedError("Variables of this type cannot be graph outputs")


class EnumList(EnumType):
    """
    **Inherit from**:
     - :class:`EnumType`

    Op parameter class that allows to create enumeration of constant values.
    Same as :class:`EnumType`, but automatically gives an unique integer value for each constant in a list of
    constants names (constant at index ``i`` in the list will receive value ``i``,
    with ``i`` from ``0`` to ``len(constants) - 1``).

    Example::

        enum = EnumList(
            "CONSTANT_1", "CONSTANT_2", "CONSTANT_3", "CONSTANT_4", "CONSTANT_5"
        )
        print(
            enum.CONSTANT_1,
            enum.CONSTANT_2,
            enum.CONSTANT_3,
            enum.CONSTANT_4,
            enum.CONSTANT_5,
        )
        # will print: 0 1 2 3 4

    Like :class:`EnumType`, you can also define the C type and a C name for the op param.
    Default C type is ``int``::

        enum = EnumList(
            "CONSTANT_1", "CONSTANT_2", "CONSTANT_3", "CONSTANT_4", ctype="unsigned int"
        )

    Like :class:`EnumType`, you can also add an alias to a constant, by replacing the only constant name
    (e.g. ``'CONSTANT_NAME'``) by a couple with constant name first and constant alias second
    (e.g. ``('CONSTANT_NAME', 'constant_alias')``).

    .. code-block:: python

        enum = EnumList(("A", "alpha"), ("B", "beta"), "C", "D", "E", "F", ("G", "gamma"))

    See test class :class:`tests.graph.test_types.TestOpEnumList` for a working example.

    """

    def __init__(self, *args, **kwargs):
        assert len(kwargs) in (0, 1, 2), (
            type(self).__name__
            + ': expected 0 to 2 extra parameters ("ctype", "cname").'
        )
        ctype = kwargs.pop("ctype", "int")
        cname = kwargs.pop("cname", None)

        for arg_rank, arg in enumerate(args):
            if isinstance(arg, list | tuple):
                if len(arg) != 2:
                    raise TypeError(
                        f"{type(self).__name__}: when using a tuple to define a constant, your tuple should contain 2 values: "
                        "constant name followed by constant alias."
                    )
                constant_name, constant_alias = arg
                if not isinstance(constant_alias, str):
                    raise TypeError(
                        f'{type(self).__name__}: constant alias should be a string, got "{constant_alias}".'
                    )
                constant_value = (constant_alias, arg_rank)
            else:
                constant_name = arg
                constant_value = arg_rank
            if not isinstance(constant_name, str):
                raise TypeError(
                    f'{type(self).__name__}: constant name should be a string, got "{constant_name}".'
                )
            if constant_name in kwargs:
                raise TypeError(
                    f'{type(self).__name__}: constant name already used ("{constant_name}").'
                )
            kwargs[constant_name] = constant_value

        kwargs.update(ctype=ctype)
        if cname is not None:
            kwargs.update(cname=cname)
        super().__init__(**kwargs)


class CEnumType(EnumList):
    """
    **Inherit from**:
     - :class:`EnumList`

    Op parameter class that allows to create enumeration of constant values that represent C-defined constants.

     - Constant should have same names as in C.
     - In Python, constants will have arbitrary-defined values.
       They should be used only for choices, not for its values.
     - In C code, the real values defined in C will be used.
       They could be used either for choices or for its real values.

    Like :class:`EnumList`, you can also define the C type and a C name for the op param.
    Default C type is ``int``.

    .. code-block:: python

        enum = CEnumType(
            "CONSTANT_CNAME_1", "CONSTANT_CNAME_2", "CONSTANT_CNAME_3", ctype="long"
        )

    Like :class:`EnumList`, you can also add an alias to a constant, with same syntax as in :class:`EnumList`.

    See test :func:`tests.graph.test_types.TestEnumTypes.test_op_with_cenumtype` for a working example.

    .. note::

        Be sure C constants are available in your C code. If they come from a C header, consider implementing
        ``c_headers()`` and ``c_header_dirs()`` in the Op class where you use CEnumType as op parameter type.

    """

    def c_support_code(self, **kwargs):
        return self.pyint_compat_code + self.c_to_string()

    def c_extract(self, name, sub, check_input=True, **kwargs):
        swapped_dict = {v: k for (k, v) in self.items()}
        # swapped_dict's keys are integers.

        fail = sub["fail"]
        cases = "".join(
            f"""
           case {i}: {name} = {swapped_dict[i]}; break;
           """
            for i in sorted(swapped_dict)
        )
        return f"""
        switch(PyInt_AsLong(py_{name})) {{
            {cases}
            default:
                PyErr_SetString(PyExc_ValueError, "CEnumType: invalid value to map to C constants.");
                {{{fail}}}
                break;
        }}
        """

    def c_code_cache_version(self):
        return (1, super().c_code_cache_version())
