# this code should be common among cwrap and ATen preprocessing
# for now, I have put it in one place but right now is copied out of cwrap

import copy
from typing import Any, Dict, Iterable, List, Union

Arg = Dict[str, Any]


def parse_arguments(args: List[Union[str, Arg]]) -> List[Arg]:
    new_args = []
    for arg in args:
        # Simple arg declaration of form "<type> <name>"
        if isinstance(arg, str):
            t, _, name = arg.partition(" ")
            new_args.append({"type": t, "name": name})
        elif isinstance(arg, dict):
            if "arg" in arg:
                arg["type"], _, arg["name"] = arg["arg"].partition(" ")
                del arg["arg"]
            new_args.append(arg)
        else:
            raise AssertionError()
    return new_args


Declaration = Dict[str, Any]


def set_declaration_defaults(declaration: Declaration) -> None:
    if "schema_string" not in declaration:
        # This happens for legacy TH bindings like
        # _thnn_conv_depthwise2d_backward
        declaration["schema_string"] = ""
    declaration.setdefault("arguments", [])
    declaration.setdefault("return", "void")
    if "cname" not in declaration:
        declaration["cname"] = declaration["name"]
    if "backends" not in declaration:
        declaration["backends"] = ["CPU", "CUDA"]
    assert "api_name" not in declaration
    declaration["api_name"] = declaration["name"]
    # NB: keep this in sync with gen_autograd.py
    if declaration.get("overload_name"):
        declaration["type_wrapper_name"] = "{}_{}".format(
            declaration["name"], declaration["overload_name"]
        )
    else:
        declaration["type_wrapper_name"] = declaration["name"]
    # TODO: Uggggh, parsing the schema string here, really???
    declaration["operator_name_with_overload"] = declaration["schema_string"].split(
        "("
    )[0]
    if declaration["schema_string"]:
        declaration["unqual_schema_string"] = declaration["schema_string"].split("::")[
            1
        ]
        declaration["unqual_operator_name_with_overload"] = declaration[
            "operator_name_with_overload"
        ].split("::")[1]
    else:
        declaration["unqual_schema_string"] = ""
        declaration["unqual_operator_name_with_overload"] = ""
    # Simulate multiple dispatch, even if it's not necessary
    if "options" not in declaration:
        declaration["options"] = [
            {
                "arguments": copy.deepcopy(declaration["arguments"]),
                "schema_order_arguments": copy.deepcopy(
                    declaration["schema_order_arguments"]
                ),
            }
        ]
        del declaration["arguments"]
        del declaration["schema_order_arguments"]
    # Parse arguments (some of them can be strings)
    for option in declaration["options"]:
        option["arguments"] = parse_arguments(option["arguments"])
        option["schema_order_arguments"] = parse_arguments(
            option["schema_order_arguments"]
        )
    # Propagate defaults from declaration to options
    for option in declaration["options"]:
        for k, v in declaration.items():
            # TODO(zach): why does cwrap not propagate 'name'? I need it
            # propagaged for ATen
            if k != "options":
                option.setdefault(k, v)


# TODO(zach): added option to remove keyword handling for C++ which cannot
# support it.

Option = Dict[str, Any]


