#!/usr/bin/env python3
# This script is part of the 64k demo project.
# Offline WGSL validation: compose #includes then run naga-cli.
# Usage: ./scripts/validate_shaders.py [--install]

import os
import sys
import subprocess
import tempfile

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# Mirrors InitShaderComposer() in src/effects/shaders.cc
# snippet_name -> path relative to REPO
SNIPPETS = {
    "common_uniforms":           "src/shaders/common_uniforms.wgsl",
    "sequence_uniforms":         "src/shaders/sequence_uniforms.wgsl",
    "postprocess_inline":        "src/shaders/postprocess_inline.wgsl",
    "camera_common":             "src/shaders/camera_common.wgsl",
    "math/sdf_shapes":           "src/shaders/math/sdf_shapes.wgsl",
    "math/sdf_utils":            "src/shaders/math/sdf_utils.wgsl",
    "math/common_utils":         "src/shaders/math/common_utils.wgsl",
    "math/noise":                "src/shaders/math/noise.wgsl",
    "math/color":                "src/shaders/math/color.wgsl",
    "math/utils":                "src/shaders/math/utils.wgsl",
    "lighting":                  "src/shaders/lighting.wgsl",
    "ray_box":                   "src/shaders/ray_box.wgsl",
    "ray_triangle":              "src/shaders/ray_triangle.wgsl",
    "render/shadows":            "src/shaders/render/shadows.wgsl",
    "render/scene_query_bvh":    "src/shaders/render/scene_query_bvh.wgsl",
    "render/scene_query_linear": "src/shaders/render/scene_query_linear.wgsl",
    "render/lighting_utils":     "src/shaders/render/lighting_utils.wgsl",
    "render/scratch_lines":      "src/shaders/render/scratch_lines.wgsl",
    "render/fullscreen_vs":      "src/shaders/render/fullscreen_vs.wgsl",
    "render/fullscreen_uv_vs":   "src/shaders/render/fullscreen_uv_vs.wgsl",
    "render/raymarching":        "src/shaders/render/raymarching.wgsl",
    "render/raymarching_id":     "src/shaders/render/raymarching_id.wgsl",
    "render/mesh":               "workspaces/main/shaders/mesh_render.wgsl",
}

# Shaders that use runtime ShaderComposer substitutions (e.g. scene_query_bvh
# vs scene_query_linear) — they reference symbols injected at runtime and
# cannot be validated standalone.
SUBSTITUTION_SHADERS = {
    "workspaces/main/shaders/renderer_3d.wgsl",
    "workspaces/main/shaders/masked_cube.wgsl",
    "workspaces/main/shaders/mesh_render.wgsl",
}

# All .wgsl files with entry points that are NOT snippets (i.e. complete shaders)
MAIN_SHADERS = [
    "src/effects/chroma_aberration.wgsl",
    "src/effects/distort.wgsl",
    "src/effects/ellipse.wgsl",
    "src/effects/flash.wgsl",
    "src/effects/gaussian_blur.wgsl",
    "src/effects/ntsc.wgsl",
    "src/effects/particle_compute.wgsl",
    "src/effects/particle_render.wgsl",
    "src/effects/particle_spray_compute.wgsl",
    "src/effects/rotating_cube.wgsl",
    "src/effects/scene1.wgsl",
    "src/effects/scene2.wgsl",
    "src/effects/scratch.wgsl",
    "src/effects/solarize.wgsl",
    "src/effects/vignette.wgsl",
    "src/shaders/combined_postprocess.wgsl",
    "src/shaders/gaussian_blur.wgsl",
    "src/shaders/heptagon.wgsl",
    "src/shaders/passthrough.wgsl",
    "src/shaders/skybox.wgsl",
    "src/shaders/compute/gen_blend.wgsl",
    "src/shaders/compute/gen_grid.wgsl",
    "src/shaders/compute/gen_mask.wgsl",
    "src/shaders/compute/gen_noise.wgsl",
    "src/shaders/compute/gen_perlin.wgsl",
    "workspaces/main/shaders/circle_mask_compute.wgsl",
    "workspaces/main/shaders/circle_mask_render.wgsl",
    "workspaces/main/shaders/main_shader.wgsl",
    "workspaces/main/shaders/masked_cube.wgsl",
    "workspaces/main/shaders/mesh_render.wgsl",
    "workspaces/main/shaders/renderer_3d.wgsl",
    "workspaces/main/shaders/visual_debug.wgsl",
]


def compose(source, included=None):
    """Recursively resolve #include directives using SNIPPETS map."""
    if included is None:
        included = set()
    result = []
    for line in source.splitlines():
        if line.strip().startswith("#include "):
            s = line.find('"')
            e = line.find('"', s + 1)
            if s == -1 or e == -1:
                result.append(line)
                continue
            name = line[s + 1:e]
            if name in included:
                continue
            included.add(name)
            if name not in SNIPPETS:
                result.append(f"// ERROR: unknown snippet '{name}'")
                continue
            snippet_path = os.path.join(REPO, SNIPPETS[name])
            try:
                snippet_src = open(snippet_path).read()
                result.append(f"// --- {name} ---")
                result.append(compose(snippet_src, included))
                result.append(f"// --- end {name} ---")
            except FileNotFoundError:
                result.append(f"// ERROR: missing file for snippet '{name}'")
        else:
            result.append(line)
    return "\n".join(result)


def validate(rel_path):
    path = os.path.join(REPO, rel_path)
    if not os.path.exists(path):
        return None, "file not found"
    composed = compose(open(path).read())
    with tempfile.NamedTemporaryFile(mode="w", suffix=".wgsl", delete=False) as f:
        f.write(composed)
        tmp = f.name
    try:
        r = subprocess.run(["naga", tmp], capture_output=True, text=True)
        if r.returncode == 0:
            return True, None
        return False, (r.stderr or r.stdout).strip()
    finally:
        os.unlink(tmp)


def main():
    if "--install" in sys.argv:
        os.execvp("cargo", ["cargo", "install", "naga-cli"])

    if subprocess.run(["which", "naga"], capture_output=True).returncode != 0:
        print("naga not found. Run: cargo install naga-cli  (or: ./scripts/validate_shaders.sh --install)")
        sys.exit(1)

    passed = failed = skipped = 0
    for shader in MAIN_SHADERS:
        label = shader.split("/", 1)[-1]
        if shader in SUBSTITUTION_SHADERS:
            print(f"  skip  {label}  (runtime substitution)")
            skipped += 1
            continue
        ok, err = validate(shader)
        if ok is None:
            print(f"  SKIP  {label}")
            skipped += 1
        elif ok:
            print(f"  ok    {label}")
            passed += 1
        else:
            print(f"  FAIL  {label}")
            for line in err.splitlines()[:8]:
                print(f"        {line}")
            failed += 1

    print(f"\n{passed} passed, {failed} failed, {skipped} skipped")
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()