# mypy: allow-untyped-defs
import collections
import inspect
import typing
from types import GenericAlias

import torch
from torch import device, dtype, Tensor, types
from torch.utils._exposed_in import exposed_in

from .opaque_object import (
    _resolve_opaque_type_info,
    is_opaque_reference_type,
    is_opaque_type,
)


# This is used as a negative test for
# test_custom_ops.py::TestTypeConversion::test_type_eval.
_TestTensor = torch.Tensor


@exposed_in("torch.library")
def infer_schema(
    prototype_function: typing.Callable,
    /,
    *,
    mutates_args,
    op_name: str | None = None,
    tags: torch.Tag | typing.Sequence[torch.Tag] | None = (),
) -> str:
    r"""Parses the schema of a given function with type hints. The schema is inferred from the
    function's type hints, and can be used to define a new operator.

    We make the following assumptions:

    * None of the outputs alias any of the inputs or each other.
    * | String type annotations "device, dtype, Tensor, types" without library specification are
      | assumed to be torch.*. Similarly, string type annotations "Optional, List, Sequence, Union"
      | without library specification are assumed to be typing.*.
    * | Only the args listed in ``mutates_args`` are being mutated. If ``mutates_args`` is "unknown",
      | it assumes that all inputs to the operator are being mutates.

    Callers (e.g. the custom ops API) are responsible for checking these assumptions.

    Args:
        prototype_function: The function from which to infer a schema for from its type annotations.
        op_name (Optional[str]): The name of the operator in the schema. If ``name`` is None, then the
            name is not included in the inferred schema. Note that the input schema to
            ``torch.library.Library.define`` requires an operator name.
        mutates_args ("unknown" | Iterable[str]): The arguments that are mutated in the function.
        tags (Tag | Sequence[Tag] | None): one or more tags to apply to the
            inferred schema. Use ``torch.Tag.inplace`` or ``torch.Tag.out`` to
            infer the conventional aliasing for those operator kinds.

    Returns:
        The inferred schema.

    Example:
        >>> def foo_impl(x: torch.Tensor) -> torch.Tensor:
        >>>     return x.sin()
        >>>
        >>> infer_schema(foo_impl, op_name="foo", mutates_args={})
        foo(Tensor x) -> Tensor
        >>>
        >>> infer_schema(foo_impl, mutates_args={})
        (Tensor x) -> Tensor
    """
    UNKNOWN_MUTATES = "unknown"
    pf_globals = prototype_function.__globals__
    pf_locals = None
    # TODO: Once our minimum version is py3.10+ pass `eval_str=True` to
    # inspect.signature() and we no longer need to deal with stringified
    # annotations below.
    sig = inspect.signature(prototype_function)
    if tags is None:
        tags = ()
    elif isinstance(tags, torch.Tag):
        tags = (tags,)
    is_inplace = torch.Tag.inplace in tags
    is_out = torch.Tag.out in tags
    if is_inplace and is_out:
        raise AssertionError(
            "torch.Tag.inplace and torch.Tag.out are mutually exclusive"
        )
    if type(mutates_args) is not str:
        mutates_args = tuple(mutates_args)

    def error_fn(what):
        raise ValueError(f"infer_schema(func): {what} Got func with signature {sig})")

    def convert_type_string(annotation_type: str):
        try:
            return eval(annotation_type, pf_globals, pf_locals)
        except NameError as e:
            error_fn(
                f"Unsupported type annotation {annotation_type}. It is not a type. "
                f"({e}). "
                f"If you are using 'from __future__ import annotations', note that "
                f"annotations are evaluated lazily as strings; make sure all types "
                f"referenced in annotations are importable at module scope, not only "
                f"inside a local function."
            )
        except Exception as e:
            error_fn(
                f"Unsupported type annotation {annotation_type}. It is not a type. "
                f"({type(e).__name__}: {e})"
            )

    def unstringify_types(
        tys: tuple[type[object] | str, ...],
    ) -> tuple[tuple[typing.Any, ...], bool]:
        res = []
        changed = False
        for ty in tys:
            ty, ty_changed = unstringify_type(ty)
            res.append(ty)
            changed |= ty_changed
        if changed:
            return tuple(res), True
        else:
            return tys, False  # type: ignore[return-value]

    def unstringify_type(ty: type[object] | str) -> tuple[typing.Any, bool]:
        # Dig through a generic type and if it contains a stringified type
        # convert that to a real type. The second return value indicates if the
        # type contained a string or not.
        if isinstance(ty, str):
            return convert_type_string(ty), True
        elif origin := typing.get_origin(ty):
            args, args_changed = unstringify_types(typing.get_args(ty))
            if args_changed:
                return GenericAlias(origin, args), True

        return ty, False

    params = []
    seen_args = set()
    saw_kwarg_only_arg = False
    first_positional_arg_name = None
    first_positional_arg_schema_type = None
    first_positional_arg_alias = None
    out_arg_aliases = []
    out_arg_names = []
    for idx, (name, param) in enumerate(sig.parameters.items()):
        if not supported_param(param):
            error_fn("We do not support positional-only args, varargs, or varkwargs.")

        if param.kind == inspect.Parameter.KEYWORD_ONLY:
            # The first time we see a kwarg-only arg, add "*" to the schema.
            if not saw_kwarg_only_arg:
                params.append("*")
                saw_kwarg_only_arg = True

        if param.annotation is inspect.Parameter.empty:
            error_fn(f"Parameter {name} must have a type annotation.")

        # The annotation might be converted to a string by annotation,
        # we convert it to the actual type.
        annotation_type, _ = unstringify_type(param.annotation)

        schema_type = None
        if annotation_type not in SUPPORTED_PARAM_TYPES:
            if is_opaque_type(annotation_type):
                schema_type = _resolve_opaque_type_info(annotation_type).class_name  # type: ignore[union-attr]
            elif annotation_type == torch._C.ScriptObject:
                error_fn(
                    f"Parameter {name}'s type cannot be inferred from the schema "
                    "as it is a ScriptObject. Please manually specify the schema "
                    "using the `schema=` kwarg with the actual type of the ScriptObject."
                )
            elif (
                hasattr(annotation_type, "__origin__")
                and annotation_type.__origin__ is tuple
            ):
                list_type = tuple_to_list(annotation_type)
                example_type_str = "\n\n"
                # Only suggest the list type if this type is supported.
                if list_type in SUPPORTED_PARAM_TYPES:
                    example_type_str = f"For example, {list_type}.\n\n"
                error_fn(
                    f"Parameter {name} has unsupported type {param.annotation}. "
                    f"We do not support Tuple inputs in schema. As a workaround, please try to use List instead. "
                    f"{example_type_str}"
                    f"The valid types are: {SUPPORTED_PARAM_TYPES.keys()}."
                )
            else:
                error_fn(
                    f"Parameter {name} has unsupported type {param.annotation}. "
                    f"The valid types are: {SUPPORTED_PARAM_TYPES.keys()}."
                )
        else:
            schema_type = SUPPORTED_PARAM_TYPES[annotation_type]

        if schema_type is None:
            raise AssertionError(f"schema_type is None for param {name}")

        if (
            param.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD
            and first_positional_arg_name is None
        ):
            first_positional_arg_name = name
            first_positional_arg_schema_type = schema_type
            first_positional_arg_alias = f"a{idx}"

        is_mutated = False
        if type(mutates_args) is str:
            if mutates_args != UNKNOWN_MUTATES:
                raise ValueError(
                    "mutates_args must either be a sequence of the names of "
                    "the arguments that are mutated or the string 'unknown'. "
                )
            if schema_type.startswith("Tensor"):
                is_mutated = True
                schema_type = f"Tensor(a{idx}!){schema_type[len('Tensor') :]}"
        elif name in mutates_args:
            is_mutated = True
            if not schema_type.startswith("Tensor"):
                error_fn(
                    f"Parameter {name} is in mutable_args but only Tensors or collections of Tensors can be mutated"
                )
            schema_type = f"Tensor(a{idx}!){schema_type[len('Tensor') :]}"
        if is_out and is_mutated:
            if param.kind != inspect.Parameter.KEYWORD_ONLY:
                error_fn("torch.Tag.out requires mutable arguments to be keyword-only.")
            if schema_type != f"Tensor(a{idx}!)":
                error_fn(
                    "torch.Tag.out only supports mutable keyword-only Tensor arguments."
                )
            out_arg_aliases.append(f"a{idx}")
            out_arg_names.append(name)
        seen_args.add(name)
        if param.default is inspect.Parameter.empty:
            # pyrefly: ignore [bad-argument-type]
            params.append(f"{schema_type} {name}")
        else:
            default_repr = None
            if param.default is None or isinstance(param.default, (int, float, bool)):
                default_repr = str(param.default)
            elif isinstance(param.default, (str, torch.device)):
                default_repr = f'"{param.default}"'
            elif isinstance(param.default, torch.dtype):
                dtype_repr = str(param.default)
                torch_dot = "torch."
                if not dtype_repr.startswith(torch_dot):
                    raise AssertionError(
                        f"dtype repr {dtype_repr!r} must start with 'torch.'"
                    )
                default_repr = dtype_repr[len(torch_dot) :]
            else:
                error_fn(
                    f"Parameter {name} has an unsupported default value type {type(param.default)}. "
                    f"Please file an issue on GitHub so we can prioritize this."
                )
            # pyrefly: ignore [bad-argument-type]
            params.append(f"{schema_type} {name}={default_repr}")
    if mutates_args != UNKNOWN_MUTATES:
        mutates_args_not_seen = set(mutates_args) - seen_args
        if len(mutates_args_not_seen) > 0:
            error_fn(
                f"{mutates_args_not_seen} in mutates_args were not found in "
                f"the custom op's signature. "
                f"mutates_args should contain the names of all args that the "
                f"custom op mutates, or just the string 'unknown' if you don't know."
            )
    if is_inplace:
        if first_positional_arg_name is None:
            error_fn("torch.Tag.inplace requires a positional Tensor argument.")
        if first_positional_arg_schema_type != "Tensor":
            error_fn(
                "torch.Tag.inplace requires the first positional argument to be a Tensor."
            )
        if type(mutates_args) is str or set(mutates_args) != {
            first_positional_arg_name
        }:
            error_fn(
                "torch.Tag.inplace requires mutates_args to contain exactly "
                f"the first positional argument, got {mutates_args}."
            )
    if is_out and len(out_arg_aliases) == 0:
        error_fn(
            "torch.Tag.out requires at least one mutable keyword-only Tensor argument."
        )
    return_annotation, _ = unstringify_type(sig.return_annotation)
    if is_out:
        ret = _infer_out_return_schema(
            return_annotation, out_arg_aliases, out_arg_names, error_fn
        )
    else:
        ret = parse_return(return_annotation, error_fn)
    if is_inplace:
        if ret != "Tensor":
            error_fn(
                "torch.Tag.inplace requires the return annotation to be torch.Tensor."
            )
        ret = f"Tensor({first_positional_arg_alias}!)"
    if op_name is not None:
        return f"{op_name}({', '.join(params)}) -> {ret}"
    return f"({', '.join(params)}) -> {ret}"