def filter_unique_options(
    options: Iterable[Option],
    allow_kwarg: bool,
    type_to_signature: Dict[str, str],
    remove_self: bool,
) -> List[Option]:
    def exclude_arg(arg: Arg) -> bool:
        return arg["type"] == "CONSTANT"  # type: ignore[no-any-return]

    def exclude_arg_with_self_check(arg: Arg) -> bool:
        return exclude_arg(arg) or (remove_self and arg["name"] == "self")

    def signature(option: Option, num_kwarg_only: int) -> str:
        if num_kwarg_only == 0:
            kwarg_only_count = None
        else:
            kwarg_only_count = -num_kwarg_only
        arg_signature = "#".join(
            type_to_signature.get(arg["type"], arg["type"])
            for arg in option["arguments"][:kwarg_only_count]
            if not exclude_arg_with_self_check(arg)
        )
        if kwarg_only_count is None:
            return arg_signature
        kwarg_only_signature = "#".join(
            arg["name"] + "#" + arg["type"]
            for arg in option["arguments"][kwarg_only_count:]
            if not exclude_arg(arg)
        )
        return arg_signature + "#-#" + kwarg_only_signature

    seen_signatures = set()
    unique = []
    for option in options:
        # if only check num_kwarg_only == 0 if allow_kwarg == False
        limit = len(option["arguments"]) if allow_kwarg else 0
        for num_kwarg_only in range(0, limit + 1):
            sig = signature(option, num_kwarg_only)
            if sig not in seen_signatures:
                if num_kwarg_only > 0:
                    for arg in option["arguments"][-num_kwarg_only:]:
                        arg["kwarg_only"] = True
                unique.append(option)
                seen_signatures.add(sig)
                break
    return unique


def sort_by_number_of_args(declaration: Declaration, reverse: bool = True) -> None:
    def num_args(option: Option) -> int:
        return len(option["arguments"])

    declaration["options"].sort(key=num_args, reverse=reverse)


class Function(object):
    def __init__(self, name: str) -> None:
        self.name = name
        self.arguments: List["Argument"] = []

    def add_argument(self, arg: "Argument") -> None:
        assert isinstance(arg, Argument)
        self.arguments.append(arg)

    def __repr__(self) -> str:
        return self.name + "(" + ", ".join(a.__repr__() for a in self.arguments) + ")"


class Argument(object):
    def __init__(self, _type: str, name: str, is_optional: bool):
        self.type = _type
        self.name = name
        self.is_optional = is_optional

    def __repr__(self) -> str:
        return self.type + " " + self.name


def parse_header(path: str) -> List[Function]:
    with open(path, "r") as f:
        lines: Iterable[Any] = f.read().split("\n")

    # Remove empty lines and prebackend directives
    lines = filter(lambda l: l and not l.startswith("#"), lines)
    # Remove line comments
    lines = (l.partition("//") for l in lines)
    # Select line and comment part
    lines = ((l[0].strip(), l[2].strip()) for l in lines)
    # Remove trailing special signs
    lines = ((l[0].rstrip(");").rstrip(","), l[1]) for l in lines)
    # Split arguments
    lines = ((l[0].split(","), l[1]) for l in lines)
    # Flatten lines
    new_lines = []
    for l, c in lines:
        for split in l:
            new_lines.append((split, c))
    lines = new_lines
    del new_lines
    # Remove unnecessary whitespace
    lines = ((l[0].strip(), l[1]) for l in lines)
    # Remove empty lines
    lines = filter(lambda l: l[0], lines)
    generic_functions = []
    for l, c in lines:
        if l.startswith("TH_API void THNN_"):
            fn_name = l[len("TH_API void THNN_") :]
            if fn_name[0] == "(" and fn_name[-2] == ")":
                fn_name = fn_name[1:-2]
            else:
                fn_name = fn_name[:-1]
            generic_functions.append(Function(fn_name))
        elif l.startswith("TORCH_CUDA_CPP_API void THNN_"):
            fn_name = l[len("TORCH_CUDA_CPP_API void THNN_") :]
            if fn_name[0] == "(" and fn_name[-2] == ")":
                fn_name = fn_name[1:-2]
            else:
                fn_name = fn_name[:-1]
            generic_functions.append(Function(fn_name))
        elif l.startswith("TORCH_CUDA_CU_API void THNN_"):
            fn_name = l[len("TORCH_CUDA_CU_API void THNN_") :]
            if fn_name[0] == "(" and fn_name[-2] == ")":
                fn_name = fn_name[1:-2]
            else:
                fn_name = fn_name[:-1]
            generic_functions.append(Function(fn_name))
        elif l:
            t, name = l.split()
            if "*" in name:
                t = t + "*"
                name = name[1:]
            generic_functions[-1].add_argument(Argument(t, name, "[OPTIONAL]" in c))
    return generic_functions