"""Checks the tilelang stub provides exactly what platform-independent DeepInfer code uses."""

import sys

import torch

import tilelang
import tilelang.language as T

FAILS = []


def chk(name, cond):
    print(f"[{'ok ' if cond else 'FAIL'}] {name}")
    if not cond:
        FAILS.append(name)


chk("PassConfigKey.TL_ENABLE_FAST_MATH", isinstance(tilelang.PassConfigKey.TL_ENABLE_FAST_MATH, str))
chk("T.float8_e4m3fn", T.float8_e4m3fn == "float8_e4m3fn")
chk("dtype map", tilelang.TL_DTYPE_TO_TORCH[T.float8_e4m3fn] is torch.float8_e4m3fn)
chk("norm.py dtype strings", T.float32 == "float" and T.bfloat16 == "bfloat16")
chk("max_value e4m3", T.max_value(T.float8_e4m3fn) == 448.0)
chk("min_value fp32", T.min_value(T.float32) == torch.finfo(torch.float32).min)

# the pass-config dict pattern of tilelang_ops/compress.py
cfg = {}
cfg[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] = True
chk("pass_configs assignment", cfg[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] is True)

try:
    T.some_dsl_primitive
    FAILS.append("unknown attr should raise")
except AttributeError as e:
    chk("unknown T attribute raises AttributeError", "Neuron stub" in str(e))

print("\nFAILURES:", FAILS if FAILS else "none")
sys.exit(1 if FAILS else 0)