#!/usr/bin/env python3
import io
import json
import sys

from code_generator import IndentManager
from code_generator_helper import *
from global_vars import global_vars
from shader_types import *

# 数据解析器（保持不变）
class ShaderLayoutParser:
    """解析JSON数据并提取到类对象"""

    def parse(self, json_file: str) -> ShaderLayout:
        """解析JSON文件并返回ShaderLayout对象"""
        with open(json_file, 'r') as f:
            json_data = json.load(f)

        layout = ShaderLayout()
        layout.vertex_fields = self._parse_vertex_input(json_data)
        layout.uniform_buffers = self._parse_uniform_buffers(json_data)

        return layout

    def _parse_vertex_input(self, json_data: Dict) -> List[VertexField]:
        """解析顶点输入字段"""
        vertex_fields = []

        entry_points = json_data.get('entryPoints', [])
        for entry in entry_points:
            if entry.get('stage') == 'vertex':
                parameters = entry.get('parameters', [])
                for param in parameters:
                    if param.get('name') == 'input' and param.get('stage') == 'vertex':
                        fields = param.get('type', {}).get('fields', [])
                        for field in fields:
                            vertex_field = VertexField(
                                name=field['name'],
                                type=FieldType.from_dict(field['type']),
                                location=field['binding']['index'],
                                semantic=field.get('semanticName', ''),
                                semantic_index=field.get('semanticIndex', 0)
                            )
                            vertex_fields.append(vertex_field)

        return vertex_fields

    def _parse_uniform_buffers(self, json_data: Dict) -> List[UniformBuffer]:
        """解析Uniform缓冲区"""
        uniform_buffers = []

        parameters = json_data.get('parameters', [])
        for param in parameters:
            binding = param.get('binding', {})
            if binding.get('kind') == 'constantBuffer':
                buffer = UniformBuffer(
                    name=param['name'],
                    binding=binding['index']
                )

                # 解析缓冲区字段
                element_type = param.get('type', {}).get('elementType', {})
                if element_type.get('kind') == 'struct':
                    fields = element_type.get('fields', [])
                    for field in fields:
                        uniform_field = UniformField(
                            name=field['name'],
                            type=FieldType.from_dict(field['type']),
                            offset=field['binding']['offset'],
                            size=field['binding']['size']
                        )
                        buffer.fields.append(uniform_field)

                uniform_buffers.append(buffer)

        return uniform_buffers

