# Data model for representing parameters in various parts of the interface.
#
# Author: Malte J. Ziebarth (malte.ziebarth@tum.de)
#
# Copyright (C) 2025 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: 2025 Technical University of Munich
# SPDX-License-Identifier: EUPL-1.2
# SPDX-FileContributor: Malte J. Ziebarth <malte.ziebarth@tum.de>

from abc import ABCMeta, abstractmethod
from ..language.arrays import NDArray
from ..language.python import is_python_type, is_ctype_type, assert_ctypes_type,\
    PythonInputType, PythonOutputType, P2Type, CTypeType
from ..language.cpp import CppArgumentType, is_cpp_type, L1Type
from ..language.conversion import NUMPY_DTYPE_MAP, BASIC_PYTHON_DTYPE_MAP, \
    equivalent_cpp_type
from ..language.abstract import Input, Output
from ..language.strings import type_string


class Forward:
    """
    Name forwarding.
    """
    name: str
    varname: str | None

    def __init__(self, name: str):
        self.name = name
        self.varname = None



class Parameter[T, NameType: str | Forward](metaclass=ABCMeta):
    """
    Parameter base class.
    """
    name: NameType
    dtype: T
    init_code: list[str] | None

    def __init__(self, name: NameType, dtype: T, init_code: list[str] | None):
        assert isinstance(name, (str, Forward))
        self.name = name
        if not self._check_dtype(dtype):
            raise TypeError(
                f"Failed the data type check for parameter {name} with "
                f"dtype='{dtype}'"
            )
        self.dtype = dtype
        self.init_code = init_code


    @abstractmethod
    def _check_dtype(self, dtype: T) -> bool:
        pass

    @staticmethod
    @abstractmethod
    def parameter_type_name() -> str:
        pass



#
# 1. Python input parameters.
# ===========================
#


class PythonInputParameter(Parameter[PythonInputType, str]):
    """
    An input parameter of the top-level Python function
    wrapper.
    """
    def __init__(self,
            name: str,
            dtype: PythonInputType,
            init_code: list[str] | None = None
        ):
        super().__init__(name, dtype, init_code)

    def _check_dtype(self, dtype: PythonInputType) -> bool:
        if not is_python_type(dtype):
            raise RuntimeError(
                f"Checking data type for PythonInputParameter {self.name} "
                f"failed with dtype='{dtype}'"
            )
        return True


    def __repr__(self) -> str:
        return f"PythonInputParameter(name={self.name}, dtype={self.dtype})"


    def type_string(self) -> str:
        """
        Return the type string for representing this parameter
        in the wrapper's function arguments.
        """
        # First check whether it's an array:
        dtype = self.dtype.val
        if isinstance(dtype, NDArray):
            nd_dtype = dtype.ctypes_type()
            if nd_dtype in BASIC_PYTHON_DTYPE_MAP:
                return (
                    f"'Iterable[{BASIC_PYTHON_DTYPE_MAP[nd_dtype]}] | "
                    f"NDArray[{NUMPY_DTYPE_MAP[nd_dtype]}]'"
                )
        # Otherwise check whether we should prevent pythonification:
        return type_string(dtype)


    @staticmethod
    def parameter_type_name() -> str:
        return "Python input parameter"


class TemporaryPythonVariable(Parameter[P2Type, str]):
    """
    A temporary python variable.
    """

    def __init__(self,
            name: str,
            #init_from: "PythonInputParameter | COutputParameter",
            dtype: P2Type,
            init_code: list[str] | None = None
        ):

        super().__init__(name, dtype, init_code)

    def _check_dtype(self, dtype: P2Type) -> bool:
        return is_python_type(dtype)


    def __repr__(self) -> str:
        return f"TemporaryPythonVariable(name={self.name}, dtype={self.dtype})"


    @staticmethod
    def parameter_type_name() -> str:
        return "temporary Python parameter"



#
# 2. CTypes input parameters.
# ===========================
# These are the input-role parameters that get passed to the
# C interface.
#

