# Wrap a function.
#
# Author: Malte J. Ziebarth (malte.ziebarth@tum.de)
#
# Copyright (C) 2026 Technische Universität München
#
# Licensed under the EUPL, Version 1.2 or – as soon they will be approved by
# the European Commission - subsequent versions of the EUPL (the "Licence");
# You may not use this work except in compliance with the Licence.
# You may obtain a copy of the Licence at:
#
# https://joinup.ec.europa.eu/collection/eupl/eupl-text-eupl-12
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the Licence is distributed on an "AS IS" basis,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the Licence for the specific language governing permissions and
# limitations under the Licence.
#
# SPDX-FileCopyrightText: 2026 Technical University of Munich
# SPDX-License-Identifier: EUPL-1.2
# SPDX-FileContributor: Malte J. Ziebarth <malte.ziebarth@tum.de>


from typing import Iterable
from ..tools.views import join, enclose
from ..parse.function import Function
from ..language.abstract import Input, Output, Pointer, Const
from ..language.arrays import NDArray, QuantityArray
from ..language.strings import (
    ctypetype_string, type_string, cpp_type_string
)
from ..language.cpp import L2Type
from ..language.conversion import equivalent_cpp_type
from .variablegraph import (
    VariableGraph,
    _VG_Parameter # pyright: ignore[reportPrivateUsage]
)
from .parameter import (
    PythonInputParameter,
    PythonOutputParameter,
    CTypeParameterByValue,
    CTypeParameterByRef,
    TemporaryPythonVariable,
    CppParameter
)
from .input import wrap_input_parameter
from .out import wrap_output_parameter
from .array import (
    wrap_input_ndarray,
    wrap_output_ndarray,
    wrap_input_quantityarray,
    wrap_output_quantityarray
)
from .string import (
    add_string_input,
    add_string_output
)


_CTypeParam = CTypeParameterByRef | CTypeParameterByValue



