# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

import functools
import inspect
import threading
import re
from collections import defaultdict, Counter
from contextlib import contextmanager
from dataclasses import dataclass
from enum import EnumMeta
from typing import Optional, NamedTuple, Tuple, Sequence, Any, Union, Callable

from cuda.tile._datatype import (
    is_integral, is_float,
    is_boolean, is_signed, DType, PointerInfo, is_pointer_dtype)
from cuda.tile._bytecode.version import BytecodeVersion
from cuda.tile._exception import TileTypeError, TileUnsupportedFeatureError
from cuda.tile._ir.ops_utils import get_dtype

from .typing_support import datatype, get_signature
from .ir import Var, Builder
from .type import TiledViewTy, TupleTy, TileTy, DTypeSpec, EnumTy, StringTy, ArrayTy, SliceType, \
    ListTy, LooselyTypedScalar, RangeIterType, FunctionTy, ClosureTy, BoundMethodTy, \
    DTypeConstructor, Type, RawArrayMemoryTy, DataclassTy, TupleValue, PointerInfoTy, TensorLikeTy


def _verify_params_match(stub_sig: inspect.Signature, func_sig: inspect.Signature):
    assert len(stub_sig.parameters) == len(func_sig.parameters), (
        f"Stub and implementation must have same number of parameters."
        f" Signatures: {stub_sig}, {func_sig}.")
    for i, (stub_param, func_param) in enumerate(zip(stub_sig.parameters.values(),
                                                     func_sig.parameters.values(), strict=True)):
        assert stub_param.name == func_param.name, (
            f"Stub and implementation have different parameter names at position {i}"
            f" Signatures: {stub_sig}, {func_sig}.")


@dataclass(frozen=True)
class WildcardClass:
    pass


# Can be used as a parameter of an overloaded implementation to match any value.
#
# For example, a getattr() implementation for a specific type may choose to handle
# all attributes in a generic way, instead of registering an individual overload
# for each attribute. In this case, one can use WILDCARD in place of the attribute name.
WILDCARD = WildcardClass()


class OverloadNotFoundError(Exception):
    def __init__(self, found_overload_matching_first_param: bool):
        self.found_overload_matching_first_param = found_overload_matching_first_param