def derived_types(
    base_type: type | typing._SpecialForm,
    cpp_type: str,
    list_base: bool,
    optional_base_list: bool,
    optional_list_base: bool,
):
    result: list[tuple[type | typing._SpecialForm | GenericAlias, str]] = [
        (base_type, cpp_type),
        # pyrefly: ignore [not-a-type]
        (typing.Optional[base_type], f"{cpp_type}?"),  # noqa: UP045
    ]

    def derived_seq_types(typ: type | typing._SpecialForm):
        return (
            typing.Sequence[typ],  # type: ignore[valid-type]
            typing.List[typ],  # type: ignore[valid-type]  # noqa: UP006
            GenericAlias(collections.abc.Sequence, (typ,)),
            GenericAlias(list, (typ,)),
        )

    if list_base:
        result.extend(
            (seq_typ, f"{cpp_type}[]") for seq_typ in derived_seq_types(base_type)
        )
    if optional_base_list:
        result.extend(
            (seq_typ, f"{cpp_type}?[]")
            # pyrefly: ignore [not-a-type]
            for seq_typ in derived_seq_types(typing.Optional[base_type])  # noqa: UP045
        )
    if optional_list_base:
        result.extend(
            (typing.Optional[seq_typ], f"{cpp_type}[]?")  # noqa: UP045
            for seq_typ in derived_seq_types(base_type)
        )
    return result


