# 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.
"""Statement AST Node in TVM.

Each statement node have subfields that can be visited from python side.

.. code-block:: python

    x = tvm.tirx.Var("n", "int32")
    buffer = tvm.tirx.decl_buffer((16,), "float32")
    st = tvm.tirx.stmt.BufferStore(buffer, 1, (x,))
    assert isinstance(st, tvm.tirx.stmt.BufferStore)
    assert(st.buffer == buffer)
"""

from collections.abc import Mapping
from enum import IntEnum
from typing import TYPE_CHECKING, Any, ClassVar

import tvm_ffi

from tvm.ir import Expr, Op, Range, Span, is_prim_expr
from tvm.runtime import Object, Scriptable, const
from tvm.tirx import FloatImm, IntImm

from . import _ffi_api
from .buffer import Buffer
from .exec_scope import ExecScope, ScopeIdDef
from .expr import IterVar, StringImm, Var

if TYPE_CHECKING:
    from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext


@tvm_ffi.register_object("tirx.Stmt")
class Stmt(Object, Scriptable):
    """Base class of all the statements."""


def _normalize_legacy_stmt(stmt: Stmt | None) -> Stmt | None:
    """Expand legacy body-carrying leaf stmt wrappers into SeqStmt form.

    Legacy python compatibility may attach a `body` attribute to leaf statements
    (Bind/DeclBuffer/AllocBuffer). This helper converts such wrappers to the new
    leaf + SeqStmt representation when embedding inside another statement node.
    """

    if stmt is None:
        return None

    prefix: list[Stmt] = []
    cur = stmt
    while True:
        if isinstance(cur, DeclBuffer) and hasattr(cur, "body"):
            prefix.append(DeclBuffer(cur.buffer, cur.span))
            cur = cur.body
            continue
        if isinstance(cur, AllocBuffer) and hasattr(cur, "body"):
            prefix.append(AllocBuffer(cur.buffer, cur.annotations, cur.span))
            cur = cur.body
            continue
        break

    if not prefix:
        return stmt

    normalized_tail = _normalize_legacy_stmt(cur)
    if normalized_tail is not None:
        prefix.append(normalized_tail)
    if len(prefix) == 1:
        return prefix[0]
    return SeqStmt(prefix)


@tvm_ffi.register_object("tirx.Bind")
class Bind(Stmt):
    """Bind node.

    Bind a variable to a value in the enclosing scope.
    Bind has no body field.
    The bound variable is visible in all subsequent statements
    within the same enclosing scope (SeqStmt, ForNode.body, etc.).

    Parameters
    ----------
    var : Var
        The variable in the binding.

    value : Expr
        The value to be bound.

    span : Optional[Span]
        The location of the stmt in the source code.
    """

    var: Var
    value: Expr
    span: Span | None

    def __init__(self, var: Var, value: Expr, span: Span | None = None) -> None:
        self.__init_handle_by_constructor__(
            _ffi_api.Bind,
            var,
            value,
            span,  # type: ignore
        )


@tvm_ffi.register_object("tirx.AssertStmt")
class AssertStmt(Stmt):
    """AssertStmt node.

    Parameters
    ----------
    kind : StringImm
        The error kind, e.g. "RuntimeError", "TypeError", "ValueError".

    condition : Expr
        The assert condition.

    message_parts : list[StringImm]
        Error message fragments, concatenated at runtime when assertion fails.

    span : Span | None
        The location of the stmt in the source code.
    """

    kind: StringImm
    condition: Expr
    message_parts: list
    span: Span | None

    def __init__(
        self,
        kind: StringImm,
        condition: Expr,
        message_parts: list | None = None,
        span: Span | None = None,
    ) -> None:
        if message_parts is None:
            message_parts = []
        self.__init_handle_by_constructor__(
            _ffi_api.AssertStmt,
            kind,
            condition,
            message_parts,
            span,  # type: ignore
        )


class ForKind(IntEnum):
    """The kind of the for loop.

    note
    ----
    ForKind can change the control flow semantics
    of the loop and need to be considered in all TIR passes.
    """

    SERIAL = 0
    PARALLEL = 1
    VECTORIZED = 2
    UNROLLED = 3
    THREAD_BINDING = 4  # pylint: disable=invalid-name


