#!/usr/bin/env python3
"""
SDL3_GPU Slang Compiler - Code Generator
生成C/C++绑定代码
"""
from contextlib import contextmanager
from pathlib import Path
from typing import TextIO

from code_generator_helper import *
from global_vars import global_vars
from shader_types import *


class IndentManager:
    """RAII风格的缩进管理器"""

    def __init__(self, file: TextIO, indent_char: str = '\t'):
        self.file = file
        self.indent_char = indent_char
        self.indent_level = 0

    def write(self, text: str = '') -> None:
        """写入带缩进的文本"""
        if text:
            self.file.write(self.indent_char * self.indent_level + text + '\n')
        else:
            self.file.write('\n')
        self.file.flush()

    def write_raw(self, text: str) -> None:
        """写入原始文本（无缩进）"""
        self.file.write(text)

    @contextmanager
    def indent(self, header: Optional[str] = None, footer: Optional[str] = None):
        """缩进上下文管理器"""
        if header:
            self.write(header)

        self.indent_level += 1
        try:
            yield self
        finally:
            self.indent_level -= 1
            if footer:
                self.write(footer)

    @contextmanager
    def block(self, header: str, footer: str = '}'):
        """代码块上下文管理器"""
        self.write(header + ' {')
        self.indent_level += 1
        try:
            yield self
        finally:
            self.indent_level -= 1
            self.write(footer)