def get_supported_param_types():
    data: list[tuple[type | typing._SpecialForm, str, bool, bool, bool]] = [
        # (python type, schema type, type[] variant, type?[] variant, type[]? variant
        (Tensor, "Tensor", True, True, False),
        (int, "SymInt", True, False, True),
        (float, "float", True, False, True),
        (bool, "bool", True, False, True),
        (str, "str", False, False, False),
        (types.Number, "Scalar", True, False, False),
        (dtype, "ScalarType", False, False, False),
        (device, "Device", False, False, False),
    ]

    if torch.distributed.is_available():
        from torch.distributed.distributed_c10d import GroupName

        data.append((typing.cast(type, GroupName), "str", False, False, False))

    result = []
    for line in data:
        result.extend(derived_types(*line))
    return dict(result)


SUPPORTED_RETURN_TYPES = {
    Tensor: "Tensor",
    typing.List[Tensor]: "Tensor[]",  # noqa: UP006
    list[Tensor]: "Tensor[]",
    int: "SymInt",
    float: "float",
    bool: "bool",
    types.Number: "Scalar",
}


def parse_return(annotation, error_fn):
    if annotation is None:
        return "()"

    if annotation is inspect.Parameter.empty:
        error_fn("No return type annotation was provided. Please add one.")

    origin = typing.get_origin(annotation)
    if origin is not tuple:
        if annotation not in SUPPORTED_RETURN_TYPES:
            if is_opaque_reference_type(annotation):
                return _resolve_opaque_type_info(annotation).class_name  # type: ignore[union-attr]
            error_fn(
                f"Return has unsupported type {annotation}. "
                f"The valid types are: {SUPPORTED_RETURN_TYPES}."
            )

        return SUPPORTED_RETURN_TYPES[annotation]

    args = typing.get_args(annotation)
    for arg in args:
        if arg not in SUPPORTED_RETURN_TYPES and not is_opaque_reference_type(arg):
            error_fn(
                f"Return has unsupported type {annotation}. "
                f"The valid types are: {SUPPORTED_RETURN_TYPES}."
            )

    def _return_type_str(arg):
        if ty := SUPPORTED_RETURN_TYPES.get(arg):
            return ty
        return _resolve_opaque_type_info(arg).class_name  # type: ignore[union-attr]

    output_ty = ", ".join(_return_type_str(arg) for arg in args)

    # use (()) to represent tuple with single element
    if len(args) == 1:
        output_ty = "(" + output_ty + ")"
    return "(" + output_ty + ")"


