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 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 gen_shader_ext(build_type):
    if build_type == 'glsl':
        return 'glsl'
    elif build_type == 'spirv':
        return 'spirv'
    elif build_type == 'dxil':
        return 'dxil'
    elif build_type == 'dxbc':
        return 'dxbc'
    elif build_type == 'metallib':
        return 'metallib'
    elif build_type == 'wgsl':
        return 'wgsl'
    else:
        return 'dat'

def gen_slangc_cmd(input_file, entry, output_file, target_type, args):
    slangc = args.slangc
    cmd = [slangc, input_file, "-entry", entry, "-o", output_file, "-target", target_type]
    if args.debug:
        cmd.append("-g3")
    else:
        cmd.append("-O3")

    if target_type == 'glsl':
        cmd.extend(["-profile", "glsl_460"])
    if target_type == 'spirv':
        cmd.extend(["-profile", "glsl_460"])
        cmd.extend(["-capability", "glsl_spirv_1_6"])
    if target_type == 'dxbc':
        cmd.extend(["-profile", "sm_5_1"])
    if target_type == 'dxil':
        cmd.extend(["-profile", "sm_6_6"])
    if target_type == 'metallib':
        cmd.extend(["-capability", "metallib_3_1"])
    return cmd

def compile_slang(input_file, target_types, output_dir, args):
    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 target_type, enabled in target_types:
            if not enabled:
                continue
            for entry in entries:
                output_file = os.path.join(abs_path, f"{base}_{entry}.{gen_shader_ext(target_type)}")
                cmd = gen_slangc_cmd(input_file, entry, output_file, target_type, args)
                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 slang shaders using slangc")
    parser.add_argument("--shader-list", help="Input shader list .txt file")
    parser.add_argument("--output-dir", help="Output directory")
    parser.add_argument("--slangc", default="slangc", help="Path to slangc")
    parser.add_argument("--debug", action="store_true", help="Compile in debug mode")
    parser.add_argument("--opengl", action="store_true", help="Compile Slang for OpenGL")
    parser.add_argument("--vulkan", action="store_true", help="Compile Slang for Vulkan")
    parser.add_argument("--d3d11", action="store_true", help="Compile Slang for D3D11")
    parser.add_argument("--d3d12", action="store_true", help="Compile Slang for D3D12")
    parser.add_argument("--metal", action="store_true", help="Compile Slang for Metal")

    args = parser.parse_args()
    target_types = [
        ['glsl', args.opengl],
        ['spirv', args.vulkan],
        ['dxbc', args.d3d11],
        ['dxil', args.d3d12],
        ['metallib', args.metal],
    ]

    output_dir = args.output_dir or "shaders"
    shader_dir = args.shader_list or "shader_paths.txt"

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

    slang_ext = ['.slang']

    all_success = True
    
    for shader_path in shader_paths:
        # Compile Slang
        for file in find_shader_files(shader_path, slang_ext):
            if not compile_slang(file, target_types, output_dir, args):
                all_success = False

    if not all_success:
        sys.exit(1)

if __name__ == "__main__":
    main()