@tvm_ffi.register_object("tirx.For")
class For(Stmt):
    """For node.

    Parameters
    ----------
    loop_var : Var
        The loop variable.

    min : Expr
        The beginning value.

    extent : Expr
        The length of the loop.

    kind : ForKind
        The type of the for.

    body : Stmt
        The body statement.

    thread_binding: Optional[tirx.IterVar]
        The thread this loop binds to. Only valid
        if kind is ThreadBinding

    step : Expr
        The loop step. Default to none which
        represent one.

    annotations: Optional[Mapping[str, Object]]
        Additional annotation hints.

    span : Optional[Span]
        The location of the stmt in the source code.
    """

    loop_var: Var
    min: Expr
    extent: Expr
    kind: ForKind
    body: Stmt
    thread_binding: IterVar | None
    annotations: Mapping[str, Object]
    step: Expr | None
    span: Span | None

    def __init__(
        self,
        loop_var: Var,
        min: Expr,  # pylint: disable=redefined-builtin
        extent: Expr,
        kind: ForKind,
        body: Stmt,
        thread_binding: IterVar | None = None,
        annotations: Mapping[str, Object] | None = None,
        step: Expr | None = None,
        span: Span | None = None,
    ) -> None:
        body = _normalize_legacy_stmt(body)
        self.__init_handle_by_constructor__(
            _ffi_api.For,  # type: ignore
            loop_var,
            min,
            extent,
            kind,
            body,
            thread_binding,
            annotations,
            step,
            span,
        )


@tvm_ffi.register_object("tirx.While")
class While(Stmt):
    """While node.

    Parameters
    ----------
    condition : Expr
        The termination condition.

    body : Stmt
        The body statement.

    span : Optional[Span]
        The location of the stmt in the source code.
    """

    condition: Expr
    body: Stmt
    span: Span | None

    def __init__(self, condition: Expr, body: Stmt, span: Span | None = None) -> None:
        body = _normalize_legacy_stmt(body)
        self.__init_handle_by_constructor__(_ffi_api.While, condition, body, span)  # type: ignore


@tvm_ffi.register_object("tirx.BufferStore")
class BufferStore(Stmt):
    """Buffer store node.

    Parameters
    ----------
    buffer : Buffer
        The buffer.

    value : Expr
        The value we to be stored.

    indices : List[Expr]
        The indices location 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.

    span : Optional[Span]
        The location of the stmt in the source code.
    """

    buffer: Buffer
    value: Expr
    indices: list[Expr]
    predicate: Expr | None
    span: Span | None

    def __init__(
        self,
        buffer: Buffer,
        value: Expr,
        indices: list[Expr],
        predicate: Expr | None = None,
        span: Span | None = None,
    ) -> None:
        self.__init_handle_by_constructor__(
            _ffi_api.BufferStore,
            buffer,
            value,
            indices,
            predicate,
            span,  # type: ignore
        )