class ImplRegistry:
    def __init__(self):
        self.op_implementations = dict()
        self._overloaded_implementations = defaultdict(dict)
        self.unflatten_aggregate_implementations = dict()

    @staticmethod
    def get_current() -> "ImplRegistry":
        ret = _current_registry.impl_registry
        assert ret is not None
        return ret

    @contextmanager
    def as_current(self):
        old = _current_registry.impl_registry
        _current_registry.impl_registry = self
        try:
            yield self
        finally:
            _current_registry.impl_registry = old

    def update(self, source: "ImplRegistry"):
        self.op_implementations.update(source.op_implementations)
        for stub, overloads in source._overloaded_implementations.items():
            self._overloaded_implementations[stub].update(overloads)
        self.unflatten_aggregate_implementations.update(source.unflatten_aggregate_implementations)

    def overload_dispatcher(self, stub, *, fixed_args: Sequence[Any] = ()):
        """
        Decorates a function to attach an overloaded implementation dispatcher to a stub.

        The decorated function must yield an "overload key", i.e. a tuple of values based
        on which the appropriate overload should be selected, and then handle
        the OverloadNotFoundError by re-raising a TileError.
        See getattr_overload_dispatcher() for an example.
        """

        def decorate(key_func):
            orig_func = key_func
            if len(fixed_args) > 0:
                key_func = functools.partial(orig_func, *fixed_args)

            @functools.wraps(key_func)
            async def implementation(*args):
                generator = key_func(*args)
                key = next(generator)
                cur_registry = ImplRegistry.get_current()
                overload_impl = cur_registry._find_overload(stub, key)
                if overload_impl is None:
                    generator.throw(OverloadNotFoundError(
                            cur_registry._have_overload_matching_first_param(stub, key[0])))
                    raise RuntimeError("Expected the overload key provider to re-throw a TileError")

                try:
                    next(generator)
                except StopIteration:
                    pass  # Good, that's what we expect
                else:
                    raise RuntimeError("Expected the overload key provider to yield exactly once")

                result = overload_impl(*args)
                if overload_impl._is_coroutine:
                    result = await result

                return result

            self.impl(stub)(implementation)
            return orig_func

        return decorate

    def _find_overload(self, stub: Callable, overload: tuple[Any, ...]) -> Callable | None:
        candidates = self._overloaded_implementations[stub]
        best_matches = []
        best_priority = -1

        for priority, predicates, impl in candidates.values():
            if priority < best_priority:
                continue

            if not all(p(arg) for p, arg in zip(predicates, overload, strict=True)):
                continue

            if priority > best_priority:
                best_matches.clear()
                best_priority = priority

            best_matches.append(impl)

        match best_matches:
            case []:
                return None
            case [x]:
                return x
            case _:
                raise ValueError(f"Multiple matching overloads found for {stub}, {overload}")

    def impl(self, stub, *, fixed_args: Sequence[Any] = (),
             min_version: Optional[BytecodeVersion] = None,
             overload: tuple[Any, ...] = ()):
        stub_sig = get_signature(stub)

        def _check_version():
            cur_version = Builder.get_current().ir_ctx.tileiras_version
            if min_version is not None and cur_version < min_version:
                raise TileUnsupportedFeatureError(
                    f"{stub.__name__} requires tileiras "
                    f"{min_version.as_string()} or later. "
                    f"Current version is {cur_version.as_string()}."
                )

        def decorate(func):
            orig_func = func
            if len(fixed_args) > 0:
                func = functools.partial(orig_func, *fixed_args)

            func_sig = inspect.signature(func)
            _verify_params_match(stub_sig, func_sig)
            is_coroutine = inspect.iscoroutinefunction(func)
            if is_coroutine:
                @functools.wraps(func)
                async def wrapper(*args, **kwargs):
                    _check_version()
                    # Memorize the stub and the args so that we can automatically
                    # provide context for error messages.
                    old = _current_stub.stub_and_args
                    _current_stub.stub_and_args = (stub, stub_sig, func_sig, args, kwargs)
                    try:
                        return await func(*args, **kwargs)
                    finally:
                        _current_stub.stub_and_args = old
            else:
                @functools.wraps(func)
                def wrapper(*args, **kwargs):
                    _check_version()
                    # Memorize the stub and the args so that we can automatically
                    # provide context for error messages.
                    old = _current_stub.stub_and_args
                    _current_stub.stub_and_args = (stub, stub_sig, func_sig, args, kwargs)
                    try:
                        return func(*args, **kwargs)
                    finally:
                        _current_stub.stub_and_args = old

            wrapper._is_coroutine = is_coroutine

            if len(overload) == 0:
                self.op_implementations[stub] = wrapper
            else:
                predicates = tuple(_predicate_from_overload_pattern(p) for p in overload)
                self._overloaded_implementations[stub][overload] = \
                    (sum(p != WILDCARD for p in overload), predicates, wrapper)

            return orig_func

        return decorate

    def _have_overload_matching_first_param(self, stub: Callable, first_param: Any) -> bool:
        candidates = self._overloaded_implementations[stub]
        return any(predicates[0](first_param)
                   for _priority, predicates, _impl in candidates.values())

    def unflatten_aggregate_impl(self, aggregate_type_class: type[Type]):
        def decorate(func):
            self.unflatten_aggregate_implementations[aggregate_type_class] = func
            return func
        return decorate


def _predicate_from_overload_pattern(pattern):
    if pattern == WILDCARD:
        return lambda _: True
    elif isinstance(pattern, type):
        return lambda x: issubclass(x, pattern)
    else:
        return lambda x: pattern == x


class _CurrentRegistry(threading.local):
    impl_registry = None


_current_registry = _CurrentRegistry()


class _CurrentStub(threading.local):
    stub_and_args = None


_current_stub = _CurrentStub()


def is_scalar(ty: Type, dtype_predicate: Callable[[DType], bool] = lambda _: True) -> bool:
    return (isinstance(ty, TensorLikeTy)
            and ty.tensor_shape() == ()
            and dtype_predicate(ty.tensor_dtype()))


