#!/usr/bin/env python3
"""
SDL3_GPU Slang Compiler - Main Compiler
主编译器类，整合所有功能模块
"""

import os
import subprocess
import tempfile
from typing import List, Dict

from binding_manager import BindingManager
from code_generator import CodeGenerator
from compiler_cmd import make_cmd
from exe_finder import dxc_path
from global_vars import global_vars
from shader_parser import ShaderParser
from shader_types import ShaderInfo, TargetFormat, ShaderStage


class SDL3GPUSlangCompiler:
    def __init__(self):
        self.parser = ShaderParser()
        self.binding_manager = BindingManager()
        self.code_generator = CodeGenerator()

    def parse_slang_shader(self) -> Dict[str, ShaderInfo]:
        """解析Slang着色器源码，提取资源信息"""
        return self.parser.parse_slang_shader()

    def compile_shader(self, shader_info: ShaderInfo, target: TargetFormat) -> tuple[str, dict]:
        """编译着色器并返回二进制路径和绑定信息"""
        output_path = tempfile.mktemp()
        # 根据目标格式分配绑定点
        if target == TargetFormat.SPIRV:
            self.binding_manager.assign_bindings_spirv(shader_info)
        elif target in [TargetFormat.DXIL, TargetFormat.DXBC]:
            self.binding_manager.assign_bindings_dxil(shader_info)
        elif target == TargetFormat.MSL:
            self.binding_manager.assign_bindings_msl(shader_info)

        # 生成带绑定信息的着色器代码
        modified_source = self.binding_manager.inject_bindings(shader_info, target)

        temp_source = os.path.join(global_vars.source_path, global_vars.source_file_name + '.slang')

        # 写入临时文件
        with open(temp_source, 'w', encoding='utf8') as f:
            f.write(modified_source)

        with tempfile.NamedTemporaryFile(mode='w', suffix='.cso', delete=False, encoding='utf8') as tmp:
            signed_shader = tmp.name

        try:
            # 编译着色器
            if target == TargetFormat.DXIL:
                cmd = make_cmd(temp_source, TargetFormat.HLSL_DX12, shader_info.stage, shader_info.entry_point, output_path)
            else:
                cmd = make_cmd(temp_source, target, shader_info.stage, shader_info.entry_point, output_path)

            print(f"Compiling shader with command: {' '.join(cmd)}")
            shader_file = output_path

            subprocess.run(cmd, check=True)
            print(f"Shader compiled successfully")

            # 如果目标是DXIL，则需要签名
            if global_vars.target == TargetFormat.DXIL:
                self._compile_dxil(shader_info, shader_file, signed_shader)
                shader_file = signed_shader

            # 生成绑定信息
            binding_info = self.binding_manager.generate_binding_info(shader_info, target, shader_file)

            return binding_info

        finally:
            # 清理临时文件
            os.unlink(signed_shader)
            os.unlink(output_path)

    def generate_binding_functions(self, binding_infos: List[Dict], output_path: str):
        """生成C/C++绑定函数"""
        self.code_generator.generate_binding_functions(binding_infos, output_path)

    def _compile_dxil(self, shader_info: ShaderInfo, input_file: str, output_file: str):
        """编译DXIL着色器"""
        # 根据stage选择不同的-T参数
        if shader_info.stage == ShaderStage.VERTEX:
            t = 'vs_6_6'
        elif shader_info.stage == ShaderStage.FRAGMENT:
            t = 'ps_6_6'
        elif shader_info.stage == ShaderStage.COMPUTE:
            t = 'cs_6_6'
        else:
            raise ValueError(f"Unsupported shader stage: {shader_info.stage}")

        dxc_cmd = [
            dxc_path,
            '-T', t,
            '-E', shader_info.entry_point,
            '-Fo', output_file,
            input_file
        ]

        result = subprocess.run(dxc_cmd, check=True)
        if result.stdout:
            print(f"Stdout: {result.stdout}")
        if result.stderr is None:
            print(f"Signed shader compiled successfully")
        else:
            raise RuntimeError(f"Failed to sign DXIL shader: {result.stderr.decode('utf-8')}")