#include "scene_renderer.h"

#include "core/log.h"

#include <glm/gtc/matrix_transform.hpp>
#include <algorithm>
#include <cmath>
#include <numbers>
#include <vector>

namespace Donut
{
    using namespace RHI;

    namespace
    {
        // std140 layouts, matching the Slang ConstantBuffer structs exactly. glm's
        // column-major matrices upload directly (the shaders decorate the members
        // row_major and use mul(M,v), so the same bytes read correctly on every
        // backend, with no transpose).
        struct GridUBO {
            glm::mat4 view_projection;    // 0
            glm::mat4 transform;          // 64
            float grid_size; float p0[3]; // 128
            glm::vec3 grid_color;         // 144
            float grid_alpha;             // 156
            glm::vec3 camera_pos;         // 160
            float p1;                     // 172
        };
        static_assert(sizeof(GridUBO) == 176, "GridUBO std140 layout mismatch");

        struct SphereUBO {
            glm::mat4 view_projection;                    // 0
            glm::mat4 transform;                          // 64
            glm::vec3 color; float specular;              // 128, 140
            float emission; float p0[3];                  // 144
            glm::vec3 light_pos; float p1;                // 160, 172
            glm::vec3 camera_pos; int is_selected;        // 176, 188
            glm::vec3 outline_color; float outline_width; // 192, 204
        };
        static_assert(sizeof(SphereUBO) == 208, "SphereUBO std140 layout mismatch");

        struct SkyboxUBO { glm::mat4 projection; glm::mat4 view; };
        static_assert(sizeof(SkyboxUBO) == 128, "SkyboxUBO std140 layout mismatch");

        // The hole's event horizon (1 r_s) as a marker at the origin.
        constexpr float BH_SCENE_RADIUS = SCENE_UNITS_PER_RS;
    }