def require_constant_int(var: Var) -> int:
    if not var.is_constant():
        raise _make_type_error("Expected an integer constant, but given value is not constant", var)
    ty = var.get_type()
    if not is_scalar(ty, is_integral):
        raise _make_type_error(f"Expected an integer constant, but given value has type {ty}", var)
    return var.get_constant()


def require_optional_constant_int(var: Var) -> Optional[int]:
    if var.is_constant() and var.get_constant() is None:
        return None
    return require_constant_int(var)


def require_constant_bool(var: Var) -> bool:
    if not var.is_constant():
        raise _make_type_error("Expected a boolean constant, but given value is not constant", var)
    ty = var.get_type()
    if not is_scalar(ty, is_boolean):
        raise _make_type_error(f"Expected a boolean constant, but given value has type {ty}", var)
    return var.get_constant()


def require_constant_scalar(var: Var) -> bool | int | float:
    ty = var.get_type()
    if not is_scalar(ty):
        raise _make_type_error(f"Expected a scalar constant, but given value has type {ty}", var)
    if not var.is_constant():
        raise _make_type_error(f"Expected a constant, but given value has non-constant type {ty}",
                               var)
    ret = var.get_constant()
    assert isinstance(ret, bool | int | float)
    return ret


def require_constant_scalar_tuple(var: Var) -> tuple[bool | int | float, ...]:
    ty = require_tuple_type(var)
    ret = []
    tuple_val = var.get_aggregate()
    assert isinstance(tuple_val, TupleValue)
    for i, (item_ty, item) in enumerate(zip(ty.value_types, tuple_val.items, strict=True)):
        if not is_scalar(item_ty):
            raise _make_type_error(f"Expected a tuple of scalar constants,"
                                   f" but item at position #{i} has type {item_ty}", var)
        if not item.is_constant():
            raise _make_type_error(f"Expected a tuple of scalar constants,"
                                   f" but item at position #{i} has non-constant type {ty}", var)
        value = item.get_constant()
        assert isinstance(value, bool | int | float)
        ret.append(value)
    return tuple(ret)


def require_optional_constant_bool(var: Var) -> Optional[bool]:
    if var.is_constant() and var.get_constant() is None:
        return None
    return require_constant_bool(var)


def require_constant_str(var: Var) -> str:
    if not var.is_constant():
        raise _make_type_error("Expected a string constant, but given value is not constant", var)
    ty = var.get_type()
    if not isinstance(ty, StringTy):
        raise _make_type_error(f"Expected a string constant, but given value has type {ty}", var)
    return ty.value


def require_optional_constant_str(var: Var) -> Optional[str]:
    if var.is_constant() and var.get_constant() is None:
        return None
    return require_constant_str(var)


def require_constant_slice(var: Var) -> slice:
    if not var.is_constant():
        raise _make_type_error("Expected a slice constant, but given value is not constant", var)
    ty = var.get_type()
    if not isinstance(ty, SliceType):
        raise _make_type_error(f"Expected a slice constant, but given value has type {ty}", var)
    return var.get_constant()


def require_dtype_spec(var: Var) -> DType:
    ty = var.get_type()
    if not isinstance(ty, DTypeSpec):
        raise _make_type_error(f"Expected a dtype constant, but given value has type {ty}", var)
    return ty.dtype


def require_constant_pointer_info(var: Var) -> PointerInfo:
    ty = var.get_type()
    if not isinstance(ty, PointerInfoTy):
        raise _make_type_error(f"Expected a PointerInfo object, but given value has type {ty}", var)
    assert var.is_constant()
    return ty.info


def require_optional_constant_enum(var: Var, enum: EnumMeta):
    if var.is_constant() and var.get_constant() is None:
        return None
    return require_constant_enum(var, enum)


def require_optional_range_type(var: Var) -> RangeIterType | None:
    if var.is_constant() and var.get_constant() is None:
        return None
    ty = var.get_type()
    if not isinstance(ty, RangeIterType):
        raise _make_type_error(f"Expected a range object, but given value has type {ty}", var)
    return ty


