#!/usr/bin/env python3
"""Compile cleanroom/glsl/shader_NN.comp with glslc and regenerate an embedded
shader header with the same interface as src/shaders_embedded.hpp.

Usage: python3 cleanroom/build_embedded.py [--out src/shaders_embedded_clean.hpp]
Env:   GLSLC / SPIRV_VAL override tool paths (default: NDK 27.3 shader-tools).
"""
import argparse
import os
import subprocess
import sys

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
GLSL_DIR = os.path.join(REPO, "cleanroom", "glsl")
NDK_TOOLS = os.path.expanduser(
    "~/Library/Android/sdk/ndk/27.3.13750724/shader-tools/darwin-x86_64")
GLSLC = os.environ.get("GLSLC", os.path.join(NDK_TOOLS, "glslc"))
SPIRV_VAL = os.environ.get("SPIRV_VAL", os.path.join(NDK_TOOLS, "spirv-val"))
REGISTRY_SLOTS = 56


def compile_one(src: str) -> bytes:
    out = src[:-5] + ".spv"
    subprocess.run(
        [GLSLC, "-fshader-stage=comp", "--target-env=vulkan1.1", "-O",
         src, "-o", out],
        check=True)
    subprocess.run([SPIRV_VAL, out], check=True)
    with open(out, "rb") as f:
        return f.read()


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=os.path.join(REPO, "src", "shaders_embedded_clean.hpp"))
    args = ap.parse_args()

    blobs = {}
    for name in sorted(os.listdir(GLSL_DIR)):
        if not name.endswith(".comp"):
            continue
        idx = int(name.replace("shader_", "").replace(".comp", ""))
        blobs[idx] = compile_one(os.path.join(GLSL_DIR, name))
        print(f"  shader_{idx:02d}: {len(blobs[idx])} bytes")

    lines = []
    lines.append("#pragma once")
    lines.append("")
    lines.append("// Generated by cleanroom/build_embedded.py from cleanroom/glsl/.")
    lines.append("// Clean-room reimplementation; see cleanroom/PROCESS.md.")
    lines.append("")
    lines.append("#include <array>")
    lines.append("#include <cstddef>")
    lines.append("#include <cstdint>")
    lines.append("")
    lines.append("namespace bionic_fg::embedded {")
    lines.append("")
    lines.append("struct ShaderBlob {")
    lines.append("    const char* name;")
    lines.append("    const std::uint8_t* data;")
    lines.append("    std::size_t size;")
    lines.append("};")
    lines.append("")
    lines.append("inline constexpr std::uint32_t kSpirvMagic = 0x07230203u;")
    lines.append("")
    for idx in sorted(blobs):
        data = blobs[idx]
        lines.append(f"inline constexpr std::uint8_t shader_{idx:02d}_data[] = {{")
        for i in range(0, len(data), 12):
            chunk = ", ".join(f"0x{b:02x}" for b in data[i:i + 12])
            lines.append(f"    {chunk},")
        lines[-1] = lines[-1].rstrip(",")
        lines.append("};")
        lines.append("")
    lines.append(f"inline constexpr std::array<ShaderBlob, {REGISTRY_SLOTS}> kShaderRegistry = {{{{")
    for idx in range(REGISTRY_SLOTS):
        if idx in blobs:
            lines.append(f'    ShaderBlob{{"shader_{idx:02d}.spv", shader_{idx:02d}_data, sizeof(shader_{idx:02d}_data)}},')
        else:
            lines.append(f'    ShaderBlob{{"shader_{idx:02d}.spv", nullptr, 0}},')
    lines[-1] = lines[-1].rstrip(",")
    lines.append("}};")
    lines.append("")
    lines.append("inline bool IsValidSpirv(const ShaderBlob& blob) {")
    lines.append("    if (blob.size < sizeof(std::uint32_t)) return false;")
    lines.append("    const std::uint32_t magic =")
    lines.append("        static_cast<std::uint32_t>(blob.data[0]) |")
    lines.append("        (static_cast<std::uint32_t>(blob.data[1]) << 8u) |")
    lines.append("        (static_cast<std::uint32_t>(blob.data[2]) << 16u) |")
    lines.append("        (static_cast<std::uint32_t>(blob.data[3]) << 24u);")
    lines.append("    return magic == kSpirvMagic;")
    lines.append("}")
    lines.append("")
    lines.append("} // namespace bionic_fg::embedded")
    with open(args.out, "w") as f:
        f.write("\n".join(lines) + "\n")
    print(f"wrote {args.out} with {len(blobs)}/{REGISTRY_SLOTS} populated slots")
    return 0


if __name__ == "__main__":
    sys.exit(main())