"""
***************************************************************************
    qgsfunction.py
    ---------------------
    Date                 : May 2018
    Copyright            : (C) 2018 by Denis Rouzaud
    Email                : denis@opengis.ch
***************************************************************************
*                                                                         *
*   This program is free software; you can redistribute it and/or modify  *
*   it under the terms of the GNU General Public License as published by  *
*   the Free Software Foundation; either version 2 of the License, or     *
*   (at your option) any later version.                                   *
*                                                                         *
***************************************************************************
"""

import inspect
import string
import traceback

from qgis._core import (
    Qgis,
    QgsExpression,
    QgsExpressionFunction,
    QgsFeatureRequest,
    QgsMessageLog,
)
from qgis.PyQt.QtCore import QCoreApplication


class QgsPyExpressionFunction(QgsExpressionFunction):
    """Python expression function"""

    def __init__(
        self,
        function,
        name,
        group,
        helptext="",
        usesgeometry=False,
        referenced_columns=QgsFeatureRequest.ALL_ATTRIBUTES,
        handlesnull=False,
        params_as_list=False,
    ):
        # Call the parent constructor
        # -1 means that function can take any number of arguments
        QgsExpressionFunction.__init__(self, name, -1, group, helptext)
        self.function = function
        self.params_as_list = params_as_list
        self.usesgeometry = usesgeometry
        self.referenced_columns = referenced_columns
        self.handlesnull = handlesnull

    def func(self, values, context, parent, node):
        feature = None
        if context:
            feature = context.feature()
        try:
            # Inspect the inner function signature to get the list of parameters
            parameters = inspect.signature(self.function).parameters
            kwvalues = {}

            # Handle special parameters
            # those will not be inserted in the arguments list
            # if they are present in the function signature
            if "context" in parameters:
                kwvalues["context"] = context
            if "feature" in parameters:
                kwvalues["feature"] = feature
            if "parent" in parameters:
                kwvalues["parent"] = parent

            # In this context, values is a list of the parameters passed to the expression.
            # If self.params_as_list is True, values is passed as is to the inner function.
            if self.params_as_list:
                return self.function(values, **kwvalues)
            # Otherwise (default), the parameters are expanded
            return self.function(*values, **kwvalues)

        except Exception as ex:
            tb = traceback.format_exception(None, ex, ex.__traceback__)
            formatted_traceback = "".join(tb)
            formatted_exception = f"{ex}:<pre>{formatted_traceback}</pre>"
            parent.setEvalErrorString(formatted_exception)
            return None

    def usesGeometry(self, node):
        return self.usesgeometry

    def referencedColumns(self, node):
        return self.referenced_columns

    def handlesNull(self):
        return self.handlesnull


def register_function(
    function,
    args="auto",
    group: str = "custom",
    usesgeometry: bool = False,
    referenced_columns=[QgsFeatureRequest.ALL_ATTRIBUTES],
    handlesnull: bool = False,
    params_as_list: bool = False,
    **kwargs,
):
    r"""
    Register a Python function to be used as a expression function.

    The function signature may contain special parameters (in any order at the end of the signature):

    - feature: the QgsFeature related to the current evaluation
    - parent: the QgsExpression parent
    - context: the QgsExpressionContext related to the current evaluation

    If those parameters are present in the function signature, they will be automatically passed to the function,
    without the need to specify them in the expression.

    Functions should return a value compatible with QVariant.

    Evaluation errors can be raised using:

    .. code-block:: python

        parent.setEvalErrorString("Error message")

    :param function: the Python function to be used as an expression function
    :param args: Argument list

      .. deprecated:: 3.32
          Use ``params_as_list`` if you want to pass parameters as a list.

    :param group: the expression group in which the function should be added
    :param usesgeometry: Defines if this expression requires the geometry. By default ``False``.
    :param referenced_columns: An array of names of fields on which this expression works. By default ``[QgsFeatureRequest.ALL_ATTRIBUTES]``. Specifying a subset of fields or an empty list will result in a faster execution.
    :param handlesnull: Defines if this expression has custom handling for NULL values. If ``False``, the result will always be NULL as soon as any parameter is NULL. ``False`` by default.
    :param params_as_list: If True, the function will receive the expression parameters as a list. If ``False``, the function will receive the parameters as individual arguments. ``False`` by default.
    :param \**kwargs:
        See below
    :Keyword Arguments:

    * *register* (``bool``): Set to ``False`` to create the QgsPyExpressionFunction without registering it. Useful for testing puposes. By default ``True``.
    * *name* (``str``): If provided, replace the function name
    * *helpText* (``str``): If provided, used in the help tooltip instead of the function docstring

    :return:
    """

    # Format the help text
    helptemplate = string.Template("<h3>function $name</h3>\n$doc")
    name = kwargs.get("name", function.__name__)
    helptext = kwargs.get("helpText") or function.__doc__ or ""
    helptext = helptext.strip()

    register = kwargs.get("register", True)
    if register and QgsExpression.isFunctionName(name):
        if not QgsExpression.unregisterFunction(name):
            msgtitle = QCoreApplication.translate("UserExpressions", "User expressions")
            msg = QCoreApplication.translate(
                "UserExpressions",
                "The user expression {0} already exists and could not be unregistered.",
            ).format(name)
            QgsMessageLog.logMessage(msg + "\n", msgtitle, Qgis.MessageLevel.Warning)
            return None

    function.__name__ = name
    helptext = helptemplate.safe_substitute(name=name, doc=helptext)

    # Legacy: if args was not 'auto', parameters were passed as a list
    params_as_list = params_as_list or args != "auto"
    f = QgsPyExpressionFunction(
        function,
        name,
        group,
        helptext,
        usesgeometry,
        referenced_columns,
        handlesnull,
        params_as_list,
    )

    if register:
        QgsExpression.registerFunction(f)
    return f