    auto SceneRenderer::init(RHI::Device& device) -> bool
    {
        m_device = &device;

        // Reference grid: line list on the XZ plane (+/-50, 1-unit cells). Grid.slang
        // scales by u_GridSize/50, so u_GridSize = 50 keeps it 1:1.
        {
            std::vector<glm::vec3> lines;
            const int N = 50;
            for (int i = -N; i <= N; ++i)
            {
                lines.push_back({ (float)i, 0.0f, (float)-N });
                lines.push_back({ (float)i, 0.0f, (float) N });
                lines.push_back({ (float)-N, 0.0f, (float)i });
                lines.push_back({ (float) N, 0.0f, (float)i });
            }
            m_grid_vertex_count = (int)lines.size();
            m_grid_vb  = device.create_buffer(BufferType::Vertex, lines.size() * sizeof(glm::vec3), lines.data());
            m_grid_ubo = device.create_buffer(BufferType::Uniform, sizeof(GridUBO));

            PipelineDesc d;
            d.shader = "grid";
            d.vertex_layout = { sizeof(glm::vec3), { { 0, 3, 0 } } };
            d.resources = { { ResourceKind::UniformBuffer, 0, "Grid" } };
            d.topology = Topology::Lines;
            d.blend = BlendMode::AlphaBlend;
            d.depth_test = true; d.depth_write = false; d.depth_op = CompareOp::LessEqual;
            d.target = { Format::Swapchain, Format::D32 };
            m_grid_pipeline = device.create_pipeline(d);
        }

        // Lit sphere: unit UV-sphere (pos + normal), placed/scaled per object.
        {
            std::vector<float> sv; std::vector<uint32_t> si;
            const int RINGS = 24, SECTORS = 48;
            for (int r = 0; r <= RINGS; ++r)
            {
                float phi = (float)std::numbers::pi * r / RINGS;
                for (int s = 0; s <= SECTORS; ++s)
                {
                    float theta = 2.0f * (float)std::numbers::pi * s / SECTORS;
                    float x = sinf(phi) * cosf(theta), y = cosf(phi), z = sinf(phi) * sinf(theta);
                    sv.push_back(x); sv.push_back(y); sv.push_back(z);   // position (unit)
                    sv.push_back(x); sv.push_back(y); sv.push_back(z);   // normal == position
                }
            }
            for (int r = 0; r < RINGS; ++r)
                for (int s = 0; s < SECTORS; ++s)
                {
                    uint32_t a = r * (SECTORS + 1) + s, b = a + SECTORS + 1;
                    si.push_back(a); si.push_back(b); si.push_back(a + 1);
                    si.push_back(b); si.push_back(b + 1); si.push_back(a + 1);
                }
            m_sphere_index_count = (int)si.size();
            m_sphere_vb = device.create_buffer(BufferType::Vertex, sv.size() * sizeof(float), sv.data());
            m_sphere_ib = device.create_buffer(BufferType::Index,  si.size() * sizeof(uint32_t), si.data());

            PipelineDesc d;
            d.shader = "sphere";
            d.vertex_layout = { 6 * sizeof(float), { { 0, 3, 0 }, { 1, 3, 3 * sizeof(float) } } };
            d.resources = { { ResourceKind::UniformBuffer, 0, "Sphere" },
                            { ResourceKind::Texture,       1, "u_HDRIEnvironment" } };
            d.topology = Topology::Triangles;
            d.depth_test = true; d.depth_write = true; d.depth_op = CompareOp::Less;
            d.target = { Format::Swapchain, Format::D32 };
            m_sphere_pipeline = device.create_pipeline(d);
        }

        // Skybox: a unit cube (36 verts) sampling the HDRI cubemap; the vertex
        // shader forces depth 1 (pos.xyww) so it sits behind all scene geometry.
        {
            const float cube[] = {
                -1,-1,-1,  1,-1,-1,  1, 1,-1,  1, 1,-1, -1, 1,-1, -1,-1,-1,
                -1,-1, 1,  1,-1, 1,  1, 1, 1,  1, 1, 1, -1, 1, 1, -1,-1, 1,
                -1, 1, 1, -1, 1,-1, -1,-1,-1, -1,-1,-1, -1,-1, 1, -1, 1, 1,
                 1, 1, 1,  1, 1,-1,  1,-1,-1,  1,-1,-1,  1,-1, 1,  1, 1, 1,
                -1,-1,-1,  1,-1,-1,  1,-1, 1,  1,-1, 1, -1,-1, 1, -1,-1,-1,
                -1, 1,-1,  1, 1,-1,  1, 1, 1,  1, 1, 1, -1, 1, 1, -1, 1,-1
            };
            m_skybox_vb  = device.create_buffer(BufferType::Vertex, sizeof(cube), cube);
            m_skybox_ubo = device.create_buffer(BufferType::Uniform, sizeof(SkyboxUBO));

            PipelineDesc d;
            d.shader = "skybox";
            d.vertex_layout = { 3 * sizeof(float), { { 0, 3, 0 } } };
            d.resources = { { ResourceKind::UniformBuffer, 0, "Skybox" },
                            { ResourceKind::Texture,       1, "u_Skybox" } };
            d.topology = Topology::Triangles;
            d.depth_test = false; d.depth_write = false;
            d.target = { Format::Swapchain, Format::D32 };
            m_skybox_pipeline = device.create_pipeline(d);
        }

        m_bh_ubo = device.create_buffer(BufferType::Uniform, sizeof(SphereUBO));

        DONUT_INFO("SceneRenderer ready ({} grid verts, {} sphere indices, skybox)",
                   m_grid_vertex_count, m_sphere_index_count);
        return true;
    }