@tvm_ffi.register_object("tirx.AllocBuffer")
class AllocBuffer(Stmt):
    """AllocBuffer node.

    Allocates a buffer and declares it in scope.

    Parameters
    ----------
    buffer: Buffer
        The buffer being allocated and declared.

    annotations: Optional[dict]
        Additional annotations about the allocation.

    span: Optional[Span]
        The location of this AllocBuffer in the source code.
    """

    buffer: Buffer
    span: Span | None

    def __init__(self, buffer: Buffer, *args, **kwargs) -> None:
        body: Stmt | None = None
        annotations: dict | None = None
        span: Span | None = None

        idx = 0
        argc = len(args)

        # Legacy form: AllocBuffer(buffer, body[, annotations][, span])
        if idx < argc and isinstance(args[idx], Stmt):
            body = args[idx]
            idx += 1

        if idx < argc:
            arg = args[idx]
            if isinstance(arg, Mapping):
                annotations = dict(arg)
                idx += 1
            elif arg is None:
                annotations = None
                idx += 1
            elif isinstance(arg, Span):
                span = arg
                idx += 1
            else:
                raise TypeError(
                    "AllocBuffer expects (buffer[, annotations][, span]) or "
                    "legacy (buffer, body[, annotations][, span])"
                )

        if idx < argc:
            arg = args[idx]
            if arg is None or isinstance(arg, Span):
                span = arg
                idx += 1
            else:
                raise TypeError("AllocBuffer span must be a Span or None")

        if idx != argc:
            raise TypeError(
                "AllocBuffer expects (buffer[, annotations][, span]) or "
                "legacy (buffer, body[, annotations][, span])"
            )

        if kwargs:
            invalid_keys = set(kwargs.keys()) - {"body", "annotations", "span"}
            if invalid_keys:
                raise TypeError(f"Unexpected keyword arguments for AllocBuffer: {invalid_keys}")
            if "body" in kwargs:
                kw_body = kwargs["body"]
                if kw_body is not None and not isinstance(kw_body, Stmt):
                    raise TypeError("AllocBuffer body must be a Stmt or None")
                if body is not None and kw_body is not None and body is not kw_body:
                    raise TypeError("AllocBuffer body specified by both args and kwargs")
                body = kw_body if kw_body is not None else body
            if "annotations" in kwargs:
                kw_ann = kwargs["annotations"]
                if kw_ann is not None and not isinstance(kw_ann, Mapping):
                    raise TypeError("AllocBuffer annotations must be Mapping or None")
                if annotations is not None and kw_ann is not None and annotations != dict(kw_ann):
                    raise TypeError("AllocBuffer annotations specified by both args and kwargs")
                annotations = dict(kw_ann) if kw_ann is not None else annotations
            if "span" in kwargs:
                kw_span = kwargs["span"]
                if kw_span is not None and not isinstance(kw_span, Span):
                    raise TypeError("AllocBuffer span must be a Span or None")
                if span is not None and kw_span is not None and span is not kw_span:
                    raise TypeError("AllocBuffer span specified by both args and kwargs")
                span = kw_span if kw_span is not None else span

        self.__init_handle_by_constructor__(_ffi_api.AllocBuffer, buffer, annotations, span)
        # Legacy compatibility. Body is carried on python side only.
        if body is not None:
            self.body = body


@tvm_ffi.register_object("tirx.DeclBuffer")
class DeclBuffer(Stmt):
    """DeclBuffer node.

    Parameters
    ----------
    buffer: Buffer
        The buffer being declared.

    span: Optional[Span]
        The location of this DeclBuffer in the source code.
    """

    buffer: Buffer
    span: Span | None

    def __init__(self, buffer: Buffer, *args, **kwargs) -> None:
        body: Stmt | None = None
        span: Span | None = None

        if len(args) == 1:
            arg0 = args[0]
            if isinstance(arg0, Stmt):
                body = arg0
            elif arg0 is None or isinstance(arg0, Span):
                span = arg0
            else:
                raise TypeError(
                    "DeclBuffer expects (buffer[, span]) or legacy (buffer, body[, span])"
                )
        elif len(args) == 2:
            body, span = args
            if body is not None and not isinstance(body, Stmt):
                raise TypeError("Legacy DeclBuffer body must be a Stmt or None")
            if span is not None and not isinstance(span, Span):
                raise TypeError("DeclBuffer span must be a Span or None")
        elif len(args) > 2:
            raise TypeError("DeclBuffer expects (buffer[, span]) or legacy (buffer, body[, span])")

        if kwargs:
            invalid_keys = set(kwargs.keys()) - {"body", "span"}
            if invalid_keys:
                raise TypeError(f"Unexpected keyword arguments for DeclBuffer: {invalid_keys}")
            if "body" in kwargs:
                kw_body = kwargs["body"]
                if kw_body is not None and not isinstance(kw_body, Stmt):
                    raise TypeError("DeclBuffer body must be a Stmt or None")
                if body is not None and kw_body is not None and body is not kw_body:
                    raise TypeError("DeclBuffer body specified by both args and kwargs")
                body = kw_body if kw_body is not None else body
            if "span" in kwargs:
                kw_span = kwargs["span"]
                if kw_span is not None and not isinstance(kw_span, Span):
                    raise TypeError("DeclBuffer span must be a Span or None")
                if span is not None and kw_span is not None and span is not kw_span:
                    raise TypeError("DeclBuffer span specified by both args and kwargs")
                span = kw_span if kw_span is not None else span

        self.__init_handle_by_constructor__(_ffi_api.DeclBuffer, buffer, span)
        # Legacy compatibility. Body is carried on python side only.
        if body is not None:
            self.body = body


