# 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
"""Relax vm primitives."""

from tvm.ir import Call

from ...expr import DataTypeImm, Expr, StringImm, Tuple, prim_value
from ...utils import convert_to_expr
from . import _ffi_api


def alloc_storage(
    shape: Expr,
    runtime_device_index: int | Expr,
    dtype: str | Expr,
    storage_scope: str | StringImm = "global",
) -> Call:
    """Construct a Call to allocate a storage with specific size,
    runtime_device_index, and dtype.

    Parameters
    ----------
    shape : Expr
        The shape of the storage to be allocated.

    runtime_device_index : Union[int, Expr]
        The device index indicating on which device the tensor is to
        be allocated at runtime. Index -1 is reserved for the host device.

    dtype : Union[str, Expr]
        The datatype of the storage to be allocated.

    storage_scope : Union[str, StringImm]
        The storage scope of the storage to allocate. Default is global.

    Returns
    -------
    result : Call
        A relax Call, which gets the allocated storage.
    """
    shape = convert_to_expr(shape)
    if isinstance(dtype, str):
        dtype = DataTypeImm(dtype)
    if isinstance(storage_scope, str):
        storage_scope = StringImm(storage_scope)
    if isinstance(runtime_device_index, int):
        runtime_device_index = prim_value(runtime_device_index)
    return _ffi_api.alloc_storage(shape, runtime_device_index, dtype, storage_scope)  # type: ignore


def alloc_tensor(
    storage: Expr,
    offset: int | Expr,
    shape: Expr,
    dtype: str | Expr,
    runtime_device_ind: int | Expr = prim_value(0),
) -> Call:
    """Construct a Call to allocate a tensor on a certain storage starting from the given offset.

    Parameters
    ----------
    storage : Expr
        The storage to allocate the tensor to.

    offset : Union[int, Expr]
        The storage offset to allocate the tensor.

    shape : Expr
        The shape of the tensor to be allocated.

    dtype : Union[str, Expr]
        The datatype of the tensor to be allocated.

    runtime_device_ind: Union[int, Expr]
        The device index indicating on which device the tensor is to be
        allocated at runtime. Index -1 is reserved for the host device.

    Returns
    -------
    result : Call
        A relax Call, which gets the allocated tensor.
    """
    if isinstance(offset, int):
        offset = prim_value(offset)
    shape = convert_to_expr(shape)
    if isinstance(dtype, str):
        dtype = DataTypeImm(dtype)
    if isinstance(runtime_device_ind, int):
        runtime_device_ind = prim_value(runtime_device_ind)
    return _ffi_api.alloc_tensor(storage, offset, shape, dtype, runtime_device_ind)  # type: ignore


def kill_object(obj: Expr) -> Call:
    """Construct a Call to set the register corresponding to the input object to
    null at runtime, in order to kill the input object.

    Parameters
    ----------
    obj : Expr
        The object to be killed.

    Returns
    -------
    result : Call
        CallNode that kills the input object.
    """
    return _ffi_api.kill_object(obj)  # type: ignore


def call_tir_dyn(func: Expr, args: Tuple) -> Call:
    """Construct a Call to call_tir_dyn (invoke the given TIR PrimFunc)
    consisting of the input tensors and the shape of the result.

    Parameters
    ----------
    func : Expr
        An expression evaluating to a TIR PrimFunc.

    args : Tuple
        The input args, includes a list of tensors, and a ShapeExpr.

    Returns
    -------
    result : Call
        A relax Call to call_tir_dyn.
    """
    func = convert_to_expr(func)
    if isinstance(args, list | tuple):
        args = Tuple(args)

    return _ffi_api.call_tir_dyn(func, args)  # type: ignore