class CodeGenerator:
    """C++绑定代码生成器"""

    # 资源类型映射表
    RESOURCE_TYPE_MAP = {
        'sampled_texture': 'resource_t::SampledTexture',
        'storage_texture': 'resource_t::StorageTexture',
        'storage_buffer': 'resource_t::StorageBuffer',
        'uniform_buffer': 'resource_t::UniformBuffer',
        'sampler': 'resource_t::Sampler'
    }

    def generate_binding_functions(self, binding_infos: List[Dict], output_path: str) -> None:
        """生成C++绑定函数的入口方法"""
        self._generate_cpp_bindings(binding_infos, output_path)

    def _generate_uniform_structs(self, writer: IndentManager, uniform_buffers: List[UniformBuffer]) -> None:
        """生成Uniform缓冲区结构体"""
        writer.write("// Uniform buffer structures")

        for buffer in uniform_buffers:
            struct_name = buffer.name.replace('_buffer', '').title().replace('_', '') + 'Buffer'

            # 计算总大小和对齐要求
            total_size = 0
            max_alignment = 16  # GPU通常要求16字节对齐

            with writer.block(f'struct {struct_name}', '};'):
                for i, field in enumerate(buffer.fields):
                    # 检查是否需要填充
                    if field.offset > total_size:
                        padding_size = field.offset - total_size
                        writer.write(f"uint8_t _padding{i}[{padding_size}];  // Padding")
                        total_size = field.offset

                    # 生成字段声明
                    declaration = get_c_type_declaration(field.type, field.name)
                    writer.write(f"{declaration};  // offset: {field.offset}, size: {field.size}")

                    total_size = field.offset + field.size

            # 确保结构体大小正确对齐
            aligned_size = get_aligned_size(total_size, max_alignment)
            if aligned_size > total_size:
                writer.write(f"// Note: Structure may need padding to {aligned_size} bytes for alignment")

            writer.write(f"// Binding: {buffer.binding}, Size: {total_size} bytes (aligned: {aligned_size})")
            writer.write()


    def _generate_cpp_bindings(self, binding_infos: List[Dict], output_path: str) -> None:
        """生成C++绑定函数"""
        output_file = Path(output_path)
        # 尝试创建输出目录
        output_file.parent.mkdir(parents=True, exist_ok=True)

        with open(output_file, 'w', encoding='utf-8') as file:
            writer = IndentManager(file)
            self._write_complete_file(writer, binding_infos)

    def _write_complete_file(self, writer: IndentManager, binding_infos: List[Dict]) -> None:
        """写入完整的文件内容"""
        self._write_header(writer)

        self._write_shader_bindings_class(writer, binding_infos)
        writer.write()
        # 查找是否有像素着色器或计算着色器的处理
        if any(info['stage'] == ShaderStage.FRAGMENT.value for info in binding_infos):
            self._write_handle_class(writer, 'pixel_shader_handle_t', binding_infos)
        else:
            self._write_handle_class(writer, 'compute_shader_handle_t', binding_infos)

    def _write_header(self, writer: IndentManager) -> None:
        """写入文件头部"""
        headers = [
            '#pragma once',
            '#include <SDL3/SDL_gpu.h>',
            '#include <string>',
            '#include <vector>',
            '#include <array>',
            '#include <unordered_map>',
            '#include <cstdint>',
            '#include "shader_handle.h"',
        ]

        for header in headers:
            writer.write(header)
        writer.write()

    def _write_shader_bindings_class(self, writer: IndentManager, binding_infos: List[Dict]) -> None:
        """写入ShaderBindings类"""
        # 获取源文件名（不带路径和拓展名）

        with writer.block(f'class {global_vars.source_file_name}_shader_bindings','};'):
            self._write_public_section(writer, binding_infos)
            self._write_private_section(writer)
            self._write_constructor(writer, binding_infos)
            self._write_public_methods(writer)

    def _write_public_section(self, writer: IndentManager, binding_infos: List[Dict]) -> None:
        """写入public部分的结构体定义"""
        writer.write('public:')
        self._generate_uniform_structs(writer, global_vars.layout.uniform_buffers)

        # 写入二进制内容
        for info in binding_infos:
            entry_point = info['entry_point']
            blob_data = info['blob']

            with writer.block(f'static constexpr std::array<uint8_t, {len(blob_data)}> {entry_point}_blob =', '};'):
                # 每行16个字节，提高可读性
                for i in range(0, len(blob_data), 16):
                    chunk = blob_data[i:i + 16]
                    hex_str = ', '.join(f'0x{byte:02x}' for byte in chunk)
                    if i + 16 < len(blob_data):
                        hex_str += ','
                    writer.write(hex_str)

            writer.write()


        # Resource结构体
        with writer.block('struct resource_t', '};'):
            with writer.block('enum type_t', '};'):
                enum_values = [
                    'SampledTexture,',
                    'StorageTexture,',
                    'StorageBuffer,',
                    'UniformBuffer,',
                    'Sampler'
                ]
                for value in enum_values:
                    writer.write(value)
            writer.write()
            writer.write('std::string name;')
            writer.write('type_t type = {};')
            writer.write('int binding = -1;')
            writer.write('int set = -1;')
            writer.write('int space = -1;')
            writer.write('int index = -1;')
        writer.write()

        # ShaderInfo结构体
        with writer.block('struct shader_info_t','};'):
            writer.write('std::string entryPoint;')
            writer.write('SDL_GPUShaderStage stage = {};')
            writer.write('std::vector<resource_t> resources;')
        writer.write()

    def _write_private_section(self, writer: IndentManager) -> None:
        """写入private部分"""
        writer.write('private:')
        with writer.indent():
            writer.write('std::unordered_map<std::string, shader_info_t> shader_infos_;')
        writer.write()

    def _write_constructor(self, writer: IndentManager, binding_infos: List[Dict]) -> None:
        """写入构造函数"""
        writer.write('public:')

        with writer.indent():
            with writer.block(f'{global_vars.source_file_name}_shader_bindings()'):
                for info in binding_infos:
                    self._write_shader_initialization(writer, info)

    def _write_shader_initialization(self, writer: IndentManager, shader_info: Dict) -> None:
        """写入单个着色器的初始化代码"""
        entry_point = shader_info['entry_point']
        stage = shader_info['stage'].upper()

        writer.write(f'// {entry_point}')

        with writer.indent('{', '}'):
            writer.write('shader_info_t info;')
            writer.write(f'info.entryPoint = "{entry_point}";')
            writer.write(f'info.stage = SDL_GPU_SHADERSTAGE_{stage};')

            # 写入资源信息
            for resource in shader_info['resources']:
                self._write_resource_initialization(writer, resource)

            writer.write(f'shader_infos_["{entry_point}"] = std::move(info);')

        writer.write()

    def _write_resource_initialization(self, writer: IndentManager, resource: Dict) -> None:
        """写入单个资源的初始化代码"""
        with writer.indent('{', '}'):
            writer.write('resource_t r;')
            writer.write(f'r.name = "{resource["name"]}";')

            # 设置资源类型
            resource_type = self.RESOURCE_TYPE_MAP.get(resource["type"])
            if resource_type:
                writer.write(f'r.type = {resource_type};')

            # 设置资源绑定信息
            if 'binding' in resource:
                writer.write(f'r.binding = {resource["binding"]};')
                writer.write(f'r.set = {resource["set"]};')
            elif 'register' in resource:
                writer.write(f'r.space = {resource["space"]};')
            elif 'index' in resource:
                writer.write(f'r.index = {resource["index"]};')

            writer.write('info.resources.push_back(r);')

    def _write_public_methods(self, writer: IndentManager) -> None:
        """写入公共方法"""
        writer.write()
        self._write_bind_resources_method(writer)
        writer.write()
        self._write_get_shader_info_method(writer)

    def _write_bind_resources_method(self, writer: IndentManager) -> None:
        """写入bindResources方法"""
        method_signature = (
            'void bind_resources(\n'
            f'{writer.indent_char * (writer.indent_level + 2)}SDL_GPUCommandBuffer* in_cmd,\n'
            f'{writer.indent_char * (writer.indent_level + 2)}const std::string& in_shader_name,\n'
            f'{writer.indent_char * (writer.indent_level + 2)}const std::unordered_map<std::string, void*>& in_resources)'
        )

        with writer.block(method_signature):
            writer.write('auto it = shader_infos_.find(in_shader_name);')

            with writer.block('if (it == shader_infos_.end())'):
                writer.write('return;')

            writer.write()
            writer.write('const auto& info = it->second;')
            writer.write()
            writer.write('// 根据着色器阶段和资源类型进行绑定')
            writer.write('// 这里需要根据SDL3 GPU的实际API来实现')

            with writer.block('for (const auto& res : info.resources)'):
                writer.write('auto resIt = in_resources.find(res.name);')

                with writer.block('if (resIt == in_resources.end())'):
                    writer.write('continue;')

                writer.write()
                writer.write('void* resource = resIt->second;')
                writer.write()
                writer.write('// 示例绑定逻辑')

                with writer.block('switch (res.type)'):
                    writer.write('case resource_t::SampledTexture:')
                    with writer.indent():
                        writer.write('// SDL_BindGPUVertexSamplers(in_cmd, ...);')
                        writer.write('break;')

                    writer.write('case resource_t::StorageBuffer:')
                    with writer.indent():
                        writer.write('// SDL_BindGPUVertexStorageBuffers(in_cmd, ...);')
                        writer.write('break;')

                    writer.write('// ... 其他资源类型')
                    writer.write('default:')
                    with writer.indent():
                        writer.write('break;')

    def _write_get_shader_info_method(self, writer: IndentManager) -> None:
        """写入getShaderInfo方法"""
        with writer.block('[[nodiscard]] const auto* get_shader_info(const std::string& in_shader_name) const'):
            writer.write('auto it = shader_infos_.find(in_shader_name);')
            writer.write('return it != shader_infos_.end() ? &it->second : nullptr;')

    def _write_handle_class(self, writer: IndentManager, parent_class: str, shader_info: List[Dict]) -> None:
        """写入ShaderHandle类"""
        # class test_shader_handle_t: public pixel_shader_handle_t
        with writer.block(f'class {global_vars.source_file_name}_shader_handle_t : public {parent_class}', '};'):
            writer.write('protected:')
            for info in shader_info:
                with writer.indent(f'virtual SDL_GPUShader* create_{info['stage'].lower()}_shader(SDL_GPUDevice* in_gpu_device) override ''{', '}'):
                    writer.write('SDL_GPUShaderCreateInfo info{};')
                    writer.write(f'info.code = {global_vars.source_file_name}_shader_bindings::{info["entry_point"]}_blob.data();')
                    writer.write(f'info.code_size = {global_vars.source_file_name}_shader_bindings::{info["entry_point"]}_blob.size();')
                    writer.write(f'info.entrypoint = "{info["entry_point"]}";')
                    writer.write(f'info.format = SDL_GPU_SHADERFORMAT_{global_vars.target.value.upper()};')
                    writer.write(f'info.stage = SDL_GPU_SHADERSTAGE_{info['stage'].upper()};')
                    writer.write(f'info.num_samplers = {info['num_resources'][ResourceType.SAMPLER]};')
                    writer.write(f'info.num_storage_textures = {info['num_resources'][ResourceType.STORAGE_TEXTURE]};')
                    writer.write(f'info.num_storage_buffers = {info['num_resources'][ResourceType.STORAGE_BUFFER]};')
                    writer.write(f'info.num_uniform_buffers = {info['num_resources'][ResourceType.UNIFORM_BUFFER]};')
                    writer.write('return SDL_CreateGPUShader(in_gpu_device, &info);')
                writer.write()