@tvm_ffi.register_object("tirx.AttrStmt")
class AttrStmt(Stmt):
    """AttrStmt node.

    Parameters
    ----------
    node : Object
        The node to annotate the attribute

    attr_key : str
        Attribute type key.

    value : Expr
        The value of the attribute

    body : Stmt
        The body statement.

    span : Optional[Span]
        The location of the stmt in the source code.
    """

    node: Object
    attr_key: str
    value: Expr
    body: Stmt
    span: Span | None

    def __init__(
        self, node: Object, attr_key: str, value: Expr, body: Stmt, span: Span | None = None
    ) -> None:
        body = _normalize_legacy_stmt(body)
        self.__init_handle_by_constructor__(
            _ffi_api.AttrStmt,
            node,
            attr_key,
            value,
            body,
            span,  # type: ignore
        )


@tvm_ffi.register_object("tirx.SeqStmt")
class SeqStmt(Stmt):
    """Sequence of statements.

    Parameters
    ----------
    seq : List[Stmt]
        The statements

    span : Optional[Span]
        The location of the stmt in the source code.
    """

    seq: list[Stmt]
    span: Span | None

    def __init__(self, seq: list[Stmt], span: Span | None = None) -> None:
        seq = [_normalize_legacy_stmt(s) for s in seq]
        self.__init_handle_by_constructor__(_ffi_api.SeqStmt, seq, span)  # type: ignore

    def __getitem__(self, i: int):
        return self.seq[i]

    def __len__(self):
        return len(self.seq)


@tvm_ffi.register_object("tirx.IfThenElse")
class IfThenElse(Stmt):
    """IfThenElse node.

    Parameters
    ----------
    condition : Expr
        The expression

    then_case : Stmt
        The statement to execute if condition is true.

    else_case : Optional[Stmt]
        The statement to execute if condition is false.

    span : Optional[Span]
        The location of the stmt in the source code.
    """

    condition: Expr
    then_case: Stmt
    else_case: Stmt | None

    def __init__(
        self, condition: Expr, then_case: Stmt, else_case: Stmt | None, span: Span | None = None
    ) -> None:
        then_case = _normalize_legacy_stmt(then_case)
        else_case = _normalize_legacy_stmt(else_case)
        self.__init_handle_by_constructor__(
            _ffi_api.IfThenElse,
            condition,
            then_case,
            else_case,
            span,  # type: ignore
        )


@tvm_ffi.register_object("tirx.Evaluate")
class Evaluate(Stmt):
    """Evaluate node.

    Parameters
    ----------
    value : Expr
        The expression to be evaluated.

    span : Optional[Span]
        The location of the stmt in the source code.
    """

    value: Expr
    span: Span | None

    def __init__(self, value: Expr, span: Span | None = None) -> None:
        self.__init_handle_by_constructor__(_ffi_api.Evaluate, value, span)  # type: ignore


@tvm_ffi.register_object("tirx.BufferRegion")
class BufferRegion(Object, Scriptable):
    """BufferRegion node.

    Parameters
    ----------
    buffer : Buffer
        The buffer of the buffer region

    region : List[Range]
        The region array of the buffer region
    """

    buffer: Buffer
    region: list[Range]

    def __init__(self, buffer: Buffer, region: list[Range]) -> None:
        self.__init_handle_by_constructor__(_ffi_api.BufferRegion, buffer, region)  # type: ignore

    def __getitem__(self, indices):
        from ..arith import Analyzer

        if not isinstance(indices, tuple | list):
            indices = [indices]

        has_step = any(
            isinstance(i, slice) and (i.step is not None and i.step != 1) for i in indices
        )
        if has_step:
            raise ValueError("BufferRegion slicing does not support steps")

        analyzer = Analyzer()
        new_region = []
        for i, index in enumerate(indices):
            old_range = self.region[i]
            if isinstance(index, slice):
                start = 0 if index.start is None else index.start
                stop = old_range.extent if index.stop is None else index.stop
                new_min = old_range.min + start
                new_extent = analyzer.simplify(stop - start)
                new_region.append(Range.from_min_extent(new_min, new_extent))
            else:
                new_min = old_range.min + index
                new_region.append(
                    Range.from_min_extent(
                        new_min, IntImm(index.ty, 1) if is_prim_expr(index) else 1
                    )
                )
        # Fill remaining dimensions with their original ranges
        for i in range(len(indices), len(self.region)):
            new_region.append(self.region[i])
        return BufferRegion(self.buffer, new_region)