def require_constant_enum(var: Var, enum: EnumMeta):
    if not var.is_constant():
        raise _make_type_error(f"Expected {enum.__name__} constant,"
                               f" but given value is not constant", var)
    ty = var.get_type()
    if not isinstance(ty, EnumTy) or ty.enum_ty is not enum:
        raise _make_type_error(f"Expected {enum.__name__}, but given value has type {ty}", var)
    return var.get_constant()


def normalize_axis(axis: int, ndim: int, var: Optional[Var] = None) -> int:
    orig_axis = axis
    if axis < 0:
        axis += ndim
    if axis < 0 or axis >= ndim:
        raise _make_type_error(f"Axis {orig_axis} is out of range for rank {ndim}'", var)
    return axis


def require_constant_int_tuple(var: Var, allow_single_int: bool = False) -> Tuple[int, ...]:
    if not var.is_constant():
        raise _make_type_error("Expected a constant integer tuple,"
                               " but given value is not constant", var)

    ty = var.get_type()
    if allow_single_int and is_scalar(ty):
        return require_constant_int(var),

    if not isinstance(ty, TupleTy):
        raise _make_type_error(f"Expected a tuple, but given value has type {ty}", var)

    for i, item_ty in enumerate(ty.value_types):
        if not is_scalar(item_ty, is_integral):
            raise _make_type_error(f"Expected a tuple of integers,"
                                   f" but element #{i} has type {item_ty}", var)

    val = var.get_constant()
    assert isinstance(val, tuple)
    assert all(isinstance(x, int) for x in val)
    return val


def require_constant_shape(var: Var,
                           allow_single_int: bool = False,
                           expected_rank: Optional[int] = None,
                           allow_0d_shape: bool = False,
                           allow_non_power_of_two: bool = False,
                           var_name: str = "shape") -> Tuple[int, ...]:
    shape = require_constant_int_tuple(var, allow_single_int=allow_single_int)

    if (expected_rank is not None and len(shape) != expected_rank
            and not (allow_0d_shape and len(shape) == 0)):
        raise _make_type_error(
            f"Expected {var_name} length to be {expected_rank}, got {len(shape)}", var)

    for i, x in enumerate(shape):
        if x <= 0:
            raise _make_type_error(
                f"Dimension #{i} of {var_name} {tuple(shape)} is not positive", var)
        if not allow_non_power_of_two and x & (x - 1) != 0:
            raise _make_type_error(
                f"Dimension #{i} of {var_name} {tuple(shape)} is not a power of two", var)

    return shape


def require_constant_axis_order(var: Var, rank: int) -> Tuple[int, ...]:
    """
    Helper for matching the 'order' argument of functions like cuda.tile.load() etc.
    The order can either be a string literal "C" or "F" (which represents a NumPy-style
    contiguous axis order), or a literal sequence of integers, e.g. (1, -2, 0).

    Returns a tuple that contains the "normalized" (i.e. non-negative) axis indices.
    """
    if not var.is_constant():
        raise _make_type_error("Expected a constant string or integer tuple", var)

    value = var.get_constant()
    if value == "C":
        return tuple(range(rank))
    elif value == "F":
        return tuple(range(rank - 1, -1, -1))
    elif isinstance(value, str):
        raise _make_type_error(f"Expected 'C' or 'F', got '{value}'", var)

    value = require_constant_int_tuple(var)
    if len(value) != rank:
        raise _make_type_error(f"Expected tuple of length {rank}, got {len(value)}", var)

    if len(value) == 0:
        return value

    ret = tuple(normalize_axis(x, rank, var) for x in value)
    [(repeating_axis, repeat_count)] = Counter(ret).most_common(1)
    if repeat_count > 1:
        raise _make_type_error(f"Axis order must be a permutation, but axis {repeating_axis}"
                               f" is used at least twice", var)
    return ret


def ensure_tile(var: Var) -> Var[TileTy]:
    ty = var.get_type()
    if not isinstance(ty, TileTy):
        raise _make_type_error(f"Expected a tile, but given value has type {ty}", var)
    return var


def require_tile_type(var: Var) -> TileTy:
    return ensure_tile(var).get_type()


def require_tile_or_tile_tuple_type(var: Var) -> TileTy | TupleTy:
    ty = var.get_type()
    if isinstance(ty, TileTy):
        return ty
    if isinstance(ty, TupleTy) and all(isinstance(x, TileTy) for x in ty.value_types):
        return ty
    raise _make_type_error(f"Expected a tile or a tuple of tiles, but given value has type {ty}",
                           var)


