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

from typing import List, Tuple, Iterator
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编码打印消息"""
    print(message)

# 着色器类型和扩展名定义
SHADER_TYPES = {
    'vertex': 'LLGL::ShaderType::Vertex',
    'pixel': 'LLGL::ShaderType::Fragment',
    'fragment': 'LLGL::ShaderType::Fragment',
    'compute': 'LLGL::ShaderType::Compute',
    'geometry': 'LLGL::ShaderType::Geometry',
    'tess_control': 'LLGL::ShaderType::TessControl',
    'tess_evaluation': 'LLGL::ShaderType::TessEvaluation'
}

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_0'],
    'dxil': ['-profile', 'sm_6_6'],
    'metallib': ['-capability', 'metallib_3_1']
}

class ShaderEntry:
    """着色器入口点信息"""
    def __init__(self, name: str, shader_type: str):
        self.name = name
        self.shader_type = shader_type

class CompiledShaderInfo:
    """编译后的着色器信息"""
    def __init__(self, output_file: str, base: str, entry: ShaderEntry):
        self.output_file = output_file
        self.base = base
        self.entry = entry
        if entry.shader_type == 'pixel':
            self.entry.shader_type = 'fragment'

# 存放所有编译成功的着色器信息
compiled_shaders: List[CompiledShaderInfo] = []

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 find_slang_entries(input_file: Path) -> List[ShaderEntry]:
    """从着色器文件中提取入口点函数名和类型"""
    # 匹配 [shader("xxx")] 形式的着色器类型声明，以及后面的函数名
    pattern = re.compile(
r'\[\s*shader\s*\(\s*"([^"]+)"\s*\)\s*\]\s*'  # 匹配 [shader("xxx")]
r'(?:\[\s*[^\]]+\])*'                       # 可选：匹配其他属性如 [numthreads(8,8,1)]
r'\s*\w+\s+(\w+)\s*\('                      # 匹配函数声明：返回类型 函数名(
    )
    try:
        content = input_file.read_text(encoding='utf-8')
        matches = pattern.findall(content)
        print_utf8(f"**调试**: 在文件 {input_file} 中找到的匹配: {matches}")

        entries = []
        for shader_type, name in matches:
            if shader_type in SHADER_TYPES:
                entries.append(ShaderEntry(name, shader_type))
            else:
                print_utf8(f"**警告**: 未知的着色器类型 {shader_type}")
        return entries
    except Exception as e:
        print_utf8(f"**错误**: 解析文件 {input_file} 失败: {e}")
        return []


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

def create_compiler_command(
        input_file: Path,
        entry: ShaderEntry,
        output_file: Path,
        target_type: str,
        args: argparse.Namespace
) -> List[str]:
    """生成着色器编译命令"""
    cmd = [args.slangc,
           str(input_file),
           "-entry", entry.name,
           "-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:
    """检查是否需要重新编译着色器"""
    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:
    """编译单个着色器文件"""
    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 entry in entries:
            compiled_shaders.append(CompiledShaderInfo(f"{base}_{entry.name}", base, entry))

        for target_type, enabled in target_types:
            if not enabled:
                continue

            for entry in entries:
                output_file = output_dir / f"{base}_{entry.name}.{get_shader_extension(target_type)}"

                if not needs_recompile(input_file, output_file):
                    print_utf8(f"**跳过**: {output_file} - 文件已是最新")
                else:
                    cmd = create_compiler_command(input_file, entry, output_file, target_type, args)
                    try:
                        subprocess.run(cmd, check=True, capture_output=True, text=True)
                        print_utf8(f"**成功**: 编译 {input_file}:{entry.name} -> {output_file}")
                    except subprocess.CalledProcessError as e:
                        print_utf8(f"**错误**: 编译 {input_file}:{entry.name} 失败")
                        print_utf8(e.stderr)
                        success = False
                        continue

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

def generate_pipeline_header_preamble() -> List[str]:
    """Generate the header file preamble"""

    return [
        "#pragma once",
        "",
        "#include \"mirage.h\"",
        "#include \"misc/mirage_type.h\"",
        "#include <LLGL/LLGL.h>",
        "#include <LLGL/Utils/VertexFormat.h>",
        "#include <memory>",
        "#include <string>",
        "#include <vector>",
        "#include <fstream>",
        "#include <stdexcept>",
        "",
        "namespace generated_pipelines {",
        "",
        "// 辅助函数：加载着色器",
        "inline auto LoadShader(",
        "    LLGL::RenderSystem* renderer,",
        "    const std::string& filename,",
        "    const LLGL::ShaderType type,",
        "    const char* entryPoint,",
        "    const LLGL::ShaderDescriptor& shaderDesc = {}) {",
        "",
        "   // 根据渲染器类型选择着色器文件后缀",
        "   auto rendererID = renderer->GetRendererID();",
        "   std::string ext;",
        "   // 选择对应的文件扩展名",
        "   if (rendererID == LLGL::RendererID::OpenGL) {",
        "       ext = \".glsl\";",
        "   } else if (rendererID == LLGL::RendererID::Vulkan) {",
        "       ext = \".spirv\";",
        "   } else if (rendererID == LLGL::RendererID::Direct3D11) {",
        "       ext = \".dxbc\";",
        "   } else if (rendererID == LLGL::RendererID::Direct3D12) {",
        "       ext = \".dxil\";",
        "   } else if (rendererID == LLGL::RendererID::Metal) {",
        "       ext = \".metallib\";",
        "   } else {",
        "       ext = \".dat\";",
        "   }",
        "",
        "    // 构造最终的文件名",
        "    std::string finalFilename = filename;",
        "    size_t pos = finalFilename.find_last_of('.');",
        "    if (pos != std::string::npos) {",
        "        finalFilename = finalFilename.substr(0, pos) + ext;",
        "    } else {",
        "        finalFilename += ext;",
        "    }",
        "    finalFilename = (mirage::get_shader_path() / finalFilename).string();",
        "",
        "    // 读取着色器文件",
        "    std::vector<char> shaderData;",
        "    try {",
        "        std::ifstream file(finalFilename, std::ios::binary | std::ios::ate);",
        "        if (!file.is_open()) {",
        "            throw std::runtime_error(\"Failed to open shader file: \" + finalFilename);",
        "        }",
        "        size_t fileSize = static_cast<size_t>(file.tellg());",
        "        shaderData.resize(fileSize);",
        "        file.seekg(0);",
        "        file.read(shaderData.data(), fileSize);",
        "    } catch (const std::exception& e) {",
        "        throw std::runtime_error(\"Failed to read shader file: \" + std::string(e.what()));",
        "    }",
        "",
        "    // 创建着色器",
        "    LLGL::ShaderDescriptor desc = shaderDesc;",
        "    desc.source = shaderData.data();",
        "    desc.sourceSize = shaderData.size();",
        "    desc.entryPoint = entryPoint;",
        "    desc.type = type;",
        "    desc.sourceType = rendererID == LLGL::RendererID::OpenGL ? LLGL::ShaderSourceType::CodeString : LLGL::ShaderSourceType::BinaryBuffer;",
        "",
        "    auto shader = renderer->CreateShader(desc);",
        "    if (auto report = shader->GetReport()) {",
        "        spdlog::error(\"Shader compilation report: {}\", report->GetText());",
        "    }",
        "    return mirage::shader_ptr(shader, mirage::llgl_deleter<LLGL::Shader>);",
        "}",
        ""
        "inline auto create_pipeline_layout(LLGL::RenderSystem* renderer, const LLGL::PipelineLayoutDescriptor& in_desc) {",
        "    auto pipelineLayout = renderer->CreatePipelineLayout(in_desc);",
        "    return mirage::pipeline_layout_ptr(pipelineLayout, mirage::llgl_deleter<LLGL::PipelineLayout>);",
        "}",
    ]

def generate_compute_pipeline(header_lines: List[str], shader: CompiledShaderInfo):
    """Generate compute pipeline creation function"""
    func_name = f"create_{shader.base}_{shader.entry.name}_pipeline"

    header_lines.extend([
        f"// 计算管线: {shader.base} - {shader.entry.name}",
        f"inline auto {func_name}(LLGL::RenderSystem* renderer) {{",
        "    // 加载计算着色器",
        f"   auto computeShader = LoadShader(renderer, \"{shader.output_file.name}\", {SHADER_TYPES['compute']}, \"{shader.entry.name}\");",
        "",
        "    // 创建管线布局",
        "    LLGL::PipelineLayoutDescriptor layoutDesc;",
        "    auto pipelineLayout = create_pipeline_layout(renderer, layoutDesc);",
        "",
        "    // 创建计算管线",
        "    LLGL::ComputePipelineDescriptor pipelineDesc;",
        "    pipelineDesc.computeShader = computeShader.get();",
        "    pipelineDesc.pipelineLayout = pipelineLayout.get();",
        "",
        "    auto pipeline = renderer->CreatePipelineState(pipelineDesc);",
        "    mirage::pipeline_info info;",
        "    info.pipeline_state = mirage::pipeline_state_ptr(pipeline, mirage::llgl_deleter<LLGL::PipelineState>);",
        "    info.pipeline_layout = pipelineLayout;",
        "    info.shaders = {computeShader};",
        "    return info;",
        "}",
        ""
    ])

def generate_graphics_pipeline(header_lines: List[str], base: str, shaders: List[CompiledShaderInfo]):
    """Generate graphics pipeline creation function"""
    func_name = f"create_{base}_pipeline"

    header_lines.extend([
        f"// 图形管线: {base}",
        f"inline auto {func_name}(LLGL::RenderSystem* renderer, const LLGL::RenderPass* render_pass, const LLGL::PipelineLayoutDescriptor& in_pipeline_layout_desc, const LLGL::VertexFormat& vertex_format = mirage::create_vertex_format()) {{",
        f"    // 加载各个阶段的着色器",
        f"    LLGL::ShaderDescriptor vertexShaderDesc, fragShaderDesc;",
        f"    vertexShaderDesc.vertex.inputAttribs = vertex_format.attributes;",
    ])

    SHADER_DESC_NAME = {
        'vertex': 'vertexShaderDesc',
        'pixel': 'fragShaderDesc',
        'fragment': 'fragShaderDesc',
    }

    # Load all shader stages
    for shader in shaders:
        shader_type = shader.entry.shader_type
        if shader_type in SHADER_TYPES:
            header_lines.append(
                f"    auto {shader_type}Shader = LoadShader(renderer, "
                f"\"{shader.output_file}\", {SHADER_TYPES[shader_type]}, "
                f"\"{shader.entry.name}\", {SHADER_DESC_NAME[shader_type]});"
            )

    # Create pipeline layout and descriptor
    header_lines.extend([
        "    // 创建管线布局",
        "    auto pipelineLayout = create_pipeline_layout(renderer, in_pipeline_layout_desc);",
        "",
        "    // 创建图形管线",
        "    LLGL::GraphicsPipelineDescriptor pipelineDesc;",
    ])

    # Set all shader stages
    for shader in shaders:
        shader_type = shader.entry.shader_type
        if shader_type in SHADER_TYPES:
            header_lines.append(f"    pipelineDesc.{shader_type}Shader = {shader_type}Shader.get();")

    # Set basic render states
    header_lines.extend([
        "    pipelineDesc.pipelineLayout = pipelineLayout.get();",
        "",
        "    // 设置基本渲染状态",
        "    pipelineDesc.renderPass = render_pass;",
        "    pipelineDesc.rasterizer.multiSampleEnabled = true;",
        "    pipelineDesc.blend.targets[0].blendEnabled = true;",
        "    pipelineDesc.depth.testEnabled = true;",
        "    pipelineDesc.depth.writeEnabled = true;",
        "",
        "    auto pipeline = renderer->CreatePipelineState(pipelineDesc);",
    ])

    # Finish function
    header_lines.extend([
        "    mirage::pipeline_info info;",
        "    info.pipeline_state = mirage::pipeline_state_ptr(pipeline, mirage::llgl_deleter<LLGL::PipelineState>);",
        "    info.pipeline_layout = pipelineLayout;",
        "    info.shaders = {",
    ])

    for shader in shaders:
        shader_type = shader.entry.shader_type
        if shader_type in SHADER_TYPES:
            header_lines.append(f"        {shader_type}Shader,")

    header_lines.append("    };")

    header_lines.extend([
        "    return info;",
        "}",
        ""
    ])

def generate_pipeline_header(header_path: Path):
    """Generate the complete pipeline header file"""
    header_lines = generate_pipeline_header_preamble()

    # Group shaders by base name
    shader_groups = {}
    for shader in compiled_shaders:
        if shader.base not in shader_groups:
            shader_groups[shader.base] = []
        shader_groups[shader.base].append(shader)

    # Generate pipeline functions
    for base, shaders in shader_groups.items():
        has_compute = any(s.entry.shader_type == "compute" for s in shaders)
        if has_compute:
            for shader in shaders:
                if shader.entry.shader_type == "compute":
                    generate_compute_pipeline(header_lines, shader)
        else:
            generate_graphics_pipeline(header_lines, base, shaders)

    # Close namespace
    header_lines.extend([
        "} // namespace generated_pipelines",
        ""
    ])

    try:
        header_path.write_text("\n".join(header_lines), encoding="utf-8")
        print_utf8(f"**成功**: 生成管线 C++ 头文件 {header_path}")
    except Exception as e:
        print_utf8(f"**错误**: 写入头文件 {header_path} 失败: {e}")

def main():
    """主函数：解析命令行参数并执行编译流程"""
    if sys.platform.startswith('win'):
        subprocess.run(['chcp', '65001'], shell=True)

    parser = argparse.ArgumentParser(description="使用 slangc 编译着色器并生成 LLGL 渲染管线 C++ 头文件")
    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 着色器")
    parser.add_argument("--header", help="生成的头文件路径")

    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")

    try:
        shader_paths = shader_list.read_text(encoding="utf-8").splitlines()
    except Exception as e:
        print_utf8(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, output_dir, args):
                all_success = False

    # 输出到shader_list所在目录
    header_file = Path(args.header or shader_list.parent / "generated_pipelines.h")
    generate_pipeline_header(header_file)

    if not all_success:
        sys.exit(1)

if __name__ == "__main__":
    main()