#include "opengl_device.h"

#include "core/log.h"
#include "core/hdri_manager.h"
#include "shader.h"
#include "texture.h"

#include <glad/glad.h>
#include <GLFW/glfw3.h>
#include <algorithm>
#include <imgui.h>
#include <imgui_impl_glfw.h>
#include <imgui_impl_opengl3.h>

namespace Donut::RHI
{
    namespace
    {
        auto gl_topology(Topology t) -> GLenum { return t == Topology::Lines ? GL_LINES : GL_TRIANGLES; }
        auto gl_compare(CompareOp o) -> GLenum { return o == CompareOp::Always ? GL_ALWAYS : o == CompareOp::LessEqual ? GL_LEQUAL : GL_LESS; }
        auto gl_filter(Filter f) -> GLint { return f == Filter::Nearest ? GL_NEAREST : GL_LINEAR; }
        auto gl_internal(Format f) -> GLint { return f == Format::RGBA16F ? GL_RGBA16F : f == Format::RGBA32F ? GL_RGBA32F : f == Format::D32 ? GL_DEPTH_COMPONENT32F : GL_RGBA8; }

        class GLBuffer : public Buffer
        {
        public:
            GLBuffer(BufferType type, size_t size, const void* data)
            {
                m_target = type == BufferType::Index ? GL_ELEMENT_ARRAY_BUFFER
                         : type == BufferType::Uniform ? GL_UNIFORM_BUFFER : GL_ARRAY_BUFFER;
                glGenBuffers(1, &m_id);
                glBindBuffer(m_target, m_id);
                glBufferData(m_target, (GLsizeiptr)size, data, GL_DYNAMIC_DRAW);
            }
            ~GLBuffer() override { if (m_id) glDeleteBuffers(1, &m_id); }
            auto update(const void* data, size_t size) -> void override
            {
                glBindBuffer(m_target, m_id);
                glBufferSubData(m_target, 0, (GLsizeiptr)size, data);
            }
            GLuint m_id = 0; GLenum m_target = GL_ARRAY_BUFFER;
        };