def require_tile_maybe_loose_type(var: Var) \
        -> TileTy | LooselyTypedScalar:
    ty = var.get_loose_type()
    if isinstance(ty, LooselyTypedScalar):
        return ty
    return require_tile_type(var)


def require_0d_tile_type(var: Var) -> TileTy:
    ty = var.get_type()
    if not isinstance(ty, TileTy) or ty.ndim != 0:
        raise _make_type_error(f"Expected a scalar or a 0D tile, but given value has type {ty}",
                               var)
    return ty


def ensure_scalar(var: Var) -> Var[TensorLikeTy]:
    ty = var.get_type()
    if not isinstance(ty, TensorLikeTy) or ty.tensor_shape() != ():
        raise _make_type_error(f"Expected a scalar value, but given value has type {ty}", var)
    return var


def require_scalar_type(var: Var) -> TensorLikeTy:
    return ensure_scalar(var).get_type()


def require_any_vector_type(var: Var) -> TileTy:
    ty = var.get_type()
    if not isinstance(ty, TileTy) or ty.ndim != 1:
        raise _make_type_error(f"Expected a vector but given value has type {ty}", var)
    return ty


def require_any_scalar_or_vector_type(var: Var) -> TileTy:
    ty = var.get_type()
    if not isinstance(ty, TileTy) or ty.ndim not in (0, 1):
        raise _make_type_error(f"Expected a scalar or a vector but given value has type {ty}", var)
    return ty


def require_integer_0d_tile_type(var: Var) -> TileTy:
    ty = require_0d_tile_type(var)
    if not datatype.is_integral(ty.dtype):
        raise _make_type_error(f"Expected an integer scalar, but got {ty}", var)
    return ty


def require_signed_integer_0d_tile_type(var: Var) -> TileTy:
    ty = require_0d_tile_type(var)
    if not datatype.is_integral(ty.dtype) or not datatype.is_signed(ty.dtype):
        raise _make_type_error(f"Expected a signed integer scalar, but got {ty}", var)
    return ty


def require_signed_integer_scalar_type(var: Var) -> TensorLikeTy:
    ty = require_scalar_type(var)
    dtype = ty.tensor_dtype()
    if not datatype.is_integral(dtype) or not datatype.is_signed(dtype):
        raise _make_type_error(f"Expected a signed integer scalar, but got {ty}", var)
    return ty


def require_pointer_tile_type(var: Var) -> TileTy:
    ty = require_tile_type(var)
    if not is_pointer_dtype(ty.dtype):
        raise TileTypeError(f"Expected a pointer, got {ty}")
    return ty


def require_bool_scalar_type(var: Var) -> TensorLikeTy:
    ty = var.get_type()
    if not is_scalar(ty, is_boolean):
        raise _make_type_error(f"Expected a bool, but given value has type {ty}", var)
    return ty


def require_0d_tile_maybe_loose_type(var: Var) -> TileTy | LooselyTypedScalar:
    ty = var.get_loose_type()
    if isinstance(ty, LooselyTypedScalar):
        return ty
    return require_0d_tile_type(var)


def require_array_type(var: Var) -> ArrayTy:
    ty = var.get_type()
    if not isinstance(ty, ArrayTy):
        raise _make_type_error(f"Expected an array, but given value has type {ty}", var)
    return ty


def require_tiled_view_type(var: Var) -> TiledViewTy:
    ty = var.get_type()
    if not isinstance(ty, TiledViewTy):
        raise TileTypeError(f"Expected a tiled view, but given value has type {ty}")
    return ty


def require_raw_array_memory_type(var: Var) -> RawArrayMemoryTy:
    ty = var.get_type()
    if not isinstance(ty, RawArrayMemoryTy):
        raise _make_type_error(f"Expected a RawArrayMemory, but given value has type {ty}", var)
    return ty


def require_list_type(var: Var) -> ListTy:
    ty = var.get_type()
    if not isinstance(ty, ListTy):
        raise _make_type_error(f"Expected a list, but given value has type {ty}", var)
    return ty


