import argparse
import os
import subprocess
import sys
import re

def find_shader_files(input_dir, extensions):
    for root, _, files in os.walk(input_dir):
        for file in files:
            if any(file.endswith(ext) for ext in extensions):
                yield os.path.join(root, file)

def has_glsl_main(input_file):
    try:
        with open(input_file, 'r', encoding='utf-8') as f:
            content = f.read()
        return "void main" in content
    except Exception as e:
        print(f"Error parsing {input_file}: {str(e)}")
        return False

def compile_glsl(input_file, output_dir, glslang_validator):
    try:
        if not has_glsl_main(input_file):
            print(f"Skipping {input_file}: No main function found")
            return True

        base = os.path.splitext(os.path.basename(input_file))[0]
        ext = os.path.splitext(input_file)[1][1:]
        abs_path = os.path.abspath(output_dir)
        output_file = os.path.join(abs_path, f"{base}.{ext}.spv")
        
        os.makedirs(abs_path, exist_ok=True)
        
        cmd = [glslang_validator, "-V", input_file, "-o", output_file]
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        print(f"Compiled GLSL: {input_file} -> {output_file}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error compiling {input_file}:")
        print(e.stderr)
        return False
    except Exception as e:
        print(f"Unexpected error with {input_file}: {str(e)}")
        return False

def find_slang_entries(input_file):
    try:
        with open(input_file, 'r', encoding='utf-8') as f:
            content = f.read()
        pattern = r'\[shader\(\s*"(?:\w+)"\s*\)\]\s*\n\s*\w+\s+(\w+)\s*\('
        return list(set(re.findall(pattern, content)))
    except Exception as e:
        print(f"Error parsing {input_file}: {str(e)}")
        return []

def compile_slang(input_file, output_dir, slangc):
    try:
        entries = find_slang_entries(input_file)
        if not entries:
            print(f"Skipping {input_file}: No shader entries found")
            return True
        
        base = os.path.splitext(os.path.basename(input_file))[0]

        abs_path = os.path.abspath(output_dir)

        os.makedirs(abs_path, exist_ok=True)
        
        success = True
        for entry in entries:
            output_file = os.path.join(abs_path, f"{base}_{entry}.spv")
            cmd = [slangc, input_file, "-entry", entry, "-o", output_file, "-target", "spirv"]
            try:
                subprocess.run(cmd, check=True, capture_output=True, text=True)
                print(f"Compiled Slang: {input_file}:{entry} -> {output_file}")
            except subprocess.CalledProcessError as e:
                print(f"Error compiling {input_file}:{entry}")
                print(e.stderr)
                success = False
        return success
    except Exception as e:
        print(f"Unexpected error with {input_file}: {str(e)}")
        return False

def main():
    parser = argparse.ArgumentParser(description="Compile shaders to SPIR-V")
    parser.add_argument("--output-dir", help="Output directory for SPIR-V files")
    parser.add_argument("--glslang", default="glslangValidator", help="Path to glslangValidator")
    parser.add_argument("--slangc", default="slangc", help="Path to slangc")
    args = parser.parse_args()

    output_dir = args.output_dir or "shaders"

    # 读取当前同级目录下的shader_paths.txt
    with open("shader_paths.txt", 'r') as f:
        shader_paths = f.readlines()

    glsl_ext = [".vert", ".frag", ".comp", ".geom", ".tesc", ".tese", ".glsl"]
    slang_ext = [".slang"]

    all_success = True
    
    for shader_path in shader_paths:
        # Compile GLSL
        for file in find_shader_files(shader_path, glsl_ext):
            if not compile_glsl(file, output_dir, args.glslang):
                all_success = False
    
        # Compile Slang
        for file in find_shader_files(shader_path, slang_ext):
            if not compile_slang(file, output_dir, args.slangc):
                all_success = False

    if not all_success:
        sys.exit(1)

if __name__ == "__main__":
    main()