from typing import Mapping
from typing import Any
import copy


def enhance_function_signatures(specification: Mapping[str, Any]) -> Mapping[str, Any]:
    """Enhance function signatures

    Add required and optional objects to signatures objects for semantic validation
    support.

    Args:
        specification (Mapping[str, Any]): bel specification dictionary

    Returns:
        Mapping[str, Any]: return enhanced bel specification dict
    """

    for func in specification["functions"]["signatures"]:

        # Add primary parent functions to modifier functions
        if specification["functions"]["signatures"][func]["func_type"] == "modifier":
            specification["functions"]["signatures"][func]["primary_function"] = specification[
                "functions"
            ]["info"][func]["primary_function"]

        for i, sig in enumerate(specification["functions"]["signatures"][func]["signatures"]):
            args = sig["arguments"]
            req_args = []
            pos_args = []
            opt_args = []
            mult_args = []

            for arg in args:
                # Multiple argument types
                if arg.get("multiple", False):
                    if arg["type"] in ["Function", "Modifier"]:
                        mult_args.extend(arg.get("values", []))
                    elif arg["type"] in ["NSArg"]:
                        # Complex and Composite signature has this
                        mult_args.extend(arg.get("values", []))
                    elif arg["type"] in ["StrArgNSArg", "StrArg"]:

                        mult_args.append(arg["type"])

                # Optional, position dependent - will be added after req_args based on order in bel_specification
                elif arg.get("optional", False) and arg.get("position", False):
                    if arg["type"] in ["Function", "Modifier"]:
                        pos_args.append(arg.get("values", []))
                    elif arg["type"] in ["StrArgNSArg", "NSArg", "StrArg"]:
                        pos_args.append(arg["type"])

                # Optional, position independent
                elif arg.get("optional", False):
                    if arg["type"] in ["Function", "Modifier"]:
                        opt_args.extend(arg.get("values", []))
                    elif arg["type"] in ["StrArgNSArg", "NSArg", "StrArg"]:
                        opt_args.append(arg["type"])

                # Required arguments, position dependent
                else:
                    if arg["type"] in ["Function", "Modifier"]:
                        req_args.append(arg.get("values", []))
                    elif arg["type"] in ["StrArgNSArg", "NSArg", "StrArg"]:
                        req_args.append(arg["type"])

            specification["functions"]["signatures"][func]["signatures"][i][
                "req_args"
            ] = copy.deepcopy(req_args)
            specification["functions"]["signatures"][func]["signatures"][i][
                "pos_args"
            ] = copy.deepcopy(pos_args)
            specification["functions"]["signatures"][func]["signatures"][i][
                "opt_args"
            ] = copy.deepcopy(opt_args)
            specification["functions"]["signatures"][func]["signatures"][i][
                "mult_args"
            ] = copy.deepcopy(mult_args)

    return specification