def derive_ctypes_name(
        py_name: str,
        given_names: set[str],
        prefix: str | None = None
    ) -> str:
    """
    Naming scheme for CTypes names derived from Python names.
    """
    if py_name.startswith("c_"):
        name = py_name
    else:
        name = "c_" + py_name
    # Prepend a prefix if given:
    if prefix is not None:
        name = f"{prefix}_{name}"
    while name in given_names:
        name = "_" + name
    return name


class CParameter(Parameter[CTypeType, str | Forward]):
    """
    A ctypes parameter.
    """
    def _check_dtype(self, dtype: CTypeType) -> bool:
        return is_ctype_type(dtype)

    @staticmethod
    def parameter_type_name() -> str:
        return "C parameter"



class CTypeParameterByValue(CParameter):
    """
    A ctypes input parameter.
    """
    _cdll_argument_name: str

    def __init__(self,
            name: str | Forward,
            dtype: CTypeType,
            init_code: list[str],
            cdll_argument_name: str | None = None
        ):
        assert isinstance(init_code, list)

        super().__init__(name, dtype, init_code)

        if cdll_argument_name is None:
            if isinstance(name, Forward):
                self._cdll_argument_name = name.name
            else:
                self._cdll_argument_name = name
        else:
            assert isinstance(cdll_argument_name, str)
            self._cdll_argument_name = cdll_argument_name


    def __repr__(self) -> str:
        return f"CTypeParameterByValue(name={self.name}, dtype={self.dtype})"


    def to_cdll_argument(self) -> str:
        """
        Returns the argument to be passed to the CDLL.
        """
        return self._cdll_argument_name


    @staticmethod
    def forward(
            pyparam: "PythonInputParameter"
        ) -> "CTypeParameterByValue":
        ctype_type = assert_ctypes_type(pyparam.dtype.val)

        return CTypeParameterByValue(
            Forward(pyparam.name),
            ctype_type,
            [],
            pyparam.name
        )


    # @staticmethod
    # def forward(pip: PythonInputParameter) -> "CTypeParameterByValue":
    #     """
    #     Forwards a Python input parameter.
    #     """
    #     # Assert it's a ctypes type and create the input parameter
    #     # with a dummy name:
    #     dtype: PythonInputType = pip.dtype
    #     ctype: CTypeType = assert_ctypes_type(dtype.val)
    #     cip: CTypeParameterByValue = CTypeParameterByValue(
    #         pip,
    #         ctype,
    #         set()
    #     )

    #     # Now override the name with the Python input parameter
    #     # name. Into the init code, insert an assert for the
    #     # correct type:
    #     cip.name = pip.name
    #     cip.init_code = [
    #         f"    assert isinstance({pip.name}, {ctypetype_string(ctype)})"
    #     ]

    #     return cip


class CTypeParameterByRef(CParameter):
    """
    A ctypes input parameter.
    """
    _cdll_argument_name: str

    def __init__(self,
            name: str | Forward,
            dtype: CTypeType,
            init_code: list[str],
            cdll_argument_name: str | None = None
        ):
        assert isinstance(init_code, list)

        super().__init__(name, dtype, init_code)

        # For argument forwarding, one may specify
        # a specific name as an argument name:
        if cdll_argument_name is None:
            if isinstance(name, Forward):
                self._cdll_argument_name = name.name
            else:
                self._cdll_argument_name = name
        else:
            assert isinstance(cdll_argument_name, str)
            self._cdll_argument_name = cdll_argument_name


    def __repr__(self) -> str:
        return f"CTypeParameterByRef(name={self.name}, dtype={self.dtype})"


    def to_cdll_argument(self) -> str:
        """
        Returns the argument to be passed to the CDLL.
        """
        return f"byref({self._cdll_argument_name})"

    @staticmethod
    def forward(
            pyparam: "PythonOutputParameter | PythonInputParameter"
        ) -> "CTypeParameterByRef":
        ctype_type = assert_ctypes_type(pyparam.dtype.val)

        return CTypeParameterByRef(
            Forward(pyparam.name),
            ctype_type,
            [],
            pyparam.name
        )


