# 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.
"""Wrapping existing transformations."""
# pylint: disable=invalid-name, unsupported-binary-operation

import enum
from collections.abc import Callable

import tvm_ffi as _ffi

from . import _ffi_api
from . import function_pass as _fpass


def Apply(ftransform):
    """Apply ftransform to each function in the Module.

    This function is a thin wrapper around tvm.tirx.transform.prim_func_pass

    Parameters
    ----------
    ftransform: tvm.tirx.PrimFunc -> tvm.tirx.PrimFunc
       The transformation pass.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """

    # pylint: disable=unused-argument
    def _transform(func, mod, ctx):
        return ftransform(func)

    return _fpass.prim_func_pass(_transform, opt_level=0, name="Apply")  # type: ignore


def VectorizeLoop(enable_vectorize: bool = True):
    """Lower vectorization loops.

    Parameters
    ----------
    enable_vectorize : bool
        Whether vectorization is enabled.
        Will lower to scalar loop when it is turned off.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.VectorizeLoop(enable_vectorize)  # type: ignore


def StorageRewrite():
    """Rewrite storage allocation pattern.

    Moves the allocation to outer most possible scope.
    Trying to share space between allocations to make
    a static allocation plan when possible.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.StorageRewrite()  # type: ignore


def InlinePrivateFunctions():
    """Inline calls to private functions

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.InlinePrivateFunctions()  # type: ignore


def PointerValueTypeRewrite():
    """
    Rewrite the pointer content type of arguments, as well as Alloc internal to the function to use
    the most frequently accessed type for load/store to avoid pointer casting in backend when
    possible.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.PointerValueTypeRewrite()  # type: ignore


@_ffi.register_object("tirx.transform.UnrollLoopConfig")
class UnrollLoopConfig(_ffi.Object):
    """Config for unroll loop pass"""


def UnrollLoop():
    """Unroll the constant loop marked by unroll.

    This pass also automatically attach pragma unroll tag to loops which meets the standard.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.UnrollLoop()  # type: ignore


@_ffi.register_object("tirx.transform.RemoveNoOpConfig")
class RemoveNoOpConfig(_ffi.Object):
    """Config for remove no op pass"""


def RemoveNoOp():
    """Remove No Op from the Stmt.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.RemoveNoOp()  # type: ignore


def RemoveAssume():
    """Remove all instances of builtin::assume

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.RemoveAssume()  # type: ignore


def BF16ComputeLegalize():
    """Legalize bf16 compute Ops.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.BF16ComputeLegalize()  # type: ignore


def FP8ComputeLegalize(promote_dtype: str = "float32"):
    """Legalize fp8 compute Ops.

    Parameters
    ----------
    promote_dtype : str
        The data type we promote fp8 to, options: float16/float32.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.FP8ComputeLegalize(promote_dtype)  # type: ignore


def BF16StorageLegalize():
    """Legalize bf16 storage types to u16.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.BF16StorageLegalize()  # type: ignore


def FP8StorageLegalize():
    """Legalize fp8 storage types to u8.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.FP8StorageLegalize()  # type: ignore


def CommonSubexprElim():
    """Replace redundant computations by new variables.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.CommonSubexprElim()  # type: ignore


@_ffi.register_object("tirx.transform.StmtSimplifyConfig")
class StmtSimplifyConfig(_ffi.Object):
    """Config for stmt simplify pass"""