@tvm_ffi.register_object("tirx.MatchBufferRegion")
class MatchBufferRegion(Object, Scriptable):
    """MatchBufferRegion node.

    Parameters
    ----------
    buffer : Buffer
        The target buffer

    source : BufferRegion
        The region of source buffer
    """

    buffer: Buffer
    source: BufferRegion

    def __init__(self, buffer: Buffer, source: BufferRegion) -> None:
        self.__init_handle_by_constructor__(
            _ffi_api.MatchBufferRegion,
            buffer,
            source,  # type: ignore
        )


@tvm_ffi.register_object("tirx.SBlock")
class SBlock(Stmt):
    """SBlock node.

    Parameters
    ----------
    iter_vars : List[IterVar]
        The block Variable.

    reads : List[BufferRegion]
        The read buffer regions of the block.

    writes: List[BufferRegion]
        The write buffer regions of the block.

    name_hint: str
        the name_hint of the block.

    body: Stmt
        The body of the block.

    init: Optional[Stmt]
        The init block of the reduction block

    alloc_buffers: Optional[list[Buffer]]
        The buffer allocations

    match_buffers: Optional[List[MatchBufferRegion]]
        The subregion buffer match

    annotations: Optional[Mapping[str, Object]]
        Additional annotation hints.

    span : Optional[Span]
        The location of this block in the source code.
    """

    iter_vars: list[IterVar]
    reads: list[BufferRegion]
    writes: list[BufferRegion]
    name_hint: str
    body: Stmt
    init: Stmt | None
    alloc_buffers: list[Buffer]
    match_buffers: list[MatchBufferRegion]
    annotations: Mapping[str, Object]
    span: Span | None

    def __init__(
        self,
        iter_vars: list[IterVar],
        reads: list[BufferRegion],
        writes: list[BufferRegion],
        name_hint: str,
        body: Stmt,
        init: Stmt | None = None,
        alloc_buffers: list[Buffer] | None = None,
        match_buffers: list[MatchBufferRegion] | None = None,
        annotations: Mapping[str, Object] | None = None,
        span: Span | None = None,
    ) -> None:
        if alloc_buffers is None:
            alloc_buffers = []
        if match_buffers is None:
            match_buffers = []
        if annotations is None:
            annotations = {}
        body = _normalize_legacy_stmt(body)
        init = _normalize_legacy_stmt(init)
        self.__init_handle_by_constructor__(
            _ffi_api.SBlock,  # type: ignore
            iter_vars,
            reads,
            writes,
            name_hint,
            body,
            init,
            alloc_buffers,
            match_buffers,
            annotations,
            span,
        )  # type: ignore


@tvm_ffi.register_object("tirx.SBlockRealize")
class SBlockRealize(Stmt):
    """SBlockRealize node.

    Parameters
    ----------
    iter_values : List[Expr]
        The binding values of the block var.

    predicate : Union[Expr, bool]
        The predicate of the block.

    block : SBlock
        The block to realize

    span : Optional[Span]
        The location of this block_realize in the source code.
    """

    iter_values: list[Expr]
    predicate: Expr
    block: SBlock
    span: Span | None

    def __init__(
        self,
        iter_values: list[Expr],
        predicate: Expr | bool,
        block: SBlock,
        span: Span | None = None,
    ) -> None:
        if isinstance(predicate, bool):
            predicate = const(predicate, "bool")
        self.__init_handle_by_constructor__(
            _ffi_api.SBlockRealize,  # type: ignore
            iter_values,
            predicate,
            block,
            span,
        )  # type: ignore


