#!/usr/bin/env python3
"""Compile every Metal shader source in the crate with the offline compiler.

    check-metal-shaders.py [--list]

`MetalCompute::compile_pipeline` builds these sources at *runtime*, so a
syntax or type error in a kernel is invisible to `cargo check` and to the
per-fabric error guard. It surfaces only when something actually dispatches
that pipeline -- and since no test yet runs a Metal kernel, that means it
never surfaces at all. `xcrun metal` is the same frontend the runtime uses,
so compiling each source offline closes exactly that gap.

This checks that the shaders *compile*, which is still not that they compute
the right thing. Numeric checks need a test that dispatches them on a device.
"""
import re, subprocess, sys, tempfile, pathlib

ROOT = pathlib.Path(__file__).resolve().parent.parent
SHADER_FILES = [
    "jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/hal/metal/shader.rs",
]
# A const is a compilable translation unit only if it declares entry points.
ENTRY = re.compile(r"\bkernel\s+void\s+\w+", re.S)
# Raw strings appear as r"..." or r#"..."# (clippy::needless_raw_string_hashes
# rewrites the second form to the first), so match any hash count. Being
# coupled to the raw-string style once took this check from 44 sources to 2,
# silently, while it still reported success -- hence MIN_SOURCES below.
CONST = re.compile(
    r'pub const (?P<name>[A-Z_0-9]+): &str = r(?P<hashes>#*)"(?P<body>.*?)"(?P=hashes);',
    re.S,
)
# This check is worthless if it finds nothing, and nearly worthless if it finds
# a handful. The count may rise; it may not fall.
MIN_SOURCES = 44


def sources():
    for rel in SHADER_FILES:
        path = ROOT / rel
        if not path.exists():
            print(f"::error::missing shader file {rel}")
            sys.exit(1)
        text = path.read_text()
        for m in CONST.finditer(text):
            line = text[: m.start()].count("\n") + 1
            yield rel, line, m.group("name"), m.group("body")


def main(argv: list[str]) -> int:
    found = list(sources())
    if not found:
        print("::error::no shader sources matched -- the check did not run")
        return 1

    compilable = [s for s in found if ENTRY.search(s[3])]
    skipped = [s for s in found if not ENTRY.search(s[3])]
    if not compilable:
        print("::error::no source declared a kernel entry point")
        return 1
    if len(found) < MIN_SOURCES:
        print(f"::error::found only {len(found)} shader sources, expected at "
              f"least {MIN_SOURCES}. Either sources were removed, or the way "
              f"they are written changed and this check stopped seeing them. "
              f"A check that silently inspects almost nothing still reports "
              f"success, so this is a failure.")
        return 1

    if "--list" in argv:
        for rel, line, name, body in found:
            kind = "kernels" if ENTRY.search(body) else "fragment"
            print(f"  {name:<28} {kind:<8} {rel}:{line}")
        return 0

    failures = []
    with tempfile.TemporaryDirectory() as td:
        tmp = pathlib.Path(td)
        for rel, line, name, body in compilable:
            src = tmp / f"{name}.metal"
            src.write_text(body)
            proc = subprocess.run(
                ["xcrun", "-sdk", "macosx", "metal", "-c",
                 str(src), "-o", str(tmp / f"{name}.air")],
                capture_output=True, text=True,
            )
            if proc.returncode != 0:
                failures.append((rel, line, name, proc.stderr.strip()))

    print(f"metal shaders: {len(compilable)} compiled, "
          f"{len(failures)} failed, {len(skipped)} skipped (no entry point)")
    for rel, line, name in ((r, l, n) for r, l, n, _ in sorted(failures)):
        print(f"  FAIL {name}  {rel}:{line}")

    if failures:
        print()
        for rel, line, name, err in failures:
            print(f"--- {name} ({rel}:{line}) ---")
            # The first few diagnostics are the actionable ones.
            print("\n".join(err.splitlines()[:25]))
            print()
        print(f"::error::{len(failures)} Metal shader source(s) do not compile. "
              f"compile_pipeline builds these at runtime, so this is a real "
              f"dispatch failure, not a style issue.")
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))