def StmtSimplify():
    """Run statement-level arithmetic simplifications on the TIR PrimFunc.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.StmtSimplify()  # type: ignore


def ConvertSSA():
    """Convert an IRModule to be SSA form.

    This pass handles cases where the same `tirx.Var` appears in
    multiple functions within the same module.  For example, after
    extracting a fragment from one function into another, where the
    same `tirx.Var` may be defined both as within the body of the
    original function, and as a parameter within the hoisted function.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass

    """
    return _ffi_api.ConvertSSA()  # type: ignore


def MakePackedAPI():
    """Transform the PrimFuncs in the module to a packed func API.

    Prior to this pass, the PrimFunc may have Buffer arguments defined
    in the `PrimFuncNode::buffer_map`.  This pass consumes the
    `buffer_map`, using it to generate arguments that implement
    the packed based TVM FFI API.

    For static shapes, the `BufferNode::shape`, `BufferNode::strides`,
    and `BufferNode::elem_offset` member variables are used to
    generate runtime checks on the corresponding member variables in
    the user-provided `DLTensor*` or `tvm.runtime.tensor` argument.  (e.g. A
    PrimFunc that accepts a buffer of shape `[16,32]` validates that
    the `DLTensor::shape` array is `[16,32]`.)

    For dynamic Buffers, in which one or more of these `BufferNode` member
    variables use `tirx.Var` that are not defined by other PrimFunc
    parameters, these are instead used to define the variables based on
    the corresponding `DLTensor` members.  (e.g. A PrimFunc that accepts a
    buffer of shape `[tirx.Var("n"), tirx.Var("m")]`, when passed a
    `DLTensor` of shape `[16,32]`, will define `n = 16` and `n=32`, based
    on the argument's shape.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.MakePackedAPI()  # type: ignore


def SplitHostDevice():
    """Annotate, split, and lower host/device functions.

    This pass first annotates device regions within host functions,
    then splits them into host and device-side PrimFuncs, and finally
    lowers host-to-device calls into the device kernel launch ABI.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.SplitHostDevice()  # type: ignore


def SkipAssert():
    """Skip assert stmt.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.SkipAssert()  # type: ignore


def LowerWarpMemory():
    """Lower warp memory access to low-level device related function calls.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.LowerWarpMemory()  # type: ignore


def LowerTVMBuiltin():
    """Lower tvm builtin intrinsics.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.LowerTVMBuiltin()  # type: ignore


def LowerIntrin():
    """Lower target specific intrinsic calls.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.LowerIntrin()  # type: ignore


def NarrowDataType(target_bits: int):
    """Narrow down Expr datatype in stmt to target_bits.

    Parameters
    ----------
    target_bits : int
        The target bit configuration.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass

    Note
    ----
    Run this pass after FlattenBuffer.
    """
    return _ffi_api.NarrowDataType(target_bits)  # type: ignore


def ForceNarrowIndexToInt32():
    """Force narrow down indexing expressions and integer buffers to int32 dtype.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass

    Note
    ----
    This pass should not be used in default cases.
    """
    return _ffi_api.ForceNarrowIndexToInt32()  # type: ignore


def VerifyMemory():
    """Verify if func contains illegal host side direct memory access.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.VerifyMemory()  # type: ignore


@_ffi.register_object("s_tir.transform.HoistIfThenElseConfig")
class HoistIfThenElseConfig(_ffi.Object):
    """Config for hoist if then else pass"""


class HoistedConditionals(enum.Flag):
    """Flags for use in HoistExpressionConfig.conditional_types

    Each bitflag represents a type of expression that should be
    hoisted to the outermost loop possible.
    """

    Never = 0
    """ No hoisting of conditionals """

    IfElseStmt = 1
    """ If set, look for hoist candidates in IfElseStmt """

    IfElseExpr = 2
    """ If set, look for hoist candidates in tirx.if_then_else """

    BooleanExpression = 4
    """ If set, look for hoist candidates in all boolean expressions """

    UsingBlockVar = 8
    """ If set, allow hoisting of conditionals that use a block variable (e.g. threadIdx.x)  """

    All = IfElseStmt | IfElseExpr | BooleanExpression | UsingBlockVar
    """ Enable all hoisting of conditionals"""


class HoistedLetBindings(enum.Flag):
    """Flags for use in HoistExpressionConfig.let_binding_types

    Each bitflag represents a type of let binding expression that should be
    hoisted to the outermost loop possible.
    """

    Never = 0
    """ No hoisting of let bindings """

    RequiredByConditional = 1
    """ Bindings that are used by a hoisted conditional """

    Bind = 2
    """ Bindings occurring in Bind nodes """

    LetExpr = 4
    """ Bindings occurring in Let expressions """

    All = RequiredByConditional | Bind | LetExpr
    """ Enable all hoisting of let bindings """


@_ffi.register_object("s_tir.transform.HoistExpressionConfig")
class HoistExpressionConfig(_ffi.Object):
    """Config for hoist expression pass"""


def FlattenBuffer():
    """Flatten the multi-dimensional BufferLoad and BufferStore to single dimensional
    BufferLoad/BufferStore for the TIR not contains opaque block.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.FlattenBuffer()  # type: ignore


def BindTarget(target):
    """Annotate a PrimFunc with a given target.
    Parameters
    -------
    target : tvm.target.Target
        target

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.BindTarget(target)  # type: ignore


def AnnotateEntryFunc():
    """Set a PrimFunc as the entry point if it is only function in IRModule.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.AnnotateEntryFunc()  # type: ignore


def Filter(fcond: Callable):
    """Filter out PrimFuncs that does not satisfy the given condition.
    `fcond` should be a function that takes a primfunc and returns boolean.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.Filter(fcond)  # type: ignore


def TilePrimitiveDispatch():
    """Lower TIRx tile primitive calls through the active backend dispatch table.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.TilePrimitiveDispatch()  # type: ignore


def LowerTIRx():
    """Lower TIR to a lower-level IR.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.LowerTIRx()  # type: ignore


def LowerTIRxOpaque():
    """Lower opaque constructs in TIRX programs.

    Handles AllocBuffer lowering, For(thread_binding) to AttrStmt(thread_extent)
    conversion, unit loop elimination, and pragma annotation handling.
    This is the tirx-specific counterpart of s_tir.LowerOpaqueBlock,
    without any SBlock/SBlockRealize handling.

    Returns
    -------
    fpass : tvm.transform.Pass
        The result pass
    """
    return _ffi_api.LowerTIRxOpaque()  # type: ignore