def require_tuple_type(var: Var) -> TupleTy:
    ty = var.get_type()
    if not isinstance(ty, TupleTy):
        raise _make_type_error(f"Expected a tuple, but given value has type {ty}", var)
    return ty


def require_dataclass_type(var: Var) -> DataclassTy:
    ty = var.get_type()
    if not isinstance(ty, DataclassTy):
        raise _make_type_error(f"Expected a dataclass instance, but given value has type {ty}", var)
    return ty


def require_index_or_index_tuple_type(var: Var,
                                      allow_nd_tiles: bool = False,
                                      allow_unsigned: bool = False) \
        -> TupleTy | TileTy:
    ty = var.get_type()
    if isinstance(ty, TupleTy):
        item_types = ty.value_types
    else:
        item_types = ty,

    for i, item_ty in enumerate(item_types):
        if not (isinstance(item_ty, TileTy)
                and (allow_nd_tiles or item_ty.ndim == 0)
                and is_integral(item_ty.dtype)
                and (allow_unsigned or is_signed(item_ty.dtype))):
            what = f"item #{i}" if isinstance(ty, TupleTy) else "given value"
            signed = "" if allow_unsigned else "signed "
            if allow_nd_tiles:
                raise _make_type_error(f"Expected a tuple of {signed}integer scalars/tiles"
                                       f" or a single {signed}integer scalar/tile,"
                                       f" but {what} has type {item_ty}", var)
            else:
                raise _make_type_error(f"Expected a tuple of {signed}integers or a single"
                                       f" {signed}integer scalar, but {what} has type {item_ty}",
                                       var)
    return ty


def require_callable_type(var: Var) -> FunctionTy | BoundMethodTy | ClosureTy | DTypeConstructor:
    ty = var.get_type()
    if not isinstance(ty, FunctionTy | BoundMethodTy | ClosureTy | DTypeConstructor):
        raise _make_type_error(f"Expected a callable object, but given value has type {ty}", var)
    return ty