# **重构后的代码生成器**
class ShaderLayoutCodeGenerator:
    """从ShaderLayout对象生成C代码"""

    def generate(self) -> str:
        """生成完整的头文件内容"""
        output = io.StringIO()
        writer = IndentManager(output, indent_char='    ')  # 使用4个空格缩进

        writer.write("#pragma once")
        writer.write()
        writer.write("#include <SDL3/SDL.h>")
        writer.write("#include <SDL3/SDL_gpu.h>")
        writer.write("#include <cstdint>  // For fixed-width integer types")
        writer.write()

        writer.write("// Auto-generated vertex structure")
        with writer.block(f'namespace {global_vars.source_file_name}_shader'):
            self._generate_vertex_struct(writer, global_vars.layout.vertex_fields)
            self._generate_vertex_attributes(writer, global_vars.layout.vertex_fields)
            self._generate_helper_functions(writer, global_vars.layout.vertex_fields)

        return output.getvalue()

    def _generate_vertex_struct(self, writer: IndentManager, vertex_fields: List[VertexField]) -> None:
        """生成顶点结构体"""
        with writer.block("struct vertex_t", "};"):
            total_offset = 0
            for field in vertex_fields:
                # 生成字段声明
                declaration = get_c_type_declaration(field.type, field.name)

                # 获取字段信息
                c_type, count, size = get_type_info(field.type)

                # 添加详细注释
                semantic = field.semantic
                if field.semantic_index > 0:
                    semantic += str(field.semantic_index)

                comment = f"// location: {field.location}, semantic: {semantic}"
                comment += f", offset: {total_offset}, size: {size} bytes"

                writer.write(comment)
                writer.write(f"{declaration};")

                total_offset += size

        # 添加结构体大小的静态断言
        writer.write()
        writer.write(f"// Total size: {total_offset} bytes")
        writer.write(f"static_assert(sizeof(vertex_t) == {total_offset}, \"Vertex struct size mismatch\");")
        writer.write()

    def _generate_vertex_attributes(self, writer: IndentManager, vertex_fields: List[VertexField]) -> None:
        """生成顶点属性数组"""
        writer.write("// Vertex attribute descriptions")
        writer.write(f"static constexpr uint32_t VERTEX_ATTRIBUTE_COUNT = {len(vertex_fields)};")
        writer.write()

        writer.write("static const SDL_GPUVertexAttribute vertex_attributes[] = {")

        offset = 0
        with writer.indent():
            for i, field in enumerate(vertex_fields):
                c_type, count, size = get_type_info(field.type)
                scalar_type = field.type.scalar_type or 'float32'

                format_key = (scalar_type, count)
                format_name = format_mapping.get(format_key, 'SDL_GPU_VERTEXELEMENTFORMAT_FLOAT4')

                with writer.indent("{", "}," if i < len(vertex_fields) - 1 else "}"):
                    writer.write(f"// {field.name}")
                    writer.write(f".location = {field.location},")
                    writer.write(f".buffer_slot = 0,")
                    writer.write(f".format = {format_name},")
                    writer.write(f".offset = {offset}")

                offset += size

        writer.write("};")
        writer.write()

        # 生成顶点缓冲区描述
        writer.write("// Vertex buffer description")
        with writer.block("static constexpr SDL_GPUVertexBufferDescription vertex_buffer_desc =", "};"):
            writer.write(".slot = 0,")
            writer.write(f".pitch = {offset},  // sizeof(vertex)")
            writer.write(".input_rate = SDL_GPU_VERTEXINPUTRATE_VERTEX,")
            writer.write(".instance_step_rate = 0")
        writer.write()

        # 生成顶点输入状态
        writer.write("// Vertex input state")
        with writer.block("static constexpr SDL_GPUVertexInputState vertex_input_state =", "};"):
            writer.write(".vertex_buffer_descriptions = &vertex_buffer_desc,")
            writer.write(".num_vertex_buffers = 1,")
            writer.write(".vertex_attributes = vertex_attributes,")
            writer.write(".num_vertex_attributes = VERTEX_ATTRIBUTE_COUNT")
        writer.write()

    def _generate_helper_functions(self, writer: IndentManager, vertex_fields: List[VertexField]) -> None:
        """生成辅助函数"""
        writer.write("// Helper functions")
        writer.write()

        # 创建顶点缓冲区函数
        writer.write("static SDL_GPUBuffer* create_vertex_buffer(SDL_GPUDevice* device,")
        writer.write("                                       const vertex_t* vertices,")

        with writer.block("                                       Uint32 vertex_count)", "}"):
            with writer.block("SDL_GPUBufferCreateInfo buffer_info =", "};"):
                writer.write(".usage = SDL_GPU_BUFFERUSAGE_VERTEX,")
                writer.write(".size = static_cast<Uint32>(sizeof(vertex_t)) * vertex_count")

            writer.write()
            writer.write("SDL_GPUBuffer* buffer = SDL_CreateGPUBuffer(device, &buffer_info);")
            writer.write()

            with writer.block("SDL_GPUTransferBufferCreateInfo transfer_info =", "};"):
                writer.write(".usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD,")
                writer.write(".size = buffer_info.size")

            writer.write()
            writer.write("// Upload vertex data")
            writer.write("SDL_GPUTransferBuffer* transfer = SDL_CreateGPUTransferBuffer(device, &transfer_info);")
            writer.write()
            writer.write("void* mapped = SDL_MapGPUTransferBuffer(device, transfer, false);")
            writer.write("SDL_memcpy(mapped, vertices, buffer_info.size);")
            writer.write("SDL_UnmapGPUTransferBuffer(device, transfer);")
            writer.write()
            writer.write("// Copy to GPU")
            writer.write("SDL_GPUCommandBuffer* cmd = SDL_AcquireGPUCommandBuffer(device);")
            writer.write("SDL_GPUCopyPass* copy = SDL_BeginGPUCopyPass(cmd);")
            writer.write()
            writer.write("SDL_GPUTransferBufferLocation src = {.transfer_buffer = transfer, .offset = 0};")
            writer.write("SDL_GPUBufferRegion dst = {.buffer = buffer, .offset = 0, .size = buffer_info.size};")
            writer.write()
            writer.write("SDL_UploadToGPUBuffer(copy, &src, &dst, false);")
            writer.write("SDL_EndGPUCopyPass(copy);")
            writer.write("SDL_SubmitGPUCommandBuffer(cmd);")
            writer.write()
            writer.write("SDL_ReleaseGPUTransferBuffer(device, transfer);")
            writer.write("return buffer;")

        writer.write()


# 主生成器类（保持不变）
class ShaderLayoutGenerator:
    """主生成器类，整合解析和代码生成功能"""

    def __init__(self):
        self.parser = ShaderLayoutParser()
        self.code_generator = ShaderLayoutCodeGenerator()

    def generate_header(self, json_file: str) -> str:
        """生成头文件内容"""
        # 解析JSON到类对象
        global_vars.layout = self.parser.parse(json_file)

        # 基于类对象生成代码
        return self.code_generator.generate()


# 使用示例
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: shader_layout_gen.py <shader.json>")
        sys.exit(1)

    generator = ShaderLayoutGenerator()
    header_content = generator.generate_header(sys.argv[1])
    print(header_content)