#!/usr/bin/env python3
"""
Shader compilation script for Vulkan SPIR-V shaders
"""

import os
import struct
import subprocess
import sys


def find_glsl_compiler():
    home = os.path.expanduser("~")
    custom_paths = [
        os.path.join(home, "sdk/glslang-13.0.0/bin/glslang"),
        os.path.join(home, "sdk/glslang-13.0.0/bin/glslangValidator"),
    ]

    for compiler_path in custom_paths:
        if os.path.exists(compiler_path) and os.access(compiler_path, os.X_OK):
            print("Found custom compiler: " + compiler_path)
            return compiler_path

    compilers = ["glslc", "glslangValidator", "glslang"]
    for compiler in compilers:
        try:
            subprocess.check_output([compiler, "--version"], stderr=subprocess.STDOUT)
            return compiler
        except (subprocess.CalledProcessError, FileNotFoundError, OSError):
            continue
    return None


def compile_shader(compiler, shader_file, output_file):
    print("Compiling " + shader_file + " -> " + output_file)

    env = os.environ.copy()
    if "/usr/local/lib64" not in env.get("LD_LIBRARY_PATH", ""):
        env["LD_LIBRARY_PATH"] = "/usr/local/lib64:" + env.get("LD_LIBRARY_PATH", "")

    compiler_name = os.path.basename(compiler)

    try:
        if compiler_name == "glslc":
            cmd = [compiler, shader_file, "-o", output_file]
        else:
            cmd = [compiler, "-V", shader_file, "-o", output_file]

        subprocess.check_output(cmd, stderr=subprocess.STDOUT, env=env)
        print("  SUCCESS")
        return True
    except subprocess.CalledProcessError as e:
        print("  FAILED: " + e.output.decode('utf-8', errors='ignore'))
        return False
    except (FileNotFoundError, OSError):
        print("  FAILED: Compiler not found")
        return False


def generate_cpp_include(spirv_file, output_inc):
    with open(spirv_file, "rb") as f:
        spirv_data = f.read()

    num_words = len(spirv_data) // 4
    if len(spirv_data) % 4 != 0:
        spirv_data += b"\x00" * (4 - len(spirv_data) % 4)
        num_words = len(spirv_data) // 4

    words = struct.unpack("<{}I".format(num_words), spirv_data)

    cpp_code = "// Auto-generated from {}\n".format(os.path.basename(spirv_file))
    cpp_code += "// Size: {} bytes ({} words)\n".format(len(spirv_data), num_words)

    for i in range(0, num_words, 8):
        chunk = words[i : i + 8]
        if i == 0:
            cpp_code += "    "
        else:
            cpp_code += ",\n    "
        cpp_code += ", ".join("0x{:08x}u".format(w) for w in chunk)

    with open(output_inc, "w") as f:
        f.write(cpp_code + "\n")

    print("  Generated " + output_inc)


def main():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    shaders_dir = os.path.join(script_dir, "shaders")
    spirv_dir = os.path.join(script_dir, "src", "shaders_spirv")

    os.makedirs(spirv_dir, exist_ok=True)

    compiler = find_glsl_compiler()
    if not compiler:
        print("ERROR: No GLSL compiler found!")
        print("Or install to: ~/sdk/glslang-13.0.0/bin/glslang")
        sys.exit(1)

    print("Using {} compiler\n".format(compiler))
    print("=" * 60)

    shaders = [
        ("background.vert", "background_vert.inc"),
        ("background.frag", "background_frag.inc"),
        ("geometry.vert", "geometry_vert.inc"),
        ("geometry.frag", "geometry_frag.inc"),
        ("text.vert", "text_vert.inc"),
        ("text.frag", "text_frag.inc"),
    ]

    success_count = 0
    fail_count = 0

    for shader_name, inc_name in shaders:
        shader_path = os.path.join(shaders_dir, shader_name)
        spirv_path = os.path.join(spirv_dir, shader_name + ".spv")
        inc_path = os.path.join(spirv_dir, inc_name)

        if not os.path.exists(shader_path):
            print("WARNING: {} not found, skipping".format(shader_path))
            fail_count += 1
            continue

        if compile_shader(compiler, shader_path, spirv_path):
            generate_cpp_include(spirv_path, inc_path)
            success_count += 1
        else:
            fail_count += 1

    print("=" * 60)
    print("\nCompilation complete: {} succeeded, {} failed".format(success_count, fail_count))

    if fail_count > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()