class PrintfValidator:
    # c-format string has the following: %[flags][width][.precision][length]specifier
    # we only support a subset which makes sense in the tile context
    float_specifiers = {'e', 'E', 'f', 'F', 'g', 'G', 'a', 'A'}
    int_specifiers = {'d', 'i', 'u', 'o', 'x', 'X'}
    flags = r"([0 #+-])?"
    width = r"([0-9]+)?"
    precision = r"(\.[0-9]+)?"
    length = r"(hh|h|ll|l)?"
    specifiers = r"([diuoxXeEfFgGaAcspn])"
    pattern = re.compile("%" + flags + width + precision + length + specifiers)

    @classmethod
    def infer_format(cls, dtype: DType) -> str:
        if is_boolean(dtype):
            return '%d'
        elif is_integral(dtype):
            result = '%'
            if dtype.bitwidth == 64:
                result += 'll'
            if is_signed(dtype):
                result += 'd'
            else:
                result += 'u'
            return result
        elif is_float(dtype):
            return '%lf' if dtype.bitwidth == 64 else '%f'
        else:
            raise TileTypeError(f"print(): cannot infer format for dtype {dtype}")

    @classmethod
    def validate_dtype(cls, dtype: DType, specifier: str) -> bool:
        if is_boolean(dtype) or is_integral(dtype):
            return specifier in cls.int_specifiers
        elif is_float(dtype):
            return specifier in cls.float_specifiers
        else:
            return False

    # Python format spec regex: [align][sign][alt][zero][width][.precision][type]
    _py_spec_pattern = re.compile(
        r'(?P<align>[<>^])?'
        r'(?P<sign>[+ -])?'
        r'(?P<alt>\#)?'
        r'(?P<zero>0)?'
        r'(?P<width>[0-9]+)?'
        r'(?:\.(?P<precision>[0-9]+))?'
        r'(?P<type>[diouxXeEfFgGaA])?'
    )

    @staticmethod
    def escape_str(s: str) -> str:
        """Escape a literal string for use in a C printf format (replace % with %%)."""
        return s.replace('%', '%%')

    @classmethod
    def apply_python_spec(cls, py_spec: str, dtype: DType) -> str:
        """Convert a Python format spec to a complete C printf specifier for the given dtype.

        If py_spec omits the type character, it is inferred from dtype.
        If py_spec includes a type character, it is validated against dtype.
        Raises TileTypeError on type mismatch; ValueError on unrecognised spec syntax.
        """
        m = cls._py_spec_pattern.fullmatch(py_spec)
        if m is None or m.group(0) != py_spec:
            raise ValueError(f"print(): unsupported format spec '{py_spec}'")

        align = m.group('align')
        sign = m.group('sign')
        alt = m.group('alt')
        zero = m.group('zero')
        width = m.group('width') or ''
        precision = ('.' + m.group('precision')) if m.group('precision') is not None else ''
        typ = m.group('type')

        flags = ''
        if align == '<':
            flags += '-'
        if sign in ('+', ' '):
            flags += sign
        if alt:
            flags += '#'
        if zero and align != '<':
            flags += '0'

        if typ is None:
            typ = cls.infer_format(dtype)[1:]  # inferred type char, e.g. 'd' from '%d'
        elif not cls.validate_dtype(dtype, typ):
            raise TileTypeError(
                f"print(): format spec '{py_spec}' is incompatible with dtype {dtype}")
        return f'%{flags}{width}{precision}{typ}'

    @classmethod
    def parse_format(cls, format: str, arg_types: Tuple[Union[TileTy, DType], ...]) -> str:
        last_pos = pos = 0
        arg_idx = 0
        tokens = []
        while pos < len(format):
            ch = format[pos]
            if ch == "%":
                tokens.append(format[last_pos:pos])
                last_pos = pos
                # escape "%%"
                if (pos + 1 < len(format) and format[pos + 1] == "%"):
                    pos += 2
                elif (m := cls.pattern.match(format, pos)):
                    # get a format match
                    _, _, _, _, sp = m.groups()
                    fmt = m.group(0)
                    if not (sp in cls.int_specifiers or sp in cls.float_specifiers):
                        raise TileTypeError(f"Specifier {sp} in {fmt} is not supported")
                    # pop argument
                    if arg_idx >= len(arg_types):
                        raise TileTypeError("Not enough arguments for format string")
                    ty = arg_types[arg_idx]
                    # validate arg type against fmt
                    if not cls.validate_dtype(get_dtype(ty), sp):
                        raise TileTypeError(f"Format {fmt} for arg #{arg_idx} got unexpected type of {ty}")  # noqa: E501
                    arg_idx += 1
                    pos = m.end()
                    tokens.append(format[last_pos:pos])
                    last_pos = pos
                else:
                    raise TileTypeError("Invalid format string")
            else:
                pos += 1
        tokens.append(format[last_pos:pos])
        if arg_idx < len(arg_types):
            raise TileTypeError("Too many arguments for format string")
        return "".join(tokens)


class _ErrorContext(NamedTuple):
    function_name: str
    param_name_or_idx: str | int


def _recover_error_context(var: Optional[Var]) -> Optional[_ErrorContext]:
    if var is None:
        return None
    cur_stub_and_args = _current_stub.stub_and_args
    if cur_stub_and_args is None:
        return None
    stub, stub_sig, func_sig, args, kwargs = cur_stub_and_args
    bound_args: inspect.BoundArguments = func_sig.bind(*args, **kwargs)
    for param_name, arg in bound_args.arguments.items():
        if arg is var:
            stub_param = stub_sig.parameters[param_name]
            if stub_param.kind == inspect.Parameter.POSITIONAL_ONLY:
                param_name_or_idx = next(i for i, pname in enumerate(stub_sig.parameters.keys(), 1)
                                         if pname == param_name)
            else:
                param_name_or_idx = param_name
            return _ErrorContext(stub.__name__, param_name_or_idx)
    return None


def _make_type_error(what: str, var: Optional[Var]) -> TileTypeError:
    context = _recover_error_context(var)
    if context is None:
        context_str = ""
    else:
        if isinstance(context.param_name_or_idx, int):
            arg_name = f"#{context.param_name_or_idx}"
        else:
            arg_name = f'"{context.param_name_or_idx}"'
        context_str = f"Invalid argument {arg_name} of {context.function_name}(): "
    return TileTypeError(context_str + what)
