#include "XCEngine/RHI/OpenGL/OpenGLDescriptorSet.h"
#include "XCEngine/RHI/OpenGL/OpenGLPipelineLayout.h"
#include "XCEngine/RHI/OpenGL/OpenGLTextureUnitAllocator.h"
#include "XCEngine/RHI/OpenGL/OpenGLResourceView.h"
#include "XCEngine/RHI/OpenGL/OpenGLEnums.h"
#include "XCEngine/RHI/OpenGL/OpenGLTexture.h"
#include "XCEngine/RHI/OpenGL/OpenGLSampler.h"
#include <glad/glad.h>

namespace XCEngine {
namespace RHI {

namespace {

uint32_t ResolveBindingPoint(
    const OpenGLPipelineLayout* pipelineLayout,
    uint32_t setIndex,
    const DescriptorBinding& binding) {
    if (pipelineLayout == nullptr || !pipelineLayout->UsesSetLayouts()) {
        return binding.binding;
    }

    switch (binding.type) {
        case DescriptorType::CBV:
            return pipelineLayout->GetConstantBufferBindingPoint(setIndex, binding.binding);
        case DescriptorType::SRV:
            return pipelineLayout->GetShaderResourceBindingPoint(setIndex, binding.binding);
        case DescriptorType::UAV:
            return pipelineLayout->GetUnorderedAccessBindingPoint(setIndex, binding.binding);
        case DescriptorType::Sampler:
            return pipelineLayout->GetSamplerBindingPoint(setIndex, binding.binding);
        default:
            return UINT32_MAX;
    }
}

bool UsesBufferBinding(ResourceViewDimension dimension) {
    return IsBufferResourceViewDimension(dimension);
}

} // namespace

OpenGLDescriptorSet::OpenGLDescriptorSet()
    : m_allocator(nullptr)
    , m_layoutBindings(nullptr)
    , m_bindingCount(0)
    , m_bound(false) {
}

OpenGLDescriptorSet::~OpenGLDescriptorSet() {
    Shutdown();
}

DescriptorBinding* OpenGLDescriptorSet::FindBinding(uint32_t binding) {
    for (auto& descriptorBinding : m_bindings) {
        if (descriptorBinding.binding == binding) {
            return &descriptorBinding;
        }
    }
    return nullptr;
}

const DescriptorBinding* OpenGLDescriptorSet::FindBinding(uint32_t binding) const {
    for (const auto& descriptorBinding : m_bindings) {
        if (descriptorBinding.binding == binding) {
            return &descriptorBinding;
        }
    }
    return nullptr;
}

bool OpenGLDescriptorSet::Initialize(OpenGLTextureUnitAllocator* allocator, uint32_t count, const DescriptorSetLayoutDesc& layout) {
    m_allocator = allocator;
    m_bindingCount = layout.bindingCount;

    if (layout.bindingCount > 0 && layout.bindings != nullptr) {
        m_layoutBindings = new DescriptorSetLayoutBinding[layout.bindingCount];
        for (uint32_t i = 0; i < layout.bindingCount; ++i) {
            m_layoutBindings[i] = layout.bindings[i];
        }
    }

    m_bindings.resize(layout.bindingCount);
    for (uint32_t i = 0; i < layout.bindingCount; ++i) {
        m_bindings[i].binding = layout.bindings[i].binding;
        m_bindings[i].type = static_cast<DescriptorType>(layout.bindings[i].type);
        m_bindings[i].count = layout.bindings[i].count;
        m_bindings[i].resourceDimension = layout.bindings[i].resourceDimension;
        m_bindings[i].textureUnits.resize(layout.bindings[i].count);
        m_bindings[i].textureIds.resize(layout.bindings[i].count, 0);
        m_bindings[i].textureTargets.resize(layout.bindings[i].count, GL_TEXTURE_2D);
        m_bindings[i].bufferIds.resize(layout.bindings[i].count, 0);
        m_bindings[i].bufferOffsets.resize(layout.bindings[i].count, 0);
        m_bindings[i].bufferSizes.resize(layout.bindings[i].count, 0);
        m_bindings[i].samplerIds.resize(layout.bindings[i].count, 0);

        if (UsesBufferBinding(m_bindings[i].resourceDimension)) {
            m_bindings[i].textureUnits.clear();
            continue;
        }

        if (m_bindings[i].type != DescriptorType::SRV &&
            m_bindings[i].type != DescriptorType::Sampler &&
            m_bindings[i].type != DescriptorType::UAV) {
            m_bindings[i].textureUnits.clear();
            continue;
        }

        if (m_allocator == nullptr) {
            Shutdown();
            return false;
        }

        for (uint32_t j = 0; j < layout.bindings[i].count; ++j) {
            int32_t unit = m_allocator->Allocate();
            if (unit < 0) {
                Shutdown();
                return false;
            }
            m_bindings[i].textureUnits[j] = static_cast<uint32_t>(unit);
        }
    }

    return true;
}

void OpenGLDescriptorSet::Shutdown() {
    if (m_constantBuffer != 0) {
        glDeleteBuffers(1, reinterpret_cast<GLuint*>(&m_constantBuffer));
        m_constantBuffer = 0;
    }

    if (m_allocator != nullptr) {
        for (auto& binding : m_bindings) {
            for (uint32_t i = 0; i < binding.textureUnits.size(); ++i) {
                m_allocator->Free(static_cast<int32_t>(binding.textureUnits[i]));
            }
        }
    }

    m_bindings.clear();
    m_allocator = nullptr;
    m_bound = false;

    if (m_layoutBindings != nullptr) {
        delete[] m_layoutBindings;
        m_layoutBindings = nullptr;
    }
    m_bindingCount = 0;
}

void OpenGLDescriptorSet::EnsureConstantBufferUploaded() {
    if (!m_constantBufferDirty || m_constantBufferData.empty()) {
        return;
    }

    if (m_constantBuffer == 0) {
        glGenBuffers(1, reinterpret_cast<GLuint*>(&m_constantBuffer));
    }
    glBindBuffer(GL_UNIFORM_BUFFER, m_constantBuffer);
    glBufferData(GL_UNIFORM_BUFFER, m_constantBufferData.size(), m_constantBufferData.data(), GL_DYNAMIC_DRAW);
    glBindBuffer(GL_UNIFORM_BUFFER, 0);
    m_constantBufferDirty = false;
}

void OpenGLDescriptorSet::Update(uint32_t offset, RHIResourceView* view) {
    if (view == nullptr) {
        return;
    }

    DescriptorBinding* binding = FindBinding(offset);
    if (binding == nullptr) {
        return;
    }

    OpenGLResourceView* glView = static_cast<OpenGLResourceView*>(view);
    const bool expectsBuffer = UsesBufferBinding(binding->resourceDimension);
    const bool viewIsBuffer = IsBufferResourceViewDimension(glView->GetDimension());
    if (expectsBuffer != viewIsBuffer) {
        return;
    }

    if (expectsBuffer) {
        if (binding->bufferIds.empty()) {
            return;
        }

        binding->bufferIds[0] = glView->GetBuffer();
        binding->bufferOffsets[0] = glView->GetBufferOffset();
        binding->bufferSizes[0] = glView->GetBufferSize();
        return;
    }

    if (binding->textureIds.empty()) {
        return;
    }

    binding->textureIds[0] = glView->GetTexture();
    if (const OpenGLTexture* texture = glView->GetTextureResource()) {
        binding->textureTargets[0] = static_cast<uint32_t>(ToOpenGL(texture->GetOpenGLType()));
    }
}

void OpenGLDescriptorSet::UpdateSampler(uint32_t offset, RHISampler* sampler) {
    if (sampler == nullptr) {
        return;
    }

    DescriptorBinding* binding = FindBinding(offset);
    if (binding == nullptr || binding->samplerIds.empty()) {
        return;
    }

    OpenGLSampler* glSampler = static_cast<OpenGLSampler*>(sampler);
    binding->samplerIds[0] = glSampler->GetID();
}

void OpenGLDescriptorSet::Bind() {
    EnsureConstantBufferUploaded();

    if (m_constantBuffer != 0) {
        for (const auto& binding : m_bindings) {
            if (binding.type != DescriptorType::CBV) {
                continue;
            }

            for (uint32_t i = 0; i < binding.count; ++i) {
                glBindBufferBase(GL_UNIFORM_BUFFER, binding.binding + i, m_constantBuffer);
            }
        }
    }

    for (size_t i = 0; i < m_bindings.size(); ++i) {
        const auto& binding = m_bindings[i];

        if (UsesBufferBinding(binding.resourceDimension)) {
            for (size_t j = 0; j < binding.bufferIds.size(); ++j) {
                const GLuint bufferId = binding.bufferIds[j];
                if (bufferId == 0) {
                    continue;
                }

                const GLuint rangeSize = j < binding.bufferSizes.size() ? binding.bufferSizes[j] : 0u;
                if (rangeSize > 0) {
                    glBindBufferRange(
                        GL_SHADER_STORAGE_BUFFER,
                        binding.binding + static_cast<uint32_t>(j),
                        bufferId,
                        static_cast<GLintptr>(binding.bufferOffsets[j]),
                        static_cast<GLsizeiptr>(rangeSize));
                } else {
                    glBindBufferBase(
                        GL_SHADER_STORAGE_BUFFER,
                        binding.binding + static_cast<uint32_t>(j),
                        bufferId);
                }
            }
            continue;
        }

        for (size_t j = 0; j < binding.textureUnits.size(); ++j) {
            uint32_t unit = binding.textureUnits[j];
            uint32_t textureId = binding.textureIds[j];
            uint32_t samplerId = binding.samplerIds[j];

            if (textureId != 0 && binding.type != DescriptorType::Sampler) {
                glActiveTexture(GL_TEXTURE0 + unit);
                glBindTexture(
                    j < binding.textureTargets.size() ? binding.textureTargets[j] : GL_TEXTURE_2D,
                    textureId);
            }

            if (samplerId != 0) {
                glBindSampler(unit, samplerId);
            }
        }
    }

    m_bound = true;
}

void OpenGLDescriptorSet::BindWithPipelineLayout(const OpenGLPipelineLayout* pipelineLayout, uint32_t setIndex) {
    EnsureConstantBufferUploaded();

    if (m_constantBuffer != 0) {
        for (const auto& binding : m_bindings) {
            if (binding.type != DescriptorType::CBV) {
                continue;
            }

            const uint32_t baseBindingPoint = ResolveBindingPoint(pipelineLayout, setIndex, binding);
            if (baseBindingPoint == UINT32_MAX) {
                continue;
            }

            for (uint32_t i = 0; i < binding.count; ++i) {
                glBindBufferBase(GL_UNIFORM_BUFFER, baseBindingPoint + i, m_constantBuffer);
            }
        }
    }

    for (const auto& binding : m_bindings) {
        const uint32_t baseBindingPoint = ResolveBindingPoint(pipelineLayout, setIndex, binding);
        if (baseBindingPoint == UINT32_MAX) {
            continue;
        }

        if (UsesBufferBinding(binding.resourceDimension)) {
            for (size_t i = 0; i < binding.bufferIds.size(); ++i) {
                const GLuint bufferId = binding.bufferIds[i];
                if (bufferId == 0) {
                    continue;
                }

                const uint32_t bindingPoint = baseBindingPoint + static_cast<uint32_t>(i);
                const GLuint rangeSize = i < binding.bufferSizes.size() ? binding.bufferSizes[i] : 0u;
                if (rangeSize > 0) {
                    glBindBufferRange(
                        GL_SHADER_STORAGE_BUFFER,
                        bindingPoint,
                        bufferId,
                        static_cast<GLintptr>(binding.bufferOffsets[i]),
                        static_cast<GLsizeiptr>(rangeSize));
                } else {
                    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, bindingPoint, bufferId);
                }
            }
            continue;
        }

        for (size_t i = 0; i < binding.textureIds.size(); ++i) {
            const uint32_t bindingPoint = baseBindingPoint + static_cast<uint32_t>(i);
            const uint32_t textureId = binding.textureIds[i];
            if (textureId != 0) {
                if (binding.type == DescriptorType::UAV) {
                    glBindImageTexture(bindingPoint, textureId, 0, GL_FALSE, 0, GL_READ_WRITE, GL_RGBA8);
                } else if (binding.type != DescriptorType::Sampler) {
                    glActiveTexture(GL_TEXTURE0 + bindingPoint);
                    glBindTexture(
                        i < binding.textureTargets.size() ? binding.textureTargets[i] : GL_TEXTURE_2D,
                        textureId);
                }
            }

            if (binding.type == DescriptorType::Sampler && i < binding.samplerIds.size()) {
                const uint32_t samplerId = binding.samplerIds[i];
                if (samplerId != 0) {
                    glBindSampler(bindingPoint, samplerId);
                }
            }
        }
    }

