#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, false, Filter::Linear);
        m_geo_hi = device.create_render_target(GEO_HI_W, GEO_HI_H, Format::RGBA8, false, 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 single "object" is the black hole itself (position 0, radius = r_s).
        std::vector<uint8_t> obj(800, 0);
        int num_objects = 1;                          std::memcpy(obj.data() +   0, &num_objects, 4);
        float pos_radius[4] = { 0, 0, 0, SagA_rs };   std::memcpy(obj.data() +  16, pos_radius, 16);
        float color[4]      = { 0, 0, 0, 1 };         std::memcpy(obj.data() + 272, color, 16);
        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.has_depth = false;   // color-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.has_depth = true;    // swapchain target carries depth (unused here)
            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, 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));

        // 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);
    }
}