        class GLTexture : public Texture
        {
        public:
            GLTexture(int w, int h, Format fmt, Filter filter, const void* data)
            {
                glGenTextures(1, &m_id); m_own = true; m_target = GL_TEXTURE_2D;
                glBindTexture(GL_TEXTURE_2D, m_id);
                GLenum ext = fmt == Format::D32 ? GL_DEPTH_COMPONENT : GL_RGBA;
                GLenum type = fmt == Format::RGBA8 ? GL_UNSIGNED_BYTE : GL_FLOAT;
                glTexImage2D(GL_TEXTURE_2D, 0, gl_internal(fmt), w, h, 0, ext, type, data);
                glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, gl_filter(filter));
                glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, gl_filter(filter));
                glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
                glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
            }
            explicit GLTexture(Ref<CubemapTexture> cube)
                : m_cube(std::move(cube))
            {
                m_target = GL_TEXTURE_CUBE_MAP;
                m_id = m_cube ? m_cube->get_renderer_id() : 0;
            }
            ~GLTexture() override { if (m_own && m_id) glDeleteTextures(1, &m_id); }
            GLuint m_id = 0; GLenum m_target = GL_TEXTURE_2D; bool m_own = false;
            Ref<CubemapTexture> m_cube;   // keepalive for cubemaps
        };

        class GLRenderTarget : public RenderTarget
        {
        public:
            GLRenderTarget(int w, int h, Format color, bool depth, Filter filter)
                : m_w(w), m_h(h)
            {
                m_color = create_scope<GLTexture>(w, h, color, filter, nullptr);
                glGenFramebuffers(1, &m_fbo);
                glBindFramebuffer(GL_FRAMEBUFFER, m_fbo);
                glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, m_color->m_id, 0);
                if (depth)
                {
                    glGenRenderbuffers(1, &m_depth);
                    glBindRenderbuffer(GL_RENDERBUFFER, m_depth);
                    glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH_COMPONENT24, w, h);
                    glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_RENDERBUFFER, m_depth);
                }
                if (glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE)
                    DONUT_ERROR("GL RHI: render target incomplete");
                glBindFramebuffer(GL_FRAMEBUFFER, 0);
            }
            ~GLRenderTarget() override
            {
                if (m_depth) glDeleteRenderbuffers(1, &m_depth);
                if (m_fbo) glDeleteFramebuffers(1, &m_fbo);
            }
            auto width()  const -> int override { return m_w; }
            auto height() const -> int override { return m_h; }
            auto color_texture() -> Texture* override { return m_color.get(); }
            GLuint m_fbo = 0, m_depth = 0; int m_w, m_h; Scope<GLTexture> m_color;
        };

        class GLPipeline : public Pipeline
        {
        public:
            explicit GLPipeline(const PipelineDesc& d) : m_desc(d)
            {
                std::string path = "assets/shaders/generated/" + d.shader + ".glsl";
                m_shader = Ref<Shader>(Shader::create(path));
                m_prog = m_shader ? m_shader->get_renderer_id() : 0;
                if (!m_prog) { DONUT_ERROR("GL RHI: shader '{}' failed", d.shader); return; }

                glUseProgram(m_prog);
                for (const auto& r : d.resources)
                {
                    if (r.kind == ResourceKind::UniformBuffer)
                    {
                        GLuint bi = glGetUniformBlockIndex(m_prog, r.name.c_str());
                        if (bi != GL_INVALID_INDEX) glUniformBlockBinding(m_prog, bi, r.binding);
                        // not found: the ConstantBuffer<> type name in the shader no longer
                        // matches this resource string, so the UBO silently won't bind.
                        else DONUT_WARN("GL RHI: shader '{}' has no uniform block '{}' (binding {})",
                                        d.shader, r.name, r.binding);
                    }
                    else // texture: point the sampler at texture unit == binding
                    {
                        GLint loc = glGetUniformLocation(m_prog, r.name.c_str());
                        if (loc >= 0) glUniform1i(loc, (GLint)r.binding);
                        else DONUT_WARN("GL RHI: shader '{}' has no sampler '{}' (unit {})",
                                        d.shader, r.name, r.binding);
                    }
                }
                glUseProgram(0);
            }
            PipelineDesc m_desc; Ref<Shader> m_shader; GLuint m_prog = 0;
        };

        // GL has no command buffers, so every call here executes immediately.
        class GLCommandList : public CommandList
        {
        public:
            explicit GLCommandList(GLuint vao) : m_vao(vao) {}

            auto begin_render_pass(RenderTarget* target, const glm::vec4& clear) -> void override
            {
                auto* rt = static_cast<GLRenderTarget*>(target);
                glBindFramebuffer(GL_FRAMEBUFFER, rt ? rt->m_fbo : 0);
                int w = rt ? rt->m_w : m_default_w, h = rt ? rt->m_h : m_default_h;
                glViewport(0, 0, w, h);
                glDisable(GL_SCISSOR_TEST);
                glDepthMask(GL_TRUE);
                glClearColor(clear.r, clear.g, clear.b, clear.a);
                glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
            }
            auto end_render_pass() -> void override {}   // GL passes are a fiction; nothing to close

            auto bind_pipeline(Pipeline* p) -> void override
            {
                m_pipe = static_cast<GLPipeline*>(p);
                glUseProgram(m_pipe->m_prog);
                const auto& d = m_pipe->m_desc;
                if (d.depth_test) { glEnable(GL_DEPTH_TEST); glDepthFunc(gl_compare(d.depth_op)); }
                else glDisable(GL_DEPTH_TEST);
                glDepthMask(d.depth_write ? GL_TRUE : GL_FALSE);
                if (d.cull == CullMode::None) glDisable(GL_CULL_FACE);
                else { glEnable(GL_CULL_FACE); glCullFace(d.cull == CullMode::Back ? GL_BACK : GL_FRONT); }
                if (d.blend == BlendMode::AlphaBlend) { glEnable(GL_BLEND); glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA); }
                else glDisable(GL_BLEND);
            }
            auto set_viewport(int x, int y, int w, int h, bool /*flip_y*/) -> void override { glViewport(x, y, w, h); }
            auto bind_uniform(uint32_t binding, Buffer* ubo) -> void override
            {
                glBindBufferBase(GL_UNIFORM_BUFFER, binding, static_cast<GLBuffer*>(ubo)->m_id);
            }
            auto bind_texture(uint32_t binding, Texture* tex) -> void override
            {
                auto* t = static_cast<GLTexture*>(tex);
                glActiveTexture(GL_TEXTURE0 + binding);
                glBindTexture(t->m_target, t->m_id);
            }
            auto bind_vertex_buffer(Buffer* vb) -> void override
            {
                glBindBuffer(GL_ARRAY_BUFFER, static_cast<GLBuffer*>(vb)->m_id);
                const auto& layout = m_pipe->m_desc.vertex_layout;
                for (const auto& a : layout.attributes)
                {
                    glEnableVertexAttribArray(a.location);
                    glVertexAttribPointer(a.location, (GLint)a.components, GL_FLOAT, GL_FALSE,
                                          (GLsizei)layout.stride, (const void*)(uintptr_t)a.offset);
                }
            }
            auto bind_index_buffer(Buffer* ib) -> void override
            {
                glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, static_cast<GLBuffer*>(ib)->m_id);
            }
            auto draw(uint32_t count) -> void override
            {
                glDrawArrays(gl_topology(m_pipe->m_desc.topology), 0, (GLsizei)count);
            }
            auto draw_indexed(uint32_t count) -> void override
            {
                glDrawElements(gl_topology(m_pipe->m_desc.topology), (GLsizei)count, GL_UNSIGNED_INT, nullptr);
            }

            GLuint m_vao; GLPipeline* m_pipe = nullptr; int m_default_w = 1, m_default_h = 1;
        };

        class GLDevice : public Device
        {
        public:
            auto init(const NativeWindow& window) -> bool override
            {
                m_window = window.glfw_handle;
                glfwMakeContextCurrent(static_cast<GLFWwindow*>(m_window));
                if (!gladLoadGLLoader((GLADloadproc)glfwGetProcAddress)) { DONUT_ERROR("GL RHI: GLAD load failed"); return false; }
                m_width = window.width; m_height = window.height;
                const char* name = (const char*)glGetString(GL_RENDERER);
                m_name = name ? name : "OpenGL";
                glGenVertexArrays(1, &m_vao); glBindVertexArray(m_vao);   // GL core wants one VAO bound at all times, so like, here's one
                m_cmds = create_scope<GLCommandList>(m_vao);
                DONUT_INFO("GL RHI device: {} ({}x{})", m_name, m_width, m_height);
                return true;
            }
            auto shutdown() -> void override
            {
                // nuke the HDRIManager's cached GL cubemaps while our context is
                // still alive. they're static-singleton-owned, so if we let them
                // destruct at program exit they'd glDeleteTextures into a dead
                // context and crash on close. no thanks.
                HDRIManager::get().clear_cache();
                if (m_imgui) { ImGui_ImplOpenGL3_Shutdown(); ImGui_ImplGlfw_Shutdown(); ImGui::DestroyContext(); m_imgui = false; }
                if (m_vao) { glDeleteVertexArrays(1, &m_vao); m_vao = 0; }
            }
            auto resize(int w, int h) -> void override { m_width = w; m_height = h; }
            auto wait_idle() -> void override { glFinish(); }
            auto set_vsync(bool enabled) -> void override { glfwSwapInterval(enabled ? 1 : 0); }

            auto create_buffer(BufferType t, size_t size, const void* data) -> Ref<Buffer> override
            { return create_ref<GLBuffer>(t, size, data); }
            auto create_texture(int w, int h, Format f, Filter fl, const void* data) -> Ref<Texture> override
            { return create_ref<GLTexture>(w, h, f, fl, data); }
            auto create_cubemap_from_hdri(const std::string& path) -> Ref<Texture> override
            {
                HDRIManager::get().set_current_hdri(path);
                return create_ref<GLTexture>(HDRIManager::get().get_current_hdri());
            }
            auto create_render_target(int w, int h, Format color, Format depth, Filter fl, int /*mips*/) -> Ref<RenderTarget> override
            { return create_ref<GLRenderTarget>(w, h, color, depth != Format::None, fl); }
            auto create_pipeline(const PipelineDesc& d) -> Ref<Pipeline> override
            { return create_ref<GLPipeline>(d); }

            auto begin_frame(const glm::vec4& /*clear*/) -> CommandList* override
            {
                glfwGetFramebufferSize(static_cast<GLFWwindow*>(m_window), &m_width, &m_height);
                m_cmds->m_default_w = m_width; m_cmds->m_default_h = m_height;
                glBindVertexArray(m_vao);
                return m_cmds.get();
            }
            auto end_frame() -> void override { glfwSwapBuffers(static_cast<GLFWwindow*>(m_window)); }

            auto run_offscreen(const std::function<void(CommandList&)>& record) -> void override
            {
                glBindVertexArray(m_vao);   // GL is immediate: record executes now
                record(*m_cmds);
                glFinish();
            }
            auto read_render_target(RenderTarget* target, std::vector<uint8_t>& out) -> void override
            {
                auto* rt = static_cast<GLRenderTarget*>(target);
                int w = rt->m_w, h = rt->m_h;
                out.resize((size_t)w * h * 4);
                glBindFramebuffer(GL_FRAMEBUFFER, rt->m_fbo);
                glReadPixels(0, 0, w, h, GL_RGBA, GL_UNSIGNED_BYTE, out.data());
                glBindFramebuffer(GL_FRAMEBUFFER, 0);
                const size_t row = (size_t)w * 4;   // GL is bottom-up -> flip to top-left
                for (int y = 0; y < h / 2; ++y)
                    std::swap_ranges(out.begin() + y * row, out.begin() + (y + 1) * row, out.begin() + (h - 1 - y) * row);
            }
            auto read_render_target_float(RenderTarget* target, std::vector<float>& out) -> void override
            {
                auto* rt = static_cast<GLRenderTarget*>(target);
                int w = rt->m_w, h = rt->m_h;
                out.resize((size_t)w * h * 4);
                glBindFramebuffer(GL_FRAMEBUFFER, rt->m_fbo);
                glReadPixels(0, 0, w, h, GL_RGBA, GL_FLOAT, out.data());   // GL converts to float32
                glBindFramebuffer(GL_FRAMEBUFFER, 0);
                const size_t row = (size_t)w * 4;   // GL is bottom-up -> flip to top-left
                for (int y = 0; y < h / 2; ++y)
                    std::swap_ranges(out.begin() + y * row, out.begin() + (y + 1) * row, out.begin() + (h - 1 - y) * row);
            }

            auto init_imgui() -> void override
            {
                IMGUI_CHECKVERSION(); ImGui::CreateContext();
                ImGui::GetIO().ConfigFlags |= ImGuiConfigFlags_DockingEnable;
                ImGui_ImplGlfw_InitForOpenGL(static_cast<GLFWwindow*>(m_window), true);
                ImGui_ImplOpenGL3_Init("#version 410");
                m_imgui = true;
                DONUT_INFO("GL RHI: ImGui backend initialized");
            }
            auto imgui_new_frame() -> void override
            {
                ImGui_ImplOpenGL3_NewFrame(); ImGui_ImplGlfw_NewFrame(); ImGui::NewFrame();
            }
            auto imgui_render(CommandList&) -> void override
            {
                ImGui::Render(); ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
            }
            auto device_name() const -> const std::string& override { return m_name; }

        private:
            void* m_window = nullptr; std::string m_name; GLuint m_vao = 0;
            int m_width = 0, m_height = 0; bool m_imgui = false;
            Scope<GLCommandList> m_cmds;
        };
    }

    auto create_opengl_device() -> Scope<Device> { return create_scope<GLDevice>(); }
}