"""`tile_kernels.utils` -- small helpers the model imports directly.

Pure Python on both paths: `align` / `ceil_div` are arithmetic, `dtype_to_str` is a lookup. There
is nothing to route, only to forward, so that a call site does not have to know whether the
compiled package is present.
"""

from __future__ import annotations

from typing import Any

try:
    from tile_kernels.utils import align, ceil_div, dtype_to_str, str_to_dtype

    _HAVE_UTILS = True
except ImportError:  # pragma: no cover - only without tile_kernels installed
    _HAVE_UTILS = False

    def align(value: int, alignment: int) -> int:
        return (value + alignment - 1) // alignment * alignment

    def ceil_div(a: int, b: int) -> int:
        return -(-a // b)

    # The compiled helper returns short names ('bf16', 'fp32'), not torch's ('bfloat16'). Callers
    # build kernel names and dtype strings from them, so the fallback has to match exactly.
    _SHORT = {
        "torch.bfloat16": "bf16",
        "torch.float16": "fp16",
        "torch.float32": "fp32",
        "torch.float8_e4m3fn": "e4m3",
        "torch.float8_e5m2": "e5m2",
        "torch.int32": "int32",
        "torch.int64": "int64",
        "torch.int8": "int8",
        "torch.uint8": "uint8",
    }

    def dtype_to_str(dtype: Any) -> str:
        key = str(dtype)
        if key in _SHORT:
            return _SHORT[key]
        return key.removeprefix("torch.")

    def str_to_dtype(name: str) -> Any:
        import torch

        long_names = {short: long.removeprefix("torch.") for long, short in _SHORT.items()}
        if name in long_names:
            return getattr(torch, long_names[name])
        return getattr(torch, name)


__all__ = ["align", "ceil_div", "dtype_to_str", "str_to_dtype"]