#!/usr/bin/env python3
"""Compile virtualgeometry GLSL compute/fragment shaders into .spv and .inc.

Run from the repo root (or pass a directory):
    python3 scripts/compile_virtualgeometry_shaders.py

Emits, next to each .glsl/.comp/.frag/.vert source:
  - <stem>.spv                 (raw SPIR-V binary, loaded at runtime)
  - <stem>_spv.inc             (embedded uint32 array for the backend)
The .inc files are #included by virtualgeometry/vulkan/VulkanVirtualGeometry.cpp.

Only VG_COMMON_GLSL is #included by the passes; the include dir is added via -I.
"""
import os
import subprocess
import sys
import struct

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SHADERS = os.path.join(ROOT, "src", "modules", "virtualgeometry", "shaders")

STAGES = {".comp": "compute", ".vert": "vertex", ".frag": "fragment"}


def find_compiler():
    for name in ("glslc", "glslangValidator"):
        path = subprocess.run(["which", name], capture_output=True, text=True).stdout.strip()
        if path:
            return name, path
    return None, None


def compile_one(compiler, src, out_spv, stage):
    inc = os.path.join(SHADERS)
    if compiler == "glslc":
        cmd = ["glslc", "-I", inc, "-fshader-stage=" + stage, "-o", out_spv, src]
    else:
        cmd = ["glslangValidator", "-I", inc, "-S", stage, "-o", out_spv, src]
    res = subprocess.run(cmd, capture_output=True, text=True)
    if res.returncode != 0:
        print("FAIL", src)
        print(res.stdout)
        print(res.stderr)
        return False
    return True


def emit_inc(spv, inc, symbol):
    with open(spv, "rb") as f:
        data = f.read()
    assert len(data) % 4 == 0, "SPIR-V size must be word-aligned"
    words = struct.unpack("<%dI" % (len(data) // 4), data)
    lines = ["// Generated by compile_virtualgeometry_shaders.py. Do not edit."]
    lines.append("#pragma once")
    lines.append("static const uint32_t %s[] = {" % symbol)
    for i in range(0, len(words), 6):
        chunk = words[i:i + 6]
        lines.append("    " + ", ".join("0x%08x" % w for w in chunk) + ",")
    lines.append("};")
    lines.append("static const size_t %s_count = sizeof(%s) / sizeof(%s[0]);" % (symbol, symbol, symbol))
    with open(inc, "w") as f:
        f.write("\n".join(lines) + "\n")


def main():
    compiler, _ = find_compiler()
    if not compiler:
        print("No GLSL compiler (glslc / glslangValidator) found on PATH.")
        return 1

    base_dir = sys.argv[1] if len(sys.argv) > 1 else SHADERS
    for root, _, files in os.walk(base_dir):
        for name in sorted(files):
            ext = os.path.splitext(name)[1]
            if ext not in STAGES or name == "vg_common.glsl":
                continue
            stem = os.path.splitext(name)[0]
            src = os.path.join(root, name)
            # Suffix with the stage so vg_resolve.vert and vg_resolve.frag don't collide.
            spv = os.path.join(root, "%s_%s.spv" % (stem, STAGES[ext]))
            inc = os.path.join(root, "%s_%s_spv.inc" % (stem, STAGES[ext]))
            if not compile_one(compiler, src, spv, STAGES[ext]):
                return 1
            emit_inc(spv, inc, "%s_%s_spv" % (stem, STAGES[ext]))
            print("compiled", os.path.relpath(spv, ROOT))
    return 0


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