# 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.
"""Abstraction for array data structures."""

import functools
from numbers import Integral

import tvm_ffi

import tvm
from tvm.ir import PointerType, PrimType, Range
from tvm.runtime import Object, Scriptable, convert

from . import _ffi_api


@tvm_ffi.register_object("tirx.Buffer")
class Buffer(Object, Scriptable):
    """Symbolic data buffer in TVM.

    Buffer provide a way to represent data layout
    specialization of data structure in TVM.

    Do not construct directly, use :py:func:`~decl_buffer` instead.
    See the documentation of :py:func:`decl_buffer` for more details.

    See Also
    --------
    decl_buffer : Declare a buffer
    """

    READ = 1
    WRITE = 2

    def access_ptr(self, access_mask, ptr_type="handle", content_lanes=1, offset=0, extent=None):
        """Get an access pointer to the head of buffer.

        This is the recommended method to get buffer data
        ptress when interacting with external functions.

        Parameters
        ----------
        access_mask : int
            The access pattern MASK. Indicate whether the
            access will read or write to the data content.

        ptr_type : str or tvm.ir.Type, optional
            The data type of the result pointer. Do not specify
            unless we want to cast pointer to specific type.

        content_lanes: int, optional
            The number of lanes for the data type. This value
            is greater than one for vector types.

        offset: Expr, optional
            The offset of pointer. We can use it to offset by
            the number of elements from the address of ptr.

        extent: Expr, optional
            The extent of pointer.

        Examples
        --------
        .. code-block:: python

          # Get access ptr for read
          buffer.access_ptr("r")
          # Get access ptr for read/write with bitmask
          buffer.access_ptr(Buffer.READ | Buffer.WRITE)
          # Get access ptr for read/write with str flag
          buffer.access_ptr("rw")
          # Get access ptr for read with offset
          buffer.access_ptr("r", offset = 100)
          # Get access ptr for read with extent
          buffer.access_ptr("r", extent = 100)
        """
        if isinstance(access_mask, str):
            mask = 0
            for value in access_mask:
                if value == "r":
                    mask = mask | Buffer.READ
                elif value == "w":
                    mask = mask | Buffer.WRITE
                else:
                    raise ValueError(f"Unknown access_mask {access_mask}")
            access_mask = mask
        if isinstance(ptr_type, str):
            ptr_type = (
                PointerType(PrimType("void"))
                if ptr_type == "handle"
                else PointerType(PrimType(ptr_type))
            )
        elif isinstance(ptr_type, PrimType):
            ptr_type = PointerType(ptr_type)
        offset = convert(offset)
        extent = convert(extent)
        return _ffi_api.BufferAccessPtr(
            self,
            access_mask,
            ptr_type,
            content_lanes,
            offset,
            extent,  # type: ignore
        )

    def vload(self, begin, dtype=None, predicate=None):
        """Generate an Expr that loads dtype from begin index.

        Parameters
        ----------
        begin : Array of Expr
            The beginning index in unit of Buffer.dtype

        dtype : str
            The data type to be loaded,
            can be vector type which have lanes that is multiple of Buffer.dtype

        predicate : Optional[Expr]
            A vector mask of boolean values indicating which lanes of a vector are to be
            loaded. The number lanes of the mask must be equal to the number of lanes being loaded.

        Returns
        -------
        load : Expr
            The corresponding load expression.
        """
        begin = (begin,) if isinstance(begin, int) or tvm.ir.is_prim_expr(begin) else begin
        dtype = dtype if dtype else self.dtype
        return _ffi_api.BufferVLoad(self, begin, dtype, predicate)  # type: ignore

    def vstore(self, begin, value, predicate=None):
        """Generate a Stmt that store value into begin index.

        Parameters
        ----------
        begin : Array of Expr
            The beginning index in unit of Buffer.dtype

        value : Expr
            The value to be stored.

        predicate : Optional[Expr]
            A vector mask of boolean values indicating which lanes of a vector are to be
            stored. The number lanes of the mask must be equal to the number of lanes in
            value.

        Returns
        -------
        store : Stmt
            The corresponding store stmt.
        """
        begin = (begin,) if isinstance(begin, int) or tvm.ir.is_prim_expr(begin) else begin
        return _ffi_api.BufferVStore(self, begin, value, predicate)  # type: ignore

    def scope(self):
        """Return the storage scope associated with this buffer.
        Returns
        -------
        scope : str
            The storage scope associated with this buffer.
        """
        return _ffi_api.BufferStorageScope(self)  # type: ignore

    def get_flattened_buffer(self):
        """Generate a Buffer that is a flattened version of this buffer.

        Returns
        -------
        flattened : Buffer
            The corresponding flat buffer.
        """
        return _ffi_api.BufferGetFlattenedBuffer(self)  # type: ignore

    def with_allocated_addr(self, allocated_addr):
        """Return a new buffer with the allocated address."""
        return _ffi_api.BufferWithAllocatedAddr(self, allocated_addr)  # type: ignore

    def with_dtype(self, dtype):
        """Return a new buffer with the dtype."""
        return _ffi_api.BufferWithDtype(self, dtype)  # type: ignore

    def with_data(self, data):
        """Return a new buffer with the data."""
        return _ffi_api.BufferWithData(self, data)  # type: ignore

    def offset_of(self, indices):
        """Determine the offset of the provided indices in the flattened buffer.

        Parameters
        ----------
        indices : Union[Expr, List[Expr]]

            The indices of the element in the original buffer.

        Returns
        -------
        flattened_indices: List[Expr]

            The offset indices of the element in the flattened buffer.
        """
        return _ffi_api.BufferOffsetOf(self, indices)  # type: ignore

    @property
    def byte_offset(self):
        """Get the byte offset of the buffer."""
        return self.elem_offset * tvm.DataType(self.dtype).bits // 8

    def elem_offset_of(self, indices, inner=True):
        """Get the element offset of the buffer at the given indices.
        Note that indices subject to buffer's layout mapping.

        Parameters
        ----------
        indices : Union[Expr, List[Expr]]
            The indices of the element in the original buffer.

        inner : bool, optional
            If False, the offset is relative to the original buffer.
            Default is True.

        Returns
        -------
        offset: Expr
            The element offset of the buffer at the given indices.
        """
        if inner:
            return _ffi_api.BufferOffsetOfp(self, indices)
        return self.elem_offset + _ffi_api.BufferOffsetOfp(self, indices)

    def byte_offset_of(self, indices, inner=True):
        """Get the byte offset of the buffer at the given indices.
        Note that indices subject to buffer's layout mapping.

        Parameters
        ----------
        indices : Union[Expr, List[Expr]]
            The indices of the element in the original buffer.

        inner : bool, optional
            If False, the offset is relative to the original buffer.
            Default is True.

        Returns
        -------
        offset: Expr
            The byte offset of the buffer at the given indices.
        """
        return self.elem_offset_of(indices, inner) * tvm.DataType(self.dtype).bits // 8

    def is_scalar(self, alloc_or_decl=True):
        """Check if the buffer is a scalar.

        Parameters
        ----------
        alloc_or_decl : bool, optional
            Whether to consider alloc_scalar and decl_scalar as scalar. True for alloc_scalar,
            False for decl_scalar.

        Returns
        -------
            bool: True if the buffer is a scalar, False otherwise.
        """
        return _ffi_api.BufferIsScalar(self, alloc_or_decl)

    def ptr_to(self, indices):
        """Get the pointer to the buffer at the given indices (logical indices).

        Note that the bufferload inside requires LowerTIPp pass to apply the layout to get the physical indices.
        """  # noqa: E501
        assert len(indices) == len(self.shape), (
            f"The number of indices {indices} does not match the shape of the buffer {self.shape}"
        )
        return tvm.tirx.address_of(self[tuple(indices)])

    def view(self, *args, **kwargs) -> "Buffer":
        """Creates a new view of the buffer. (used by parser)

        Supported signatures are ``view(*shape, layout=None)``, where shape can contain
        ``-1`` to indicate that the dimension size is auto-inferred, and
        ``view(dtype: Union[str, tvm.DataType])``.

        Returns
        -------
        view : DeclBufferFrame
            The corresponding view buffer.
        """

        def _infer_shape(shape):
            shape = list(shape)
            if -1 in shape and shape.count(-1) == 1:
                size = functools.reduce(lambda x, y: x * y, self.shape)
                n_size = functools.reduce(lambda x, y: x * y, [s for s in shape if s != -1], 1)
                shape[shape.index(-1)] = size // n_size
            else:
                # Only validate the shape product when both old and new shapes
                # are fully concrete: a Expr `==` returns an `EQ` node, not
                # a Python bool, and `assert <Expr>` raises (no __bool__).
                if all(isinstance(s, int) for s in shape) and all(
                    isinstance(s, int) for s in self.shape
                ):
                    assert functools.reduce(lambda x, y: x * y, shape) == functools.reduce(
                        lambda x, y: x * y, self.shape
                    ), (
                        "The shape of the buffer "
                        + str(self.shape)
                        + " and the new shape "
                        + str(shape)
                        + " are not compatible"
                    )
            return shape

        if len(args) == 1 and isinstance(args[0], str | tvm.DataType) and not kwargs:
            cast_dtype = tvm.DataType(args[0])
            cur_dtype = tvm.DataType(self.dtype)
            if cast_dtype.bits > cur_dtype.bits:
                # cast up
                assert cast_dtype.bits % cur_dtype.bits == 0
                ratio = cast_dtype.bits // cur_dtype.bits
                layout = self.layout.pack(ratio)
                shape = [s for s in self.shape[:-1]] + [self.shape[-1] // ratio]
                new_elem_offset = self.elem_offset // ratio
            else:
                # cast down
                assert cur_dtype.bits % cast_dtype.bits == 0
                ratio = cur_dtype.bits // cast_dtype.bits
                layout = self.layout.unpack(ratio)
                shape = [s for s in self.shape[:-1]] + [self.shape[-1] * ratio]
                new_elem_offset = self.elem_offset * ratio
            return tvm.tirx.script.builder.decl_buffer(
                shape,
                cast_dtype,
                self.data,
                self.strides,
                new_elem_offset,
                None,
                self.scope(),
                self.data_alignment,
                self.offset_factor,
                "",
                self.axis_separators,
                layout,
            )
        else:
            # --- Signature 1: view(*shape, **opts) ---
            # Check if all positional args are integers/PrimExprs with dtype int32 or int64 (the shape)  # noqa: E501
            shape = args
            assert all(
                isinstance(arg, int)
                or (tvm.ir.is_prim_expr(arg) and arg.ty.dtype in ["int32", "int64"])
                for arg in shape
            ), "shape must be a list of integers or PrimExprs with dtype int32 or int64"
            # Safely get optional keyword arguments
            layout = kwargs.get("layout", None)
            # Assert there are no other kwargs
            assert set(kwargs.keys()).issubset({"layout"}), (
                f"Unsupported kwargs for view: {set(kwargs.keys()) - {'layout'}}"
            )

            if layout is None:
                shape = _infer_shape(shape)

            return tvm.tirx.script.builder.decl_buffer(
                shape,
                self.dtype,
                self.data,
                self.strides,
                self.elem_offset,
                None,
                self.scope(),
                self.data_alignment,
                self.offset_factor,
                "",
                self.axis_separators,
                self.layout if layout is None else layout,
            )

    def local(self, *shape, layout=None) -> "Buffer":
        """Create a thread-local view of this buffer.

        When called with no shape arguments, auto-infers a 1D shape from
        the layout's non-thread component (i.e. ``layout.storage().shard``).

        Parameters
        ----------
        shape : tuple of Expr
            The shape of the local view for indexing. If omitted, a 1D
            shape is computed automatically.

        layout : optional
            Override layout. If None, uses the storage layout
            (parent layout with thread axes removed).

        Returns
        -------
        local : DeclBufferFrame
            The corresponding local buffer.
        """
        if not shape:
            local_layout = self.layout.storage()
            total = functools.reduce(
                lambda x, y: x * y, [it.extent for it in local_layout.shard], 1
            )
            shape = (total,)
        return tvm.tirx.script.builder.decl_buffer(
            shape,
            self.dtype,
            self.data,
            self.strides,
            self.elem_offset,
            None,
            self.scope(),
            self.data_alignment,
            self.offset_factor,
            "",
            self.axis_separators,
            self.layout.storage() if layout is None else layout,
        )

    def permute(self, *dims) -> "Buffer":
        """Permute the dimensions of the buffer.

        Parameters
        ----------
        dims : tuple of int
            The permutation of dimensions.

        Returns
        -------
        permuted : DeclBufferFrame
            The buffer with permuted dimensions.
        """
        new_shape = [self.shape[d] for d in dims]
        # Permute *logical* dims, not the layout's fine-grained shard iters: a
        # tcgen05/atom layout maps several shard iters to each logical axis, so
        # group by the current shape first and permute whole groups. ``group``
        # returns a regrouped layout (degenerate extent-1 iters folded away)
        # plus seps over *that* layout — permute the regrouped one, not
        # ``self.layout``. For a simple layout (one shard iter per axis) this
        # reduces to ``permute_dims(dims)``.
        grouped, seps = self.layout.group(list(self.shape))
        new_layout = grouped.permute_by_groups(seps, list(dims))
        return tvm.tirx.script.builder.decl_buffer(
            new_shape,
            self.dtype,
            self.data,
            self.strides,
            self.elem_offset,
            None,
            self.scope(),
            self.data_alignment,
            self.offset_factor,
            "",
            self.axis_separators,
            new_layout,
        )

    def __getitem__(self, indices):
        from ..arith import Analyzer  # pylint: disable=import-outside-toplevel
        from .expr import BufferLoad, Ramp  # pylint: disable=import-outside-toplevel
        from .stmt import BufferRegion  # pylint: disable=import-outside-toplevel

        if not isinstance(indices, tuple | list):
            indices = [indices]
        has_slice = any(isinstance(i, slice) for i in indices)
        has_step = any(
            isinstance(i, slice) and (i.step is not None and i.step != 1) for i in indices
        )
        has_implicit_slice = len(indices) < len(self.shape)
        analyzer = Analyzer()
        if (has_slice and not has_step) or has_implicit_slice:
            region = []
            for i, index in enumerate(indices):
                if isinstance(index, slice):
                    start = 0 if index.start is None else index.start
                    stop = self.shape[i] if index.stop is None else index.stop
                    region.append(Range.from_min_extent(start, analyzer.simplify(stop - start)))
                else:
                    region.append(
                        Range.from_min_extent(
                            index,
                            tvm.tirx.expr.IntImm(index.ty, 1) if tvm.ir.is_prim_expr(index) else 1,
                        )
                    )
            if has_implicit_slice:
                for i in range(len(indices), len(self.shape)):
                    region.append(Range.from_min_extent(0, self.shape[i]))
            return BufferRegion(self, region)
        else:
            expr_indices = []
            for i, index in enumerate(indices):
                if isinstance(index, slice):
                    start = 0 if index.start is None else index.start
                    stop = self.shape[i] if index.stop is None else index.stop
                    step = 1 if index.step is None else index.step
                    # We should ensure the dtype of start is the same with that of step.
                    if tvm.ir.is_prim_expr(start) and isinstance(step, int):
                        step = tvm.tirx.expr.IntImm(start.ty, step)
                    lanes = analyzer.simplify((stop - start + step - 1) // step)
                    if lanes == 1:
                        expr_indices.append(start)
                    else:
                        expr_indices.append(Ramp(start, step, int(lanes)))
                else:
                    expr_indices.append(index)
            return BufferLoad(self, expr_indices)


def decl_buffer(
    shape,
    dtype=None,
    name="buffer",
    data=None,
    strides=None,
    elem_offset=None,
    scope="",
    data_alignment=-1,
    offset_factor=0,
    buffer_type="",
    axis_separators=None,
    span=None,
    layout="default",
):
    # pylint: disable=import-outside-toplevel
    from .expr import Var
    from .layout import S, TileLayout

    shape = (shape,) if tvm.ir.is_prim_expr(shape) or isinstance(shape, Integral) else shape
    dtype = "float32" if dtype is None else dtype
    strides = () if strides is None else strides

    if axis_separators is None:
        axis_separators = []

    if layout == "default":
        layout = TileLayout(S[tuple(shape)]) if shape else None

    if offset_factor != 0 and elem_offset is None:
        shape_ty = shape[0].ty if shape and tvm.ir.is_prim_expr(shape[0]) else "int32"
        elem_offset = Var(f"{name}_elem_offset", shape_ty)
    if data is None:
        # Bool is represented as uint1 in the IR, but stored as int8
        storage_type = dtype if isinstance(dtype, PrimType) else PrimType(dtype)
        storage_type = PrimType("int8") if storage_type.dtype == "bool" else storage_type
        data = Var(name, PointerType(storage_type, scope), span)
    return _ffi_api.Buffer(  # type: ignore
        data,
        dtype,
        shape,
        strides,
        elem_offset,
        name,
        data_alignment,
        offset_factor,
        buffer_type,
        axis_separators,
        span,
        layout,
    )


@tvm_ffi.register_object("tirx.DataProducer")
class DataProducer(Object):
    pass