def _infer_out_return_schema(
    return_annotation,
    out_arg_aliases,
    out_arg_names,
    error_fn,
):
    if len(out_arg_aliases) == 1:
        return_annotation_matches_out_args = return_annotation is Tensor
    else:
        return_annotation_matches_out_args = typing.get_origin(
            return_annotation
        ) is tuple and typing.get_args(return_annotation) == (Tensor,) * len(
            out_arg_aliases
        )

    if not return_annotation_matches_out_args:
        error_fn(
            "torch.Tag.out requires the return annotation to match the "
            f"mutable keyword-only Tensor arguments {tuple(out_arg_names)}."
        )

    aliased_returns = [f"Tensor({out_arg_alias}!)" for out_arg_alias in out_arg_aliases]
    if len(aliased_returns) == 1:
        return aliased_returns[0]
    return f"({', '.join(aliased_returns)})"


SUPPORTED_PARAM_TYPES = get_supported_param_types()


def supported_param(param: inspect.Parameter) -> bool:
    return param.kind in (
        inspect.Parameter.POSITIONAL_OR_KEYWORD,
        inspect.Parameter.KEYWORD_ONLY,
    )


def tuple_to_list(tuple_type: type[tuple]) -> type[list]:
    """
    Convert `tuple_type` into a list type with the same type arguments. Assumes that `tuple_type` is typing.Tuple type.
    """
    type_args = typing.get_args(tuple_type)
    if (
        tuple_type is typing.Tuple  # noqa: UP006
        or tuple_type is tuple
        or not type_args
    ):
        # Handle the case of an empty tuple type
        return list
    elif len(type_args) == 1:
        # General case: create a List with the same type arguments
        return list[type_args[0]]  # type: ignore[valid-type]
    elif len(type_args) == 2 and type_args[1] is Ellipsis:
        return list[type_args[0]]  # type: ignore[valid-type]
    else:
        return list[typing.Union[tuple(type_args)]]  # type: ignore[misc, return-value]  # noqa: UP007