#include "black_hole_renderer.h"

#include "core/log.h"
#include "core/settings_manager.h"

#include <algorithm>
#include <cstdint>
#include <cstring>

namespace Donut
{
    using namespace RHI;

    namespace
    {
        constexpr float SagA_rs = 1.269e10f;   // Schwarzschild radius of the modelled hole

        struct CamUBO {
            glm::vec3 pos;  float p0; glm::vec3 right; float p1;
            glm::vec3 up;   float p2; glm::vec3 fwd;   float p3;
            float tan_half_fov; float aspect; uint32_t moving; int p4;
        };
        struct SimUBO { int steps_moving; int steps_static; float early_exit; float time; };
    }

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

        m_geo_lo = device.create_render_target(GEO_LO_W, GEO_LO_H, Format::RGBA8, Format::None, Filter::Linear);
        m_geo_hi = device.create_render_target(GEO_HI_W, GEO_HI_H, Format::RGBA8, Format::None, Filter::Linear);

        // Fullscreen quad shared by the geodesic and present passes: pos.xy + uv.
        const float quad[] = {
            -1.f,  1.f, 0.f, 1.f,  -1.f, -1.f, 0.f, 0.f,   1.f, -1.f, 1.f, 0.f,
            -1.f,  1.f, 0.f, 1.f,   1.f, -1.f, 1.f, 0.f,   1.f,  1.f, 1.f, 1.f,
        };
        m_quad_vb  = device.create_buffer(BufferType::Vertex, sizeof(quad), quad);

        m_cam_ubo  = device.create_buffer(BufferType::Uniform, 128);
        m_disk_ubo = device.create_buffer(BufferType::Uniform, 32);
        m_obj_ubo  = device.create_buffer(BufferType::Uniform, 800);
        m_sim_ubo  = device.create_buffer(BufferType::Uniform, 16);

        // The hole itself is the hardcoded singularity in the shader; the Objects
        // UBO carries the Scene's spheres, refilled each frame in render_geodesic.
        std::vector<uint8_t> obj(800, 0);   // numObjects = 0 until the first frame
        m_obj_ubo->update(obj.data(), obj.size());

        {
            PipelineDesc d;
            d.shader = "geodesic";
            d.vertex_layout = { 16, { { 0, 2, 0 }, { 1, 2, 8 } } };
            d.resources = {
                { ResourceKind::UniformBuffer, 0, "Camera" },
                { ResourceKind::UniformBuffer, 1, "Disk" },
                { ResourceKind::UniformBuffer, 2, "Objects" },
                { ResourceKind::UniformBuffer, 3, "Simulation" },
                { ResourceKind::Texture,       4, "u_HDRIEnvironment" },
            };
            d.topology = Topology::Triangles;
            d.target = { Format::RGBA8, Format::None };   // colour-only off-screen target
            m_geo_pipeline = device.create_pipeline(d);
        }
        {
            PipelineDesc d;
            d.shader = "textured_quad";
            d.vertex_layout = { 16, { { 0, 2, 0 }, { 1, 2, 8 } } };
            d.resources = { { ResourceKind::Texture, 0, "u_ScreenTexture" } };
            d.topology = Topology::Triangles;
            d.target = { Format::Swapchain, Format::D32 }; // swapchain target (depth unused)
            m_present_pipeline = device.create_pipeline(d);
        }