def qgsfunction(args="auto", group: str = "custom", **kwargs):
    r"""
    Decorator function used to define a user expression function.

    The decorated function signature may contain special parameters (in any order at the end of the signature):

    - ``feature``: the QgsFeature related to the current evaluation
    - ``parent``: the QgsExpression parent
    - ``context``: the QgsExpressionContext related to the current evaluation

    If those parameters are present in the function signature, they will be automatically passed to the function,
    without the need to specify them in the expression.

    Functions should return a value compatible with QVariant.

    Evaluation errors can be raised using:

    .. code-block:: python

        parent.setEvalErrorString("Error message")

    :param args: Argument list

      .. deprecated:: 3.32
          Use ``params_as_list`` if you want to pass parameters as a list.

    :param str group: The expression group to which this expression should be added.
    :param \**kwargs:
        See below
    :Keyword Arguments:
        * **usesgeometry** (*bool*) --
          Defines if this expression requires the geometry. By default ``False``.
        * **referenced_columns** (*List[str]*) --
          An array of names of fields on which this expression works. By default ``[QgsFeatureRequest.ALL_ATTRIBUTES]``. Specifying a subset of fields or an empty list will result in a faster execution.
        * **handlesnull** (*bool*) --
          Defines if this expression has custom handling for NULL values. If ``False``, the result will always be NULL as soon as any parameter is NULL. ``False`` by default.
        * **params_as_list** (*bool*) --
          Defines if the parameters are passed to the function as a list. If set to ``False``, they will be expanded. By default ``False``. Since QGIS 3.32.
        * **register** (*bool*) --
          Set to ``False`` to create the QgsPyExpressionFunction without registering it. Useful for testing puposes. By default ``True``.
        * **name** (*str*) --
          If provided, replace the function name.
        * **helpText** (*str*) --
          If provided, used in the help tooltip instead of the function docstring.


    Example 1 (Basic function, with default parameters)
    ---------------------------------------------------

    .. code-block:: python

        @qgsfunction(group="python")
        def center(text, width, fillchar=" "):
            return text.center(width, fillchar)

    This registers a function called "center" in the "python" group.
    This expression requires two parameters: ``text`` and ``width``.
    An optional third parameter can be provided: ``fillchar``.

    .. code-block:: text

        >>> center('Hello', 10)
        '  Hello   '
        >>> center('Hello', 15, '*')
        '*****Hello*****'


    Example 2 (variable number of parameters)
    -----------------------------------------

    .. code-block:: python

        @qgsfunction(group="custom")
        def fibonnaci_numbers(*args):
            def fibonnaci(n):
                if n <= 1:
                    return n
                else:
                    return fibonnaci(n - 1) + fibonnaci(n - 2)
            return [fibonnaci(arg) for arg in args]

    This registers a function called "fibonnaci_numbers" in the "custom" group.
    This expression can be called with any number of parameters, and will return a list of the fibonnaci numbers
    corresponding to the parameters.

    .. code-block:: text

        >>> fibonnaci_numbers()
        []
        >>> fibonnaci_numbers(1)
        [1]
        >>> fibonnaci_numbers(3)
        [2]
        >>> fibonnaci_numbers(1, 2, 3, 4, 5, 6, 7)
        [ 1, 1, 2, 3, 5, 8, 13 ]


    Example 3 (feature and parent)
    ------------------------------

    .. code-block:: python

        @qgsfunction(group="custom")
        def display_field(fieldname, feature, parent):
            try:
                return f"<b>{fieldname}</b>: {feature[fieldname]}"
            except KeyError:
                parent.setEvalErrorString(f"Field {fieldname} not found")

    This registers a function called "display_field" in the "custom" group.
    This expression requires an unique parameter: ``fieldname``. ``feature`` is automatically passed to the function.
    ``parent`` is the QgsExpression parent, and can be used to raise an error.

    .. code-block:: text

        >>> display_field("altitude")
        <b>altitude</b>: 164
        >>> display_field("lat")
        <b>lat</b>: 45.4

    Example 4 (context)
    -------------------

    .. code-block:: python

        @qgsfunction(group="custom")
        def title_layer_name(context):
            return context.variable("layer_name").title()

    This registers a function called "title_layer_name" in the "custom" group. It takes no parameters,
    but extracts the layer name from the expression context and returns it as a title case string.

    """

    def wrapper(func):
        return register_function(func, args, group, **kwargs)

    return wrapper
