#!/usr/bin/env python3
# -*- coding: utf-8 -*-

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

# 支持的后端列表
BACKENDS = [
    "glsl", # SG_BACKEND_GLCORE
    "essl", # SG_BACKEND_GLES3
    "hlsl", # SG_BACKEND_D3D11
    "metal", # SG_BACKEND_METAL_IOS, SG_BACKEND_METAL_MACOS, SG_BACKEND_METAL_SIMULATOR
    "wgsl", # SG_BACKEND_WGPU
]

SHADER_EXTENSIONS = {
    'glsl': 'glsl',
    'essl': 'glsl',
    'hlsl': 'dxil',
    'metal': 'metallib',
    'wgsl': 'wgsl'
}

TARGET_PROFILES = {
    # 'glsl': ['-profile', 'glsl_460'],
    # 'essl': ['-profile', 'glsl_300es'],
    # 'hlsl': ['-profile', 'sm_5_0'],
    # 'metal': ['-capability', 'metallib'],
    # 'wgsl': ['-profile', 'wgsl']
}

def need_recompile(input_file: Path, output_path: Path) -> bool:
    """检查是否需要重新编译"""
    # 着色器输出文件名的格式为input.file.stem + '.' + backend + '.h'
    # 所以需要检查是否有input.file.stem + '.*.h'文件存在
    for backend in BACKENDS:
        output_file = output_path / f"{input_file.stem}.{backend}.h"
        if not output_file.exists():
            continue
        else:
            input_time = input_file.stat().st_mtime
            output_time = output_file.stat().st_mtime
            return input_time > output_time
    return True


def find_shader_files(input_dir: Path, extensions: List[str]) -> Iterator[Path]:
    """递归查找指定目录下的着色器文件"""
    for file_path in Path(input_dir).rglob('*'):
        if file_path.suffix in extensions:
            yield file_path

def create_compiler_command(
        input_file: Path,
        output_file: Path,
        target_type: str,
        args: argparse.Namespace
) -> List[str]:
    """生成着色器编译命令"""
    cmd = [args.shdc,
           str(input_file),
           "-o", str(output_file),
           "-t", target_type,
        ]

    if args.debug:
        cmd.append('-d')

    if target_type in TARGET_PROFILES:
        cmd.extend(TARGET_PROFILES[target_type])

    return cmd

def compile_shader(
        input_file: Path,
        target_types: List[Tuple[str, bool]],
        args: argparse.Namespace
) -> bool:
    """编译单个着色器文件并读取二进制数据"""
    try:
        base = input_file.stem
        output_dir = input_file.parent
        success = True

        for target_type, enabled in target_types:
            if not enabled:
                continue
            if not need_recompile(input_file, output_dir):
                print(f"**跳过**: {input_file} 已经是最新的")
                continue

            cmd = create_compiler_command(input_file, output_dir, target_type, args)
            try:
                # 修改这里: 明确指定编码为utf-8，并添加errors参数处理无法解码的字符
                subprocess.run(cmd, check=True, capture_output=True, text=True, encoding='utf-8', errors='replace')
                print(f"**成功**: 编译 {input_file}")
            except subprocess.CalledProcessError as e:
                print(f"**错误**: 编译 {input_file} 失败")
                print(e.stderr)
                success = False
                continue

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

def main():
    parser = argparse.ArgumentParser(description='编译slang着色器为C++头文件')
    parser.add_argument('--shdc', help='着色器编译器路径')
    parser.add_argument('--shader_list', help='着色器目录列表文件路径')
    parser.add_argument("--hlsl", action="store_true", help="编译HLSL着色器")
    parser.add_argument("--glsl", action="store_true", help="编译GLSL着色器")
    parser.add_argument("--essl", action="store_true", help="编译ESSL着色器")
    parser.add_argument("--metal", action="store_true", help="编译Metal着色器")
    parser.add_argument("--wgsl", action="store_true", help="编译WGSL着色器")
    parser.add_argument("--debug", action="store_true", help="编译调试版本")
    args = parser.parse_args()

    # 确定要编译的目标后端
    target_types: List[Tuple[str, bool]] = [
        ('glsl', args.glsl),
        ('essl', args.essl),
        ('hlsl', args.hlsl),
        ('metal', args.metal),
        ('wgsl', args.wgsl),
    ]

    # 如果没有指定任何后端，默认启用所有后端
    if not any(enabled for _, enabled in target_types):
        target_types = [(backend, True) for backend, _ in target_types]

    shader_list = Path(args.shader_list or "shader_paths.txt")

    try:
        shader_paths = shader_list.read_text(encoding="utf-8").splitlines()
    except Exception as e:
        print(f"**错误**: 读取着色器列表文件 {shader_list} 失败: {e}")
        sys.exit(1)

    all_success = True
    for shader_path in shader_paths:
        shader_path = shader_path.strip()
        if not shader_path:
            continue
        for file in find_shader_files(Path(shader_path), ['.slang']):
            if not compile_shader(file, target_types, args):
                all_success = False

    if not all_success:
        sys.exit(1)

if __name__ == "__main__":
    main()