#
# Within the C++ wrapper code, we need to transform
# the multiplexed CType parameters back into C++ parameters.
#
class CppParameter(Parameter[CppArgumentType, str | Forward]):
    """
    A ctypes parameter.
    """
    _cpp_argument_name: str
    _cpp_argument_id: int
    post_call_code: list[str] | None

    def __init__(self,
            name: str | Forward,
            dtype: CppArgumentType,
            init_code: list[str],
            post_call_code: list[str] | None,
            cpp_argument_name: str,
            cpp_argument_id: int
        ):
        assert isinstance(init_code, list)

        super().__init__(name, dtype, init_code)

        # The C++ argument name is mandatory since it coincides
        # with the Python argument name (which is already given in
        # the variable graph)
        assert isinstance(cpp_argument_name, str)
        self._cpp_argument_name = cpp_argument_name

        # The C++ argument id is mandatory since it determines the
        # order in which the arguments shall be passed to
        # the user-provided function.
        assert isinstance(cpp_argument_id, int)
        self._cpp_argument_id = cpp_argument_id

        # Optional post call code:
        assert isinstance(post_call_code, list) or post_call_code is None
        self.post_call_code = post_call_code


    def __repr__(self) -> str:
        return (
            f"CppParameter(name={self.name}, dtype={self.dtype}, "
            f"argname={self._cpp_argument_name})"
        )


    def _check_dtype(self, dtype: CppArgumentType) -> bool:
        return is_cpp_type(dtype)

    def __lt__(self, other: "CppParameter") -> bool:
        return self._cpp_argument_id < other._cpp_argument_id

    @staticmethod
    def parameter_type_name() -> str:
        return "C++ parameter"

    @staticmethod
    def _forward_type_name(
            cparam: CTypeParameterByRef | CTypeParameterByValue
        ) -> tuple[L1Type, str]:
        ctype_type = assert_ctypes_type(cparam.dtype)
        cpp_type = equivalent_cpp_type(ctype_type)
        name = cparam.name
        if isinstance(name, Forward):
            name = name.name

        return cpp_type, name


    @staticmethod
    def forward_input(
            cparam: CTypeParameterByRef | CTypeParameterByValue,
            cpp_argument_id: int
        ) -> "CppParameter":
        cpp_type, name = CppParameter._forward_type_name(cparam)

        return CppParameter(
            Forward(name),
            Input(cpp_type),
            [],
            None,
            name,
            cpp_argument_id
        )


    @staticmethod
    def forward_output(
            cparam: CTypeParameterByRef | CTypeParameterByValue,
            cpp_argument_id: int
        ) -> "CppParameter":
        cpp_type, name = CppParameter._forward_type_name(cparam)

        return CppParameter(
            Forward(name),
            Output(cpp_type),
            [],
            None,
            name,
            cpp_argument_id
        )


    def to_cpp_argument(self) -> str:
        """
        Returns the argument to be passed to the user defined function.
        """
        return self._cpp_argument_name




#
# 4. Python output parameters.
# ===========================
# These are the final output parameters to the Python world.
#

class PythonOutputParameter(Parameter[PythonOutputType, str]):
    """
    A Python output parameter
    """
    # Code that shall be performed after calling the function:
    post_call_code: list[str] | None

    def __init__(self,
            name: str,
            dtype: PythonOutputType,
            init_code: list[str],
            post_call_code: list[str] | None
        ):
        super().__init__(name, dtype, init_code)
        self.post_call_code = post_call_code


    def __repr__(self) -> str:
        return f"PythonOutputParameter(name={self.name}, dtype={self.dtype})"


    def _check_dtype(self, dtype: PythonOutputType) -> bool:
        return is_python_type(dtype)


    @staticmethod
    def parameter_type_name() -> str:
        return "Python output parameter"
