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

namespace XCEngine {
namespace RHI {

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].textureUnits.resize(layout.bindings[i].count);
        m_bindings[i].textureIds.resize(layout.bindings[i].count, 0);
        m_bindings[i].samplerIds.resize(layout.bindings[i].count, 0);

        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;

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

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

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

    OpenGLResourceView* glView = static_cast<OpenGLResourceView*>(view);
    binding->textureIds[0] = glView->GetTexture();
}

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() {
    if (m_constantBufferDirty && !m_constantBufferData.empty()) {
        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);
        glBindBufferBase(GL_UNIFORM_BUFFER, 0, m_constantBuffer);
        m_constantBufferDirty = false;
    }

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

        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(GL_TEXTURE_2D, textureId);
            }

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

    m_bound = true;
}

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

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

            glActiveTexture(GL_TEXTURE0 + unit);
            glBindTexture(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