    m_bound = true;
}

void OpenGLDescriptorSet::Unbind() {
    for (size_t i = 0; i < m_bindings.size(); ++i) {
        const auto& binding = m_bindings[i];

        if (binding.type == DescriptorType::CBV) {
            for (uint32_t j = 0; j < binding.count; ++j) {
                glBindBufferBase(GL_UNIFORM_BUFFER, binding.binding + j, 0);
            }
        }

        if (UsesBufferBinding(binding.resourceDimension)) {
            for (size_t j = 0; j < binding.bufferIds.size(); ++j) {
                glBindBufferBase(GL_SHADER_STORAGE_BUFFER, binding.binding + static_cast<uint32_t>(j), 0);
            }
            continue;
        }

        for (size_t j = 0; j < binding.textureUnits.size(); ++j) {
            uint32_t unit = binding.textureUnits[j];

            glActiveTexture(GL_TEXTURE0 + unit);
            glBindTexture(
                j < binding.textureTargets.size() ? binding.textureTargets[j] : GL_TEXTURE_2D,
                0);
            glBindSampler(unit, 0);
        }
    }

    m_bound = false;
}

uint32_t OpenGLDescriptorSet::GetBindingPoint(uint32_t binding) const {
    const DescriptorBinding* descriptorBinding = FindBinding(binding);
    if (descriptorBinding != nullptr && !descriptorBinding->textureUnits.empty()) {
        return descriptorBinding->textureUnits[0];
    }
    return 0;
}

void OpenGLDescriptorSet::WriteConstant(uint32_t binding, const void* data, size_t size, size_t offset) {
    size_t requiredSize = offset + size;
    if (m_constantBufferData.size() < requiredSize) {
        m_constantBufferData.resize(requiredSize);
    }
    memcpy(m_constantBufferData.data() + offset, data, size);
    m_constantBufferDirty = true;
}

} // namespace RHI
} // namespace XCEngine