@tvm_ffi.register_object("tirx.ScopeIdDefStmt")
class ScopeIdDefStmt(Stmt):
    """ScopeIdDefStmt node.

    Leaf statement that introduces scope-identifier vars
    (``wg_id = Tx.warpgroup_id([N])``, ``warp_id = Tx.warp_id_in_wg([4])``,
    ``lane_id = Tx.lane_id([32])``, …) at the kernel-body top level. The
    underlying ``ScopeIdDef`` carries the def vars, their extents, and
    the parent/child scope binding.

    Note: the C++ field is named ``def`` (a Python keyword). Access it
    via ``getattr(stmt, "def")`` or ``stmt.__getattribute__("def")`` —
    the type-annotation alias here is purely for documentation.

    Parameters
    ----------
    def_ : ScopeIdDef
        The scope-id definition (def vars, extents, scope binding).

    span : Optional[Span]
        The location of this statement in the source code.
    """

    span: Span | None

    def __init__(self, def_: ScopeIdDef, span: Span | None = None) -> None:
        self.__init_handle_by_constructor__(
            _ffi_api.ScopeIdDefStmt,  # type: ignore
            def_,
            span,
        )  # type: ignore


@tvm_ffi.register_object("tirx.Break")
class Break(Stmt):
    """Break node.

    Parameters
    ----------
    """

    def __init__(self, span: Span | None = None) -> None:
        self.__init_handle_by_constructor__(_ffi_api.Break, span)  # type: ignore


@tvm_ffi.register_object("tirx.Continue")
class Continue(Stmt):
    """Continue node.

    Parameters
    ----------
    """

    def __init__(self, span: Span | None = None) -> None:
        self.__init_handle_by_constructor__(_ffi_api.Continue, span)  # type: ignore


def stmt_seq(*args: Expr | Stmt) -> SeqStmt:
    """Make sequence of statements

    Parameters
    ----------
    *args : Union[Expr, Stmt]
        List of statements to be combined as sequence.

    Returns
    -------
    stmt : Stmt
        The combined statement.
    """
    ret = []
    for value in args:
        if not isinstance(value, Stmt):
            value = Evaluate(value)
        ret.append(value)
    if len(ret) == 1:
        return ret[0]
    return SeqStmt(ret)


def stmt_list(stmt: Stmt) -> list[Stmt]:
    """Make list of stmt from blocks.

    Parameters
    ----------
    stmt : Stmt
        The input statement.

    Returns
    -------
    stmt_list : List[Stmt]
        The unpacked list of statements
    """
    if isinstance(stmt, SeqStmt):
        res = []
        for x in stmt:
            res += stmt_list(x)
        return res
    return [stmt]


def normalize_const_arg(arg) -> Expr:
    if isinstance(arg, float):
        return FloatImm("float32", arg)
    return arg