    auto SceneRenderer::render(RHI::CommandList& cmd, const CameraView& cam,
                               const std::vector<SceneObject>& objects, int selected,
                               RHI::Texture* cubemap) -> void
    {
        // Full-window viewport; flip_y lets the backend match the shared top-left
        // orientation (Vulkan flips via a negative-height viewport, GL is a no-op).
        cmd.set_viewport(0, 0, cam.fb_width, cam.fb_height, true);

        glm::mat4 vp = cam.projection * cam.view;

        // Skybox background first (depth 1, no depth test/write).
        {
            SkyboxUBO sky{};
            sky.projection = cam.projection;
            sky.view       = glm::mat4(glm::mat3(cam.view));  // strip translation
            m_skybox_ubo->update(&sky, sizeof(sky));
            cmd.bind_pipeline(m_skybox_pipeline.get());
            cmd.bind_uniform(0, m_skybox_ubo.get());
            cmd.bind_texture(1, cubemap);
            cmd.bind_vertex_buffer(m_skybox_vb.get());
            cmd.draw(36);
        }

        // Opaque spheres (write depth). One UBO per object slot avoids aliasing the
        // per-draw uniforms across the deferred command stream.
        {
            int count = std::min((int)objects.size(), MAX_OBJECTS);
            while ((int)m_sphere_ubos.size() < count)
                m_sphere_ubos.push_back(m_device->create_buffer(BufferType::Uniform, sizeof(SphereUBO)));

            cmd.bind_pipeline(m_sphere_pipeline.get());
            for (int i = 0; i < count; ++i)
            {
                const SceneObject& o = objects[i];
                SphereUBO s{};
                s.view_projection = vp;
                s.transform       = glm::translate(glm::mat4(1.0f), o.position)
                                  * glm::scale(glm::mat4(1.0f), glm::vec3(o.radius));
                s.color           = o.color;
                s.specular        = 0.6f;
                s.emission        = 0.0f;
                s.light_pos       = glm::vec3(10.0f, 20.0f, 10.0f);
                s.camera_pos      = cam.position;
                s.is_selected     = (i == selected) ? 1 : 0;
                s.outline_color   = glm::vec3(1.0f, 1.0f, 0.0f);
                s.outline_width   = 0.15f;
                m_sphere_ubos[i]->update(&s, sizeof(s));

                cmd.bind_uniform(0, m_sphere_ubos[i].get());
                cmd.bind_texture(1, cubemap);
                cmd.bind_vertex_buffer(m_sphere_vb.get());
                cmd.bind_index_buffer(m_sphere_ib.get());
                cmd.draw_indexed(m_sphere_index_count);
            }

            // Black-hole marker: a near-black sphere at the origin (radius = r_s in
            // scene units), so the Scene shows where the hole the Simulation renders
            // sits. 1 grid unit == 1 Schwarzschild radius.
            SphereUBO bh{};
            bh.view_projection = vp;
            bh.transform       = glm::scale(glm::mat4(1.0f), glm::vec3(BH_SCENE_RADIUS));
            bh.color           = glm::vec3(0.02f);            // near-black event horizon
            bh.specular        = 0.10f;
            bh.light_pos       = glm::vec3(10.0f, 20.0f, 10.0f);
            bh.camera_pos      = cam.position;
            bh.is_selected     = 1;                            // reuse the rim as an
            bh.outline_color   = glm::vec3(0.98f, 0.62f, 0.20f); // amber accretion glow
            bh.outline_width   = 0.30f;
            m_bh_ubo->update(&bh, sizeof(bh));
            cmd.bind_uniform(0, m_bh_ubo.get());
            cmd.bind_texture(1, cubemap);
            cmd.bind_vertex_buffer(m_sphere_vb.get());
            cmd.bind_index_buffer(m_sphere_ib.get());
            cmd.draw_indexed(m_sphere_index_count);
        }

        // Transparent grid on top (tests depth, doesn't write).
        {
            GridUBO g{};
            g.view_projection = vp;
            g.transform       = glm::mat4(1.0f);
            g.grid_size       = 50.0f;
            g.grid_color      = glm::vec3(0.55f, 0.55f, 0.6f);
            g.grid_alpha      = 0.75f;
            g.camera_pos      = cam.position;
            m_grid_ubo->update(&g, sizeof(g));
            cmd.bind_pipeline(m_grid_pipeline.get());
            cmd.bind_uniform(0, m_grid_ubo.get());
            cmd.bind_vertex_buffer(m_grid_vb.get());
            cmd.draw(m_grid_vertex_count);
        }
    }
}