//
// Created by atdunbg on 2026/9/20.
//

#include <Engine/Vulkan/ShaderManager.h>
#include <Engine/Vulkan/Device.h>
#include <Engine/Core/Log.h>

#include <array>
#include <fstream>

namespace Engine {

ShaderManager::ShaderManager(Device& device)
    : m_Device(device)
{
}

ShaderManager::~ShaderManager()
{
}

ShaderObject ShaderManager::CreateFromFile(const std::string& spvPath,
                                            vk::ShaderStageFlagBits stage,
                                            const std::vector<vk::DescriptorSetLayout>& layouts,
                                            u32 pushConstantSize)
{
    auto spirv = ReadSPIRV(spvPath);
    if (spirv.empty()) {
        ENGINE_ERROR("Failed to read SPIR-V: {}", spvPath);
        return { nullptr, stage };
    }
    return CreateFromSPIRV(spirv, stage, layouts, pushConstantSize);
}

ShaderObject ShaderManager::CreateFromSPIRV(const std::vector<u32>& spirv,
                                             vk::ShaderStageFlagBits stage,
                                             const std::vector<vk::DescriptorSetLayout>& layouts,
                                             u32 pushConstantSize)
{
    vk::ShaderStageFlags nextStage = {};
    if (stage == vk::ShaderStageFlagBits::eVertex) {
        nextStage = vk::ShaderStageFlagBits::eFragment;
    }

    vk::ShaderCreateInfoEXT createInfo;
    createInfo.setStage(stage)
              .setNextStage(nextStage)
              .setCodeType(vk::ShaderCodeTypeEXT::eSpirv)
              .setCodeSize(spirv.size() * sizeof(u32))
              .setPCode(spirv.data())
              .setPName("main");

    if (!layouts.empty()) {
        createInfo.setSetLayouts(layouts);
    }

    std::array<vk::PushConstantRange, 1> pushRanges;
    if (pushConstantSize > 0) {
        // push range 是"程序级"声明：shader object 要求同一程序绑定的
        // 全部 shader（vert+frag）push range 定义完全一致（VUID-08878），
        // 因此统一按 vertex|fragment 声明；shader 实际读取哪些字节由
        // GLSL 的 push_constant block 决定，声明范围只是能力上限
        pushRanges[0].setStageFlags(vk::ShaderStageFlagBits::eVertex |
                                    vk::ShaderStageFlagBits::eFragment)
                     .setOffset(0)
                     .setSize(pushConstantSize);
        createInfo.setPushConstantRanges(pushRanges);
    }

    vk::ShaderEXT shader;
    try {
        auto result = m_Device.GetDevice().createShadersEXT(createInfo);
        shader = result.value[0];
    }
    catch (const vk::SystemError& e) {
        ENGINE_ERROR("Failed to create shader object: {}", e.what());
        return { nullptr, stage };
    }

    return { shader, stage };
}

void ShaderManager::Destroy(const ShaderObject& shader)
{
    if (shader.Handle) {
        m_Device.GetDevice().destroyShaderEXT(shader.Handle);
    }
}

std::vector<u32> ShaderManager::ReadSPIRV(const std::string& path)
{
    std::ifstream file(path, std::ios::binary | std::ios::ate);
    if (!file) {
        ENGINE_ERROR("Failed to open: {}", path);
        return {};
    }

    size_t size = file.tellg();
    std::vector<u32> buffer(size / sizeof(u32));
    file.seekg(0);
    file.read(reinterpret_cast<char*>(buffer.data()), size);
    file.close();

    return buffer;
}

} // namespace Engine