@tvm_ffi.register_object("tirx.TilePrimitiveCall")
class TilePrimitiveCall(Stmt):
    """TilePrimitiveCall node.

    Parameters
    ----------
    op : Op
        The operator.

    args : List[Expr]
        The arguments.

    workspace : Map[str, Buffer]
        The workspace.

    config : Map[str, ObjectRef]
        The scheduler/config dictionary.

    dispatch : Optional[str]
        The explicit variant name to dispatch to.

    scope : ExecScope
        The cooperation scope of this call. Defaults to ``thread`` (an unscoped call).
    """

    args: list[Expr]
    workspace: dict[str, Buffer]
    config: dict[str, Any]
    dispatch: str | None
    scope: ExecScope
    _registry: ClassVar[dict[Op, type["TilePrimitiveCall"]]] = {}

    def __init__(
        self,
        *args: list[Expr],
        op: Op | None = None,
        workspace: dict[str, Buffer] | None = None,
        config: dict[str, Any] | None = None,
        dispatch: str | None = None,
        scope: ExecScope | None = None,
    ) -> None:
        if workspace is None:
            workspace = {}
        if config is None:
            config = {}
        if scope is None:
            scope = ExecScope("thread")
        if op is None:
            assert self.__class__ != TilePrimitiveCall, (
                "Directly instantiating TilePrimitiveCall needs to specify the op"
            )
            op = self.__class__.op
        args = list(map(normalize_const_arg, args))
        self.__init_handle_by_constructor__(
            _ffi_api.TilePrimitiveCall,
            op,
            args,
            workspace,
            config,
            dispatch,
            scope,  # pylint: disable=no-member
        )

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        if hasattr(cls, "op"):
            cls._registry[cls.op] = cls

    @classmethod
    def downcast(cls, instance: "TilePrimitiveCall") -> "TilePrimitiveCall":
        subclass = cls._registry.get(instance.op)
        if subclass is None:
            return instance  # Unknown op: return as-is
        new_instance = subclass.__new__(subclass)
        new_instance.__init_handle_by_constructor__(
            _ffi_api.TilePrimitiveCallCopyHandle,
            instance,  # pylint: disable=no-member
        )
        return new_instance

    def replace(self, **changes: Any) -> "TilePrimitiveCall":
        """Return a copy of this call with selected fields replaced.

        Every field that is not overridden in ``changes`` is preserved from
        ``self`` (including ``scope``), so rebuilds never silently drop fields.
        The returned node is downcast to the registered subclass for ``op``.

        Parameters
        ----------
        **changes : Any
            Field overrides; any of ``op``, ``args``, ``workspace``, ``config``,
            ``dispatch``, ``scope``.

        Returns
        -------
        new_call : TilePrimitiveCall
            A new call with the requested fields replaced.
        """
        unknown = set(changes) - {"op", "args", "workspace", "config", "dispatch", "scope"}
        if unknown:
            raise TypeError(f"Unknown field(s) for TilePrimitiveCall.replace: {sorted(unknown)}")
        new_call = TilePrimitiveCall(
            *changes.get("args", self.args),
            op=changes.get("op", self.op),
            workspace=changes.get("workspace", self.workspace),
            config=changes.get("config", self.config),
            dispatch=changes.get("dispatch", self.dispatch),
            scope=changes.get("scope", self.scope),
        )
        return TilePrimitiveCall.downcast(new_call)

    def with_workspace(self, workspace: dict[str, Buffer]) -> "TilePrimitiveCall":
        """Return a copy with ``workspace`` replaced, preserving all other fields."""
        return self.replace(workspace=workspace)

    @property
    def srcs(self) -> list[Expr]:
        raise NotImplementedError("Subclass must implement this method")

    @property
    def dsts(self) -> list[Expr]:
        raise NotImplementedError("Subclass must implement this method")

    def get_private_buffers(
        self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext"
    ) -> dict[str, Any]:
        """
        Create private (intermediate) buffers needed in this operator.

        Parameters
        ----------
        buffer_dict: Dict[Any, Tuple[Buffer, Optional[Stmt]]]
            A dictionary containing private buffers (and their init stmts) in other operators.
            Key can be anything to reference the buffer.
            This is used to reuse private buffers in other operators (like identity tensor etc.).
            If the buffer is not found in the buffer_dict, it will be created and added to
            the buffer_dict.
            If the buffer is found in the buffer_dict but smaller than required, it will be
            enlarged and updated.

        sctx: DispatchContext
            The dispatch context.
            This is used to get the target and reuse op dispatch implementations.

        Returns:
            private_buffer_refs: Dict[str, Any]
            The references to private buffers created in this operator.
            Key will be the name to add into workspace.
            private buffer can be accessed by buffer_dict[private_buffer_refs[name]]
        """
        if sctx.target.kind.name == "trn":
            return self.get_private_buffers_trn(buffer_dict, sctx)
        elif sctx.target.kind.name == "cuda":
            return self.get_private_buffers_cuda(buffer_dict, sctx)
        else:
            raise ValueError(f"Unsupported target: {sctx.target.kind.name}")

    def get_private_buffers_trn(
        self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext"
    ) -> dict[str, Any]:
        return {}

    def get_private_buffers_cuda(
        self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext"
    ) -> dict[str, Any]:
        return {}

    def validate(self) -> None:
        pass