class FunctionWrapper:
    """
    Wraps a `funktionsklempner.language.function.Function` parsed from
    a C++ header, and generates the initialization code.
    """
    function: Function
    vargraph: VariableGraph

    # Wrapping tasks: disentangle the C++ function parameters into
    # Python input arguments, temporary variables, CTypesIO-parameters,
    # and Python output variables.
    init_order: list[_VG_Parameter]
    python_input: list[PythonInputParameter]
    python_output: list[PythonOutputParameter]
    ctypes_parameters: list[CTypeParameterByValue | CTypeParameterByRef]
    cpp_parameters: list[CppParameter]

    def __init__(self,
            function: Function
        ):
        # Initialize the parameters:
        self.function = function
        self.vargraph = VariableGraph()

        def is_char_pointer(cpptype: L2Type) -> bool:
            if not isinstance(cpptype, Pointer):
                return False
            val = cpptype.val
            return val == Const("char") or val == "char"

        # Now parse the function
        for i, (category, name) in enumerate(function.cparams):
            cpptype = category.val
            if isinstance(category, Input):
                # Input parameter.
                if isinstance(cpptype, NDArray):
                    wrap_input_ndarray(
                        name,
                        Input(cpptype),
                        self.vargraph,
                        function.fk_namespace,
                        i
                    )
                elif isinstance(cpptype, QuantityArray):
                    wrap_input_quantityarray(
                        name,
                        Input(cpptype),
                        self.vargraph,
                        function.fk_namespace,
                        i
                    )
                elif is_char_pointer(cpptype) or cpptype == "std::string":
                    add_string_input(name, self.vargraph, None, i)
                else:
                    wrap_input_parameter(
                        name, Input(cpptype), self.vargraph, None, i
                    )

            else:
                # Output parameter.
                assert isinstance(category, Output)
                if isinstance(cpptype, NDArray):
                    wrap_output_ndarray(
                        name,
                        Output(cpptype),
                        self.vargraph,
                        function.size_control,
                        function.fk_namespace,
                        i
                    )
                elif isinstance(cpptype, QuantityArray):
                    wrap_output_quantityarray(
                        name,
                        Output(cpptype),
                        self.vargraph,
                        function.size_control,
                        function.fk_namespace,
                        i
                    )
                elif is_char_pointer(cpptype) or cpptype == "std::string":
                    add_string_output(name, self.vargraph, None, i)
                else:
                    wrap_output_parameter(
                        name, Output(cpptype), self.vargraph, None, i
                    )

        # Finally, resolve the parameter order:
        self.python_input = []
        self.python_output = []
        self.ctypes_parameters = []
        self.cpp_parameters = []
        self.init_order = self.vargraph.derive_initialization_order()
        for param in self.init_order:
            if isinstance(param, PythonInputParameter):
                self.python_input.append(param)
            elif isinstance(param, PythonOutputParameter):
                self.python_output.append(param)
            elif isinstance(param, _CTypeParam):
                self.ctypes_parameters.append(param)
            elif isinstance(param, CppParameter):
                self.cpp_parameters.append(param)
            else:
                assert isinstance(param, TemporaryPythonVariable)



    def compose_ctypes_function(self, cdll_name: str) -> Iterable[str]:
        """
        Declare the ctypes function as a part of the loaded CDLL.
        """
        fun = self.function
        yield f"_{fun.name} = {cdll_name}.FK_{fun.name}\n"
        yield f"_{fun.name}.argtypes = (\n    POINTER(c_uint64),\n"

        def cpp_type_string(p: CTypeParameterByValue | CTypeParameterByRef) -> str:
            """
            Here we handle the fact that using the two different types
            'CTypeParameterByValue' and 'CTypeParameterByRef' hides the outer
            pointer layer of the output parameters.
            When classifying the C++ function signature, we need to
            reintroduce this layer!
            """
            if isinstance(p, CTypeParameterByRef):
                # We save the output parameters without the outer pointer
                # layer. However, for correctly classifying the C++ library
                # function signature, we need to add this pointer here:
                return ctypetype_string(Pointer(p.dtype))

            return ctypetype_string(p.dtype)


        yield from (
            f"    {cpp_type_string(p)},\n"
            for p in self.ctypes_parameters
        )
        yield ")\n"


    def compose_python_body(self) -> Iterable[str]:
        """
        Compose the Python function body.
        """
        fun = self.function
        yield f"def {fun.name}("
        p_in = self.python_input
        p_out = self.python_output
        n_in: int = len(p_in)
        n_out: int = len(p_out)

        # Input arguments:
        out_indent = ""
        if n_in > 0:
            out_indent = "    "
            yield "\n"
            yield from (f"        {p.name}: {p.type_string()},\n" for p in p_in)


        # Output parameters:
        if n_out == 0:
            yield f"{out_indent}):\n"
        elif n_out == 1:
            yield f"{out_indent}) -> {type_string(p_out[0].dtype.val)}:\n"
        else:
            yield f"{out_indent}) -> tuple[\n"
            yield from (f"        {type_string(p.dtype.val)},\n" for p in p_out)
            yield "    ]:\n"

        # Function body:
        # First initialization of ctypes parameters:
        yield from (
            f"{l}\n" for p in self.init_order
            if not isinstance(p, CppParameter) and p.init_code is not None
            for l in p.init_code
        )

        # Call:
        yield f"    _call_with_error(\n        _{fun.name},\n"
        if n_in + n_out == 0:
            yield ")\n"
        else:
            yield from (
                f"        {p.to_cdll_argument()},\n"
                for p in self.ctypes_parameters)
            yield "    )\n"



        # Output parameter handling:
        yield from (
            "\n".join(l for l in p.post_call_code) + "\n" for p in p_out
            if p.post_call_code is not None
        )


        # Return statement:
        if n_out == 1:
            yield f"    return {p_out[0].name}\n"

        elif n_out > 1:
            yield "    return (\n"
            yield from (f"        {p.name},\n" for p in p_out)
            yield "    )\n"

        yield "\n"



    def compose_wrapper_c_args(self,
            indent: str,
            names: bool,
            types: bool
        ) -> Iterable[str]:
        """
        Compose the arguments
        """
        i_last: int = len(self.ctypes_parameters) - 1
        for i,param in enumerate(self.ctypes_parameters):
            yield indent
            if types:
                if isinstance(param, CTypeParameterByValue):
                    yield cpp_type_string(equivalent_cpp_type(param.dtype))
                else:
                    yield cpp_type_string(
                        Pointer(equivalent_cpp_type(param.dtype))
                    )
            if names:
                name = param.name
                if types:
                    yield " "
                yield name if isinstance(name, str) else name.name
            if i < i_last:
                yield ",\n"
            else:
                yield "\n"


    def compose_cpp_head(self,
            declaration: bool
        ) -> Iterable[str]:
        """
        This shall build a function signature like

        FUNKTIONSKLEMPNER_DLLEXPORT int FUNKTIONSKLEMPNER_CDECL
        FK_f_name(
            uint64_t* __err_oid,
            type arg0,
            type arg1
        );

        where the two macros are switched on/off via the `declaration`
        parameter.
        """
        if declaration:
            yield "FUNKTIONSKLEMPNER_DLLEXPORT int FUNKTIONSKLEMPNER_CDECL\nFK_"
        else:
            yield "int FK_"
        yield self.function.name

        # Early exit if we have no parameters:
        if len(self.ctypes_parameters) == 0:
            yield "(uint64_t* __err_oid\n)\n"
            return

        yield "(\n    uint64_t* __err_oid,\n"

        yield from self.compose_wrapper_c_args("    ", not declaration, True)
        if declaration:
            yield ");\n"
        else:
            yield ")\n"




    def compose_cpp_body(self) -> Iterable[str]:
        """

        """
        #
        # This shall build a function like
        #
        #   FUNKTIONSKLEMPNER_DLLEXPORT int FUNKTIONSKLEMPNER_CDECL
        #   FK_f_name(
        #       uint64_t* __err_oid,
        #       type arg0,
        #       type arg1
        #   )
        #   {
        #       return funktionsklempner::wrap_call(
        #           [](
        #               type arg0,
        #               type arg1
        #           ) -> int
        #           {
        #               f_name(
        #                   arg0,
        #                   arg1
        #               );
        #           },
        #           __err_oid,
        #           arg0,
        #           arg1
        #       );
        #   }

        yield from self.compose_cpp_head(False)
        yield (
            "{\n"
            "    return funktionsklempner::wrap_call(\n"
            "        [](\n"
        )
        yield from self.compose_wrapper_c_args(" " * 12, True, True)
        yield (
            "        )\n"
            "        {\n"
        )
        # Now initialize all the C++ arguments:
        yield from (
            s for p in self.cpp_parameters if p.init_code is not None
            for s in p.init_code
        )



        # Now call the function:
        yield "            "
        yield self.function.name
        yield "(\n"
        # Now the argument forwarding:
        yield from join(
            (f"                {p.to_cpp_argument()}"
             for p in sorted(self.cpp_parameters)
            ),
            ",\n"
        )
        yield (
            "\n"
            "            );\n"
        )
        # Now execute any post call code:
        yield from enclose(
            (
                s for p in self.cpp_parameters
                if p.post_call_code is not None for s in p.post_call_code
            ),
            "            ",
            ";\n"
        )
        # Now complete the wrap_call:
        yield (
            "        },\n"
            "        __err_oid"
        )
        if len(self.ctypes_parameters) > 0:
            yield ",\n"
        yield from self.compose_wrapper_c_args("        ", True, False)
        yield "\n    );\n}\n\n"
