from typing import List, Tuple, Iterator, Optional
from pathlib import Path
import argparse
import subprocess
import sys
import re

# 设置控制台输出编码
if sys.platform.startswith('win'):
    # Windows系统下设置
    sys.stdout.reconfigure(encoding='utf-8')
    sys.stderr.reconfigure(encoding='utf-8')

def print_utf8(message: str):
    """
    以UTF-8编码打印消息

    Args:
        message: 要打印的消息
    """
    print(message.encode('utf-8').decode(sys.stdout.encoding))

# 常量定义
SHADER_EXTENSIONS = {
    'glsl': 'glsl',
    'spirv': 'spirv',
    'dxil': 'dxil',
    'dxbc': 'dxbc',
    'metallib': 'metallib',
    'wgsl': 'wgsl'
}

# 不同目标平台的编译配置
TARGET_PROFILES = {
    'glsl': ['-profile', 'glsl_460'],
    'spirv': ['-profile', 'glsl_460', '-capability', 'glsl_spirv_1_6'],
    'dxbc': ['-profile', 'sm_5_1'],
    'dxil': ['-profile', 'sm_6_6'],
    'metallib': ['-capability', 'metallib_3_1']
}

def find_shader_files(input_dir: Path, extensions: List[str]) -> Iterator[Path]:
    """
    递归查找指定目录下的着色器文件
    
    Args:
        input_dir: 输入目录路径
        extensions: 着色器文件扩展名列表
    
    Yields:
        符合扩展名的着色器文件路径
    """
    for file_path in Path(input_dir).rglob('*'):
        if file_path.suffix in extensions:
            yield file_path

def find_slang_entries(input_file: Path) -> List[str]:
    """
    从着色器文件中提取入口点函数名
    
    Args:
        input_file: 着色器文件路径
    
    Returns:
        入口点函数名列表
    """
    # 匹配 [shader("xxx")] 形式的着色器入口点声明
    pattern = re.compile(r'\[shader\(\s*"(?:\w+)"\s*\)\]\s*\n\s*\w+\s+(\w+)\s*\(')
    try:
        content = input_file.read_text(encoding='utf-8')
        return list(set(pattern.findall(content)))
    except Exception as e:
        print_utf8(f"**错误**: 解析文件 {input_file} 失败: {e}")
        return []

def get_shader_extension(build_type: str) -> str:
    """
    根据构建类型获取对应的着色器文件扩展名
    
    Args:
        build_type: 构建类型
    
    Returns:
        对应的文件扩展名
    """
    return SHADER_EXTENSIONS.get(build_type, 'dat')

def create_compiler_command(
    input_file: Path,
    entry: str,
    output_file: Path,
    target_type: str,
    args: argparse.Namespace
) -> List[str]:
    """
    生成着色器编译命令
    
    Args:
        input_file: 输入文件路径
        entry: 着色器入口点
        output_file: 输出文件路径
        target_type: 目标平台类型
        args: 命令行参数
    
    Returns:
        编译命令列表
    """
    cmd = [args.slangc,
           str(input_file),
           "-entry", entry,
           "-o", str(output_file),
           "-target", target_type,
           "-g3" if args.debug else "-O3"
    ]
    
    # 添加优化或调试标志

    # 添加目标平台特定的编译选项
    if target_type in TARGET_PROFILES:
        cmd.extend(TARGET_PROFILES[target_type])
    
    return cmd

def needs_recompile(input_file: Path, output_file: Path) -> bool:
    """
    检查是否需要重新编译着色器
    
    Args:
        input_file: 输入文件路径
        output_file: 输出文件路径
    
    Returns:
        是否需要重新编译
    """
    if not output_file.exists():
        return True
    try:
        return input_file.stat().st_mtime > output_file.stat().st_mtime
    except OSError:
        return True

def compile_shader(
    input_file: Path,
    target_types: List[Tuple[str, bool]],
    output_dir: Path,
    args: argparse.Namespace
) -> bool:
    """
    编译单个着色器文件
    
    Args:
        input_file: 输入文件路径
        target_types: 目标平台类型列表
        output_dir: 输出目录
        args: 命令行参数
    
    Returns:
        编译是否成功
    """
    try:
        # 获取着色器入口点
        entries = find_slang_entries(input_file)
        if not entries:
            print_utf8(f"**跳过**: {input_file} - 未找到着色器入口点")
            return True

        # 创建输出目录
        output_dir.mkdir(parents=True, exist_ok=True)
        base = input_file.stem
        success = True

        # 针对每个目标平台编译
        for target_type, enabled in target_types:
            if not enabled:
                continue
                
            for entry in entries:
                output_file = output_dir / f"{base}_{entry}.{get_shader_extension(target_type)}"
                
                # 检查是否需要重新编译
                if not needs_recompile(input_file, output_file):
                    print_utf8(f"**跳过**: {output_file} - 文件已是最新")
                    continue

                # 执行编译
                cmd = create_compiler_command(input_file, entry, output_file, target_type, args)
                try:
                    result = subprocess.run(cmd, check=True, capture_output=True, text=True)
                    print_utf8(f"**成功**: 编译 {input_file}:{entry} -> {output_file}")
                except subprocess.CalledProcessError as e:
                    print_utf8(f"**错误**: 编译 {input_file}:{entry} 失败")
                    print_utf8(e.stderr)
                    success = False

        return success
    except Exception as e:
        print_utf8(f"**错误**: 处理 {input_file} 时发生异常: {e}")
        return False

def main():
    """
    主函数：解析命令行参数并执行编译流程
    """
    # 设置UTF-8编码
    if sys.platform.startswith('win'):
        # Windows下设置控制台代码页
        subprocess.run(['chcp', '65001'], shell=True)

    # 设置命令行参数
    parser = argparse.ArgumentParser(description="使用 slangc 编译着色器")
    parser.add_argument("--shader-list", help="着色器列表文件路径")
    parser.add_argument("--output-dir", help="输出目录")
    parser.add_argument("--slangc", default="slangc", help="slangc 编译器路径")
    parser.add_argument("--debug", action="store_true", help="启用调试模式编译")
    parser.add_argument("--opengl", action="store_true", help="编译 OpenGL 着色器")
    parser.add_argument("--vulkan", action="store_true", help="编译 Vulkan 着色器")
    parser.add_argument("--d3d11", action="store_true", help="编译 D3D11 着色器")
    parser.add_argument("--d3d12", action="store_true", help="编译 D3D12 着色器")
    parser.add_argument("--metal", action="store_true", help="编译 Metal 着色器")

    args = parser.parse_args()

    # 配置目标平台
    target_types = [
        ['glsl', args.opengl],
        ['spirv', args.vulkan],
        ['dxbc', args.d3d11],
        ['dxil', args.d3d12],
        ['metallib', args.metal],
    ]

    # 设置输出目录和着色器列表文件
    output_dir = Path(args.output_dir or "shaders")
    shader_list = Path(args.shader_list or "shader_paths.txt")

    # 读取着色器路径列表
    shader_paths = shader_list.read_text().splitlines()

    # 编译所有着色器
    all_success = True
    for shader_path in shader_paths:
        shader_path = shader_path.strip()
        for file in find_shader_files(shader_path, ['.slang']):
            if not compile_shader(file, target_types, output_dir, args):
                all_success = False

    if not all_success:
        sys.exit(1)

if __name__ == "__main__":
    main()