        DONUT_INFO("BlackHoleRenderer ready ({}x{} moving / {}x{} settled)",
                   GEO_LO_W, GEO_LO_H, GEO_HI_W, GEO_HI_H);
        return true;
    }

    auto BlackHoleRenderer::render_geodesic(RHI::CommandList& cmd, const GeodesicView& view,
                                            const BlackHoleParams& params,
                                            const std::vector<SceneObject>& objects, RHI::Texture* cubemap) -> void
    {
        CamUBO cam{};
        cam.pos = view.position; cam.right = view.right; cam.up = view.up; cam.fwd = view.forward;
        cam.tan_half_fov = view.tan_half_fov;
        cam.aspect = view.aspect;
        cam.moving = view.moving ? 1u : 0u;
        m_cam_ubo->update(&cam, sizeof(cam));

        // Same integration budget whether moving or settled (the disk vanishes at
        // steep poses below ~15000 steps); responsiveness comes from the lower-res
        // target instead. Sourced from the UI, clamped GPU-safe.
        constexpr int kStepCeil = 15000;
        SimUBO sim{};
        sim.steps_static = std::clamp(params.quality_steps, 1000, kStepCeil);
        sim.steps_moving = sim.steps_static;
        sim.early_exit   = SettingsManager::get_early_exit_distance();
        sim.time         = view.time;
        m_sim_ubo->update(&sim, sizeof(sim));

        // Disk struct: r_in, r_out, turbulence, slab half-thickness, brightness,
        // temperature (radii in Schwarzschild radii).
        float disk[8] = {
            std::max(params.disk_inner_rs, 3.0f) * SagA_rs,
            std::max(params.disk_outer_rs, params.disk_inner_rs + 0.5f) * SagA_rs,
            std::max(params.turbulence, 0.0f),
            SagA_rs * 0.1f,
            std::max(params.brightness, 0.0f),
            std::max(params.temperature, 1000.0f),
            0.0f, 0.0f,
        };
        m_disk_ubo->update(disk, sizeof(disk));

        // Scene objects → the geodesic's Objects UBO (std140: numObjects@0,
        // objPosRadius[i]@16+16i, objColor[i]@272+16i). They render as spheres the
        // curved rays intersect, so the hole lenses them. 1 scene unit = 1 r_s.
        {
            const float k = SagA_rs / SCENE_UNITS_PER_RS;   // scene units -> metres
            std::vector<uint8_t> objbuf(800, 0);
            int n = std::min((int)objects.size(), 16);
            std::memcpy(objbuf.data(), &n, 4);
            for (int i = 0; i < n; ++i)
            {
                const SceneObject& o = objects[i];
                float pr[4]  = { o.position.x * k, o.position.y * k, o.position.z * k, o.radius * k };
                float col[4] = { o.color.r, o.color.g, o.color.b, 1.0f };
                std::memcpy(objbuf.data() + 16  + 16 * i, pr,  16);
                std::memcpy(objbuf.data() + 272 + 16 * i, col, 16);
            }
            m_obj_ubo->update(objbuf.data(), objbuf.size());
        }

        // Progressive resolution: small target while moving, large once settled.
        RenderTarget* target = view.moving ? m_geo_lo.get() : m_geo_hi.get();
        m_last_target = target;

        cmd.begin_render_pass(target, glm::vec4(0, 0, 0, 1));
        cmd.set_viewport(0, 0, target->width(), target->height(), false);
        cmd.bind_pipeline(m_geo_pipeline.get());
        cmd.bind_uniform(0, m_cam_ubo.get());
        cmd.bind_uniform(1, m_disk_ubo.get());
        cmd.bind_uniform(2, m_obj_ubo.get());
        cmd.bind_uniform(3, m_sim_ubo.get());
        cmd.bind_texture(4, cubemap);
        cmd.bind_vertex_buffer(m_quad_vb.get());
        cmd.draw(6);
        cmd.end_render_pass();
    }

    auto BlackHoleRenderer::blit(RHI::CommandList& cmd, int fb_width, int fb_height) -> void
    {
        if (!m_last_target) return;
        // flip_y matches the shared top-left orientation (Vulkan flips via a
        // negative-height viewport; GL is a no-op).
        cmd.set_viewport(0, 0, fb_width, fb_height, true);
        cmd.bind_pipeline(m_present_pipeline.get());
        cmd.bind_texture(0, m_last_target->color_texture());
        cmd.bind_vertex_buffer(m_quad_vb.get());
        cmd.draw(6);
    }
}