# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name
"""Thin device-helper registration for TIRx intrinsic codegens.

Exposes one entry point: :func:`device_intrinsic`. Given an op name plus
``(helper_name, c_signature, body)`` (each a string or a ``(*args) -> str``
callable), it:

* wraps the body in
  ``__forceinline__ __device__ <ret> <name><sig> { <body> }``,
* registers a codegen function under the op name so
  ``call_intrin("", "tirx.<op_name>", *args)`` resolves to a call to that
  helper. TVM Op registration is static C++ only.

Args passed to the codegen are split into ``(forward_args, attr_args)``:
the trailing ``n_attrs`` are attrs (consumed by the ``helper_name`` /
``c_signature`` / ``body`` callables but **not** forwarded to the helper),
and the rest are operand args (forwarded). The default ``n_attrs=0`` means
every arg is forwarded — appropriate for fixed-arity ops with literal
``c_signature`` and ``helper_name``.

Coerce / validate attrs explicitly inside the callables — there is no
``Choice`` / ``Bool`` / ``IntAttr`` machinery; just call ``parse_str`` /
``int`` / ``bool`` on the raw arg as needed.
"""

from __future__ import annotations

from collections.abc import Callable

from tvm.backend.cuda.intrinsics.registry import register_codegen
from tvm.backend.cuda.op import cuda_func_call

# C primitive type → TVM dtype string. Used when the caller specifies a
# non-void ``return_type`` but no explicit ``tvm_return_type`` — the helper
# knows the TVM-side dtype from the C return type.
_C_TO_TVM_DTYPE = {
    "float": "float32",
    "double": "float64",
    "uint32_t": "uint32",
    "int32_t": "int32",
    "uint64_t": "uint64",
    "int64_t": "int64",
    "uint16_t": "uint16",
    "int16_t": "int16",
    "unsigned long long": "uint64",
    "long long": "int64",
    "unsigned short": "uint16",
    "float2": "float32x2",
    "bool": "bool",
    "unsigned int": "uint32",
    "int": "int32",
}


def device_intrinsic(
    op_name: str,
    *,
    helper_name: str | Callable | None = None,
    c_signature: str | Callable = "()",
    body: str | Callable,
    n_attrs: int = 0,
    return_type: str | Callable = "void",
    tvm_return_type: str | Callable | None = None,
    templated: bool = False,
    extra_deps: tuple = (),
) -> None:
    """Register a CUDA device-helper intrinsic.

    Parameters
    ----------
    op_name :
        Registry key — ``call_intrin("", "tirx.<op_name>", ...)`` resolves
        here. Also used as the default helper name (``tvm_builtin_<op_name>``)
        when ``helper_name`` is not provided.
    helper_name :
        Literal C function name, OR ``(*args) -> str`` to compute it from
        attr values. Defaults to ``f"tvm_builtin_{op_name}"``.
    c_signature :
        Literal C parameter list including outer parens (``"(int x, int y)"``),
        OR ``(*args) -> str`` to compute it from attr values. Defaults to
        ``"()"``.
    body :
        Literal C body string (already indented), OR ``(*args) -> str``.
    n_attrs :
        Number of trailing args that are attrs (consumed by ``helper_name``
        / ``c_signature`` / ``body`` callables, NOT forwarded to the helper
        as call arguments). The first ``len(args) - n_attrs`` args are the
        operand args forwarded to the helper.
    return_type :
        C return type. Default ``"void"``. Either a literal string or
        ``(*args) -> str`` when the helper return type depends on attrs.
    tvm_return_type :
        TVM dtype for the call result, when the helper has a non-void
        return. Either a literal string (``"int32"``) or ``(*args) -> str``.
        If omitted and ``return_type`` is non-void, it is auto-derived from
        the ``_C_TO_TVM_DTYPE`` table.
    templated :
        Prefix the helper with ``template <typename T>``.
    extra_deps :
        Helper-tag list (e.g. ``("get_tmem_addr",)``) forwarded as the second
        element of the codegen result so the header generator emits the
        prerequisite snippets.
    """
    if helper_name is None:
        helper_name = f"tvm_builtin_{op_name}"
    extra_deps = tuple(extra_deps)

    def codegen(*args):
        forward = args if n_attrs == 0 else args[:-n_attrs]
        name = helper_name(*args) if callable(helper_name) else helper_name
        sig = c_signature(*args) if callable(c_signature) else c_signature
        body_str = body(*args) if callable(body) else body
        ret_type = return_type(*args) if callable(return_type) else return_type
        prefix = "template <typename T>\n" if templated else ""
        source_code = (
            f"\n{prefix}__forceinline__ __device__ {ret_type} {name}{sig} {{\n{body_str}\n}}\n"
        )
        kwargs = {"source_code": source_code}
        if tvm_return_type is not None:
            kwargs["return_type"] = (
                tvm_return_type(*args) if callable(tvm_return_type) else tvm_return_type
            )
        elif ret_type != "void":
            kwargs["return_type"] = _C_TO_TVM_DTYPE.get(ret_type, ret_type)
        result = cuda_func_call(name, *forward, **kwargs)
        return (result, list(extra_deps)) if extra_deps else result

    codegen.__name__ = f"codegen_{op_name}"
    register_codegen(op_name)(codegen)
