#!/usr/bin/env python3
"""Regenerate src/vr/vr_shaders_spirv.h from shaders/vr with glslc.

The header is checked in so a build never needs a shader compiler. The SMAA
passes include SMAA.hlsl from the third_party/smaa submodule. Run this
after editing a shader; VrShaders.SpirvMatchesSources checks it when glslc is
installed.
"""

import argparse
import pathlib
import struct
import subprocess
import sys
import tempfile

ROOT = pathlib.Path(__file__).resolve().parent.parent
SHADERS = [
    ("kFullscreenVertexSpirv", "shaders/vr/fullscreen.vert"),
    ("kFxaaFragmentSpirv", "shaders/vr/fxaa.frag"),
    ("kSmaaEdgesVertexSpirv", "shaders/vr/smaa_edges.vert"),
    ("kSmaaEdgesFragmentSpirv", "shaders/vr/smaa_edges.frag"),
    ("kSmaaWeightsVertexSpirv", "shaders/vr/smaa_weights.vert"),
    ("kSmaaWeightsFragmentSpirv", "shaders/vr/smaa_weights.frag"),
    ("kSmaaBlendVertexSpirv", "shaders/vr/smaa_blend.vert"),
    ("kSmaaBlendFragmentSpirv", "shaders/vr/smaa_blend.frag"),
]
# SMAA.hlsl comes from the third_party/smaa submodule.
INCLUDES = [ROOT / "third_party/smaa"]
OUTPUT = ROOT / "src/vr/vr_shaders_spirv.h"


def compile_words(glslc, source):
    with tempfile.TemporaryDirectory() as directory:
        output = pathlib.Path(directory) / "shader.spv"
        subprocess.run(
            [glslc, "-O", "--target-env=vulkan1.0"] +
            ["-I" + str(path) for path in INCLUDES] +
            [str(ROOT / source), "-o", str(output)],
            check=True)
        data = output.read_bytes()
    return struct.unpack("<%dI" % (len(data) // 4), data)


def render(glslc):
    lines = [
        "// Generated by scripts/generate_vr_shaders.py from shaders/vr. Do not edit.",
        "#ifndef MOCKTAIL_VR_VR_SHADERS_SPIRV_H_",
        "#define MOCKTAIL_VR_VR_SHADERS_SPIRV_H_",
        "",
        "#include <cstdint>",
        "",
        "namespace mocktail::vr {",
        "",
    ]
    for name, source in SHADERS:
        words = compile_words(glslc, source)
        lines.append("// %s" % source)
        lines.append("inline constexpr std::uint32_t %s[] = {" % name)
        for start in range(0, len(words), 6):
            chunk = ", ".join("0x%08x" % word for word in words[start:start + 6])
            lines.append("    %s," % chunk)
        lines.append("};")
        lines.append("")
    lines += ["}  // namespace mocktail::vr", "",
              "#endif  // MOCKTAIL_VR_VR_SHADERS_SPIRV_H_", ""]
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--glslc", default="glslc")
    parser.add_argument("--check", action="store_true",
                        help="exit 1 when the header is out of date")
    arguments = parser.parse_args()
    text = render(arguments.glslc)
    if arguments.check:
        current = OUTPUT.read_text() if OUTPUT.exists() else ""
        if current != text:
            print("%s is out of date; run %s" % (OUTPUT, __file__),
                  file=sys.stderr)
            return 1
        return 0
    OUTPUT.write_text(text)
    return 0


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