#include "ui_layer.h"

#include "scene/scene.h"
#include "core/settings_manager.h"
#include "core/hdri_manager.h"
#include "rendering/render_api.h"

#include <imgui.h>
#include <ImGuizmo.h>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <algorithm>
#include <cstdlib>
#include <cstdio>
#include <string>

namespace Donut
{
    // One-time ImGui theme: rounded, roomy, dark with a warm accretion-disk accent.
    static auto apply_donut_style() -> void
    {
        ImGuiStyle& s = ImGui::GetStyle();
        s.WindowRounding = 7.0f; s.ChildRounding = 5.0f; s.FrameRounding = 4.0f;
        s.GrabRounding = 4.0f; s.PopupRounding = 4.0f; s.ScrollbarRounding = 5.0f; s.TabRounding = 4.0f;
        s.WindowPadding = ImVec2(12, 12); s.FramePadding = ImVec2(9, 5);
        s.ItemSpacing = ImVec2(9, 8); s.ItemInnerSpacing = ImVec2(7, 5);
        s.WindowBorderSize = 0.0f; s.FrameBorderSize = 0.0f; s.WindowTitleAlign = ImVec2(0.02f, 0.5f);

        const ImVec4 amber   = ImVec4(0.98f, 0.62f, 0.20f, 1.00f);
        const ImVec4 amberHi = ImVec4(1.00f, 0.73f, 0.36f, 1.00f);
        ImVec4* c = s.Colors;
        c[ImGuiCol_WindowBg]        = ImVec4(0.07f, 0.08f, 0.10f, 0.97f);
        c[ImGuiCol_ChildBg]         = ImVec4(0.10f, 0.11f, 0.13f, 0.55f);
        c[ImGuiCol_PopupBg]         = ImVec4(0.09f, 0.10f, 0.12f, 0.98f);
        c[ImGuiCol_TitleBg]         = ImVec4(0.09f, 0.10f, 0.12f, 1.00f);
        c[ImGuiCol_TitleBgActive]   = ImVec4(0.13f, 0.14f, 0.17f, 1.00f);
        c[ImGuiCol_Text]            = ImVec4(0.90f, 0.91f, 0.93f, 1.00f);
        c[ImGuiCol_TextDisabled]    = ImVec4(0.48f, 0.50f, 0.54f, 1.00f);
        c[ImGuiCol_FrameBg]         = ImVec4(0.16f, 0.17f, 0.20f, 1.00f);
        c[ImGuiCol_FrameBgHovered]  = ImVec4(0.22f, 0.24f, 0.28f, 1.00f);
        c[ImGuiCol_FrameBgActive]   = ImVec4(0.26f, 0.28f, 0.33f, 1.00f);
        c[ImGuiCol_Header]          = ImVec4(0.17f, 0.19f, 0.23f, 1.00f);
        c[ImGuiCol_HeaderHovered]   = ImVec4(0.24f, 0.27f, 0.32f, 1.00f);
        c[ImGuiCol_HeaderActive]    = ImVec4(0.28f, 0.31f, 0.37f, 1.00f);
        c[ImGuiCol_Button]          = ImVec4(0.20f, 0.22f, 0.26f, 1.00f);
        c[ImGuiCol_ButtonHovered]   = ImVec4(0.27f, 0.30f, 0.35f, 1.00f);
        c[ImGuiCol_ButtonActive]    = amber;
        c[ImGuiCol_SliderGrab]      = amber;
        c[ImGuiCol_SliderGrabActive]= amberHi;
        c[ImGuiCol_CheckMark]       = amberHi;
        c[ImGuiCol_Tab]             = ImVec4(0.13f, 0.14f, 0.17f, 1.00f);
        c[ImGuiCol_TabHovered]      = ImVec4(0.28f, 0.31f, 0.37f, 1.00f);
        c[ImGuiCol_TabActive]       = ImVec4(0.22f, 0.24f, 0.28f, 1.00f);
        c[ImGuiCol_Separator]       = ImVec4(0.20f, 0.22f, 0.26f, 1.00f);
        c[ImGuiCol_ScrollbarGrab]   = ImVec4(0.24f, 0.26f, 0.30f, 1.00f);
    }

    namespace
    {
        // Setup: display / renderer / overlay / environment settings
        class SetupWorkspace : public Workspace
        {
        public:
            auto name() const -> const char* override { return "Setup"; }
            auto view() const -> View override { return View::None; }
            auto on_ui(const UIContext& ctx) -> void override
            {
                auto& s = SettingsManager::get_settings();

                if (ImGui::CollapsingHeader("Display", ImGuiTreeNodeFlags_DefaultOpen))
                {
                    const bool fs = ctx.actions.is_fullscreen();
                    int cw = 0, ch = 0; ctx.actions.get_window_size(cw, ch);

                    // Resolution presets (windowed only).
                    struct Res { const char* label; int w, h; };
                    static const Res kRes[] = {
                        { "1280 x 720",  1280, 720 }, { "1600 x 900",  1600, 900 },
                        { "1920 x 1080", 1920, 1080 }, { "2560 x 1440", 2560, 1440 },
                    };
                    char cur_label[32]; std::snprintf(cur_label, sizeof(cur_label), "%d x %d", cw, ch);
                    ImGui::BeginDisabled(fs);
                    if (ImGui::BeginCombo("Resolution", cur_label))
                    {
                        for (const auto& r : kRes)
                            if (ImGui::Selectable(r.label, cw == r.w && ch == r.h))
                                ctx.actions.set_resolution(r.w, r.h);
                        ImGui::EndCombo();
                    }
                    ImGui::EndDisabled();

                    bool fs_v = fs;
                    if (ImGui::Checkbox("Fullscreen", &fs_v)) ctx.actions.set_fullscreen(fs_v);

                    if (ImGui::Checkbox("V-Sync", &s.graphics.v_sync_enabled))
                    {
                        ctx.actions.set_vsync(s.graphics.v_sync_enabled);
                        SettingsManager::save_settings();
                    }

                    ImGui::BeginDisabled(s.graphics.v_sync_enabled);
                    if (ImGui::SliderInt("Frame cap", &s.simulation.target_fps, 30, 240, "%d FPS"))
                        SettingsManager::save_settings();
                    ImGui::EndDisabled();
                    if (s.graphics.v_sync_enabled)
                        ImGui::TextDisabled("Frame cap applies only with V-Sync off.");

                    if (ImGui::SliderFloat("UI scale", &m_ui_scale, 0.8f, 1.6f, "%.2fx"))
                        ctx.actions.set_ui_scale(m_ui_scale);
                }

                if (ImGui::CollapsingHeader("Renderer", ImGuiTreeNodeFlags_DefaultOpen))
                {
                    const char* apis[] = { "OpenGL", "Vulkan" };
                    int cur = (s.graphics.render_api == "Vulkan") ? 1 : 0;
                    if (ImGui::Combo("Graphics API", &cur, apis, 2))
                    {
                        s.graphics.render_api = apis[cur];
                        SettingsManager::save_settings();
                    }
                    const char* running = (RendererAPI::get_api() == RendererAPI::API::Vulkan) ? "Vulkan" : "OpenGL";
                    if (s.graphics.render_api != running)
                        ImGui::TextColored(ImVec4(0.98f, 0.62f, 0.20f, 1.0f), "Restart to apply (%s running)", running);
                    else
                        ImGui::TextDisabled("Active backend: %s (%s)", running, ctx.device_name.c_str());
                }

                if (ImGui::CollapsingHeader("Overlay", ImGuiTreeNodeFlags_DefaultOpen))
                {
                    bool dirty = false;
                    dirty |= ImGui::Checkbox("Show FPS",        &s.graphics.show_fps);
                    dirty |= ImGui::Checkbox("Show frame time", &s.graphics.show_performance_metrics);
                    if (dirty) SettingsManager::save_settings();
                }

                if (ImGui::CollapsingHeader("Environment", ImGuiTreeNodeFlags_DefaultOpen))
                {
                    auto& hdri = HDRIManager::get();
                    std::string preview = hdri.get_hdri_name(ctx.scene.hdri_path);
                    if (ImGui::BeginCombo("Starfield", preview.c_str()))
                    {
                        for (const auto& path : hdri.get_available_hdri())
                        {
                            bool selected = (path == ctx.scene.hdri_path);
                            if (ImGui::Selectable(hdri.get_hdri_name(path).c_str(), selected))
                                ctx.scene.hdri_path = path;   // RenderPath rebuilds the cubemap next frame
                            if (selected) ImGui::SetItemDefaultFocus();
                        }
                        ImGui::EndCombo();
                    }
                    ImGui::TextDisabled("Switching rebuilds the cubemap (brief pause).");
                }
            }
        private:
            float m_ui_scale = 1.0f;
        };

        // Scene: place and edit objects
        class SceneWorkspace : public Workspace
        {
        public:
            auto name() const -> const char* override { return "Scene"; }
            auto view() const -> View override { return View::Scene; }
            auto on_ui(const UIContext& ctx) -> void override
            {
                Scene& sc = ctx.scene;
                ImGui::TextDisabled("Drag to orbit, scroll to zoom. HDRI skybox behind.");
                auto& objs = sc.objects;
                if (sc.selected_object >= (int)objs.size()) sc.selected_object = (int)objs.size() - 1;

                ImGui::Text("Spheres (%d)", (int)objs.size());
                if (ImGui::Button("Add") && objs.size() < 64)
                {
                    SceneObject o;
                    o.position = glm::vec3(((int)objs.size() % 5) * 5.0f - 10.0f, 2.0f, ((int)objs.size() / 5) * 5.0f);
                    o.radius   = 1.5f;
                    o.color    = glm::vec3(0.35f + 0.12f * (objs.size() % 5), 0.55f, 0.9f - 0.12f * (objs.size() % 4));
                    objs.push_back(o);
                    sc.selected_object = (int)objs.size() - 1;
                }
                ImGui::SameLine();
                if (ImGui::Button("Delete") && !objs.empty())
                {
                    objs.erase(objs.begin() + sc.selected_object);
                    if (sc.selected_object >= (int)objs.size()) sc.selected_object = (int)objs.size() - 1;
                }

                ImGui::BeginChild("obj_list", ImVec2(0, 120), true);
                for (int i = 0; i < (int)objs.size(); ++i)
                {
                    std::string label = "Sphere " + std::to_string(i);
                    if (ImGui::Selectable(label.c_str(), sc.selected_object == i)) sc.selected_object = i;
                }
                ImGui::EndChild();

                if (sc.selected_object >= 0 && sc.selected_object < (int)objs.size())
                {
                    SceneObject& o = objs[sc.selected_object];
                    ImGui::DragFloat3("Position", &o.position.x, 0.1f);
                    ImGui::DragFloat("Radius", &o.radius, 0.05f, 0.1f, 20.0f);
                    ImGui::ColorEdit3("Color", &o.color.x);

                    ImGui::Spacing();
                    if (ImGui::RadioButton("Move",  m_op == ImGuizmo::TRANSLATE)) m_op = ImGuizmo::TRANSLATE;
                    ImGui::SameLine();
                    if (ImGui::RadioButton("Scale", m_op == ImGuizmo::SCALE))     m_op = ImGuizmo::SCALE;
                    ImGui::SameLine(); ImGui::TextDisabled("(drag the handles in the viewport)");

                    // Viewport gizmo for the selected sphere. Uses the scene camera
                    // (its projection was set for this frame by the Application) and
                    // draws into the background draw list so it sits over the render
                    // but under the panels.
                    ImGuizmo::SetOrthographic(false);
                    ImGuizmo::SetDrawlist(ImGui::GetBackgroundDrawList());
                    ImGuiIO& io = ImGui::GetIO();
                    ImGuizmo::SetRect(0.0f, 0.0f, io.DisplaySize.x, io.DisplaySize.y);
                    glm::mat4 view = sc.scene_camera.get_view_matrix();
                    glm::mat4 proj = sc.scene_camera.get_projection_matrix();
                    glm::mat4 xform = glm::translate(glm::mat4(1.0f), o.position)
                                    * glm::scale(glm::mat4(1.0f), glm::vec3(o.radius));
                    if (ImGuizmo::Manipulate(&view[0][0], &proj[0][0], m_op, ImGuizmo::WORLD, &xform[0][0]))
                    {
                        float t[3], r[3], s[3];
                        ImGuizmo::DecomposeMatrixToComponents(&xform[0][0], t, r, s);
                        o.position = glm::vec3(t[0], t[1], t[2]);
                        o.radius   = std::max(0.05f, (s[0] + s[1] + s[2]) / 3.0f);
                    }
                }
            }
        private:
            ImGuizmo::OPERATION m_op = ImGuizmo::TRANSLATE;
        };

        // Simulation: camera + accretion disk + quality
        class SimulationWorkspace : public Workspace
        {
        public:
            auto name() const -> const char* override { return "Simulation"; }
            auto view() const -> View override { return View::BlackHole; }
            auto on_ui(const UIContext& ctx) -> void override
            {
                Scene& sc = ctx.scene;
                BlackHoleParams& bh = sc.black_hole;

                if (ImGui::CollapsingHeader("Camera", ImGuiTreeNodeFlags_DefaultOpen))
                {
                    bool ff = sc.is_free_fly();
                    if (ImGui::RadioButton("Orbital", !ff)) sc.set_free_fly(false);
                    ImGui::SameLine();
                    if (ImGui::RadioButton("Free-fly", ff)) sc.set_free_fly(true);
                    ImGui::SameLine();
                    if (ImGui::Button("Reset view")) sc.reset_sim_camera();
                    ImGui::SliderFloat("FOV", &bh.fov_degrees, 20.0f, 90.0f, "%.0f\xc2\xb0");
                    ImGui::TextDisabled(ff ? "WASD move | Q/E down-up | Shift boost | drag look"
                                           : "Drag to orbit | scroll to zoom");
                }

                if (ImGui::CollapsingHeader("Accretion disk", ImGuiTreeNodeFlags_DefaultOpen))
                {
                    ImGui::SliderFloat("Temperature", &bh.temperature, 2000.0f, 15000.0f, "%.0f K");
                    ImGui::SliderFloat("Brightness",  &bh.brightness, 0.0f, 3.0f, "%.2f");
                    ImGui::SliderFloat("Inner radius", &bh.disk_inner_rs, 3.0f, 12.0f, "%.1f r_s");
                    if (bh.disk_outer_rs < bh.disk_inner_rs + 0.5f) bh.disk_outer_rs = bh.disk_inner_rs + 0.5f;
                    ImGui::SliderFloat("Outer radius", &bh.disk_outer_rs, bh.disk_inner_rs + 0.5f, 25.0f, "%.1f r_s");
                    ImGui::SliderFloat("Turbulence",  &bh.turbulence, 0.0f, 2.0f, "%.2f");
                    ImGui::TextDisabled("Inner edge is the ISCO (3 r_s). Novikov-Thorne profile.");
                }

                if (ImGui::CollapsingHeader("Quality"))
                {
                    ImGui::SliderInt("Integration steps", &bh.quality_steps, 2000, 15000);
                    ImGui::TextDisabled("Deeper lensing at higher cost. Settled frame is 4x supersampled.");
                }
            }
        };

        // Export: pick what to render out (not wired up yet)
        class ExportWorkspace : public Workspace
        {
        public:
            auto name() const -> const char* override { return "Export"; }
            auto view() const -> View override { return View::None; }
            auto on_ui(const UIContext& ctx) -> void override
            {
                ImGui::TextDisabled("Render the black hole + observables to disk for analysis.");
                ImGui::Spacing();
                if (ImGui::CollapsingHeader("Channels", ImGuiTreeNodeFlags_DefaultOpen))
                {
                    ImGui::Checkbox("Colour (tone-mapped)",  &m_color);
                    ImGui::Checkbox("Redshift (g)",          &m_redshift);
                    ImGui::Checkbox("Emission temperature",  &m_temperature);
                    ImGui::Checkbox("Impact parameter",      &m_impact);
                    ImGui::TextDisabled("Redshift & temperature are disk-only; impact is whole-frame (false-colour).");
                }
                if (ImGui::CollapsingHeader("Output", ImGuiTreeNodeFlags_DefaultOpen))
                {
                    ImGui::Combo("Resolution", &m_resolution, "1280 x 720\0" "1920 x 1080\0" "2560 x 1440\0" "3840 x 2160\0");
                    ImGui::Combo("Format",     &m_format,     "Display image (PNG)\0" "Raw float image (PFM)\0" "Raw values (CSV)\0");
                    if (m_format == 0) ImGui::TextDisabled("Tone-mapped colour / false-coloured observables.");
                    else               ImGui::TextDisabled("Actual physical values: g, temperature (K), impact (r_s).");
                }

                ImGui::Spacing();
                const bool any = m_color || m_redshift || m_temperature || m_impact;
                ImGui::BeginDisabled(!any);
                if (ImGui::Button("Export frame", ImVec2(-1.0f, 0.0f)))
                {
                    static const int dims[][2] = { {1280,720}, {1920,1080}, {2560,1440}, {3840,2160} };
                    ExportConfig cfg;
                    cfg.color = m_color; cfg.redshift = m_redshift;
                    cfg.temperature = m_temperature; cfg.impact = m_impact;
                    cfg.width = dims[m_resolution][0]; cfg.height = dims[m_resolution][1];
                    cfg.format = static_cast<ExportFormat>(m_format);
                    m_written = ctx.actions.export_frame(cfg);
                    m_reported = true;
                }
                ImGui::EndDisabled();

                if (m_reported)
                {
                    if (m_written > 0)
                        ImGui::TextColored(ImVec4(0.55f, 0.85f, 0.45f, 1.0f), "Wrote %d PNG(s) to exports/", m_written);
                    else
                        ImGui::TextColored(ImVec4(0.95f, 0.45f, 0.35f, 1.0f), "Export failed (see log)");
                }
            }
        private:
            bool m_color = true, m_redshift = false, m_temperature = false, m_impact = false;
            int  m_resolution = 1, m_format = 0;
            int  m_written = 0; bool m_reported = false;
        };
    }

    UILayer::UILayer()
    {
        m_workspaces.push_back(create_scope<SetupWorkspace>());
        m_workspaces.push_back(create_scope<SceneWorkspace>());
        m_workspaces.push_back(create_scope<SimulationWorkspace>());
        m_workspaces.push_back(create_scope<ExportWorkspace>());

        // Open on Simulation so the black hole greets the user on launch.
        for (int i = 0; i < (int)m_workspaces.size(); ++i)
            if (std::string(m_workspaces[i]->name()) == "Simulation") { m_landing = i; m_active = i; }
    }

    auto UILayer::draw(const UIContext& ctx) -> View
    {
        static bool styled = false;
        if (!styled) { apply_donut_style(); styled = true; }

        ImGuizmo::BeginFrame();   // gizmos are drawn by the Scene workspace

        // Full-viewport dock space with a pass-through centre, so the panel docks to
        // any edge while the render shows through the middle.
        ImGui::DockSpaceOverViewport(0, ImGui::GetMainViewport(), ImGuiDockNodeFlags_PassthruCentralNode);

        ImGui::SetNextWindowSize(ImVec2(360, 620), ImGuiCond_FirstUseEver);
        ImGui::Begin("Donut");

        ImGuiIO& io = ImGui::GetIO();
        ImGui::TextDisabled("%s", ctx.device_name.empty() ? "GPU" : ctx.device_name.c_str());
        const bool show_fps = SettingsManager::get_show_fps();
        const bool show_ms  = SettingsManager::get_show_performance_metrics();
        if (show_fps) ImGui::Text("%.0f FPS", io.Framerate);
        if (show_ms)
        {
            if (show_fps) ImGui::SameLine();
            ImGui::TextDisabled("%.1f ms/frame", io.Framerate > 0 ? 1000.0f / io.Framerate : 0.0f);
        }
        ImGui::Spacing();

        if (ImGui::BeginTabBar("workspaces"))
        {
            for (int i = 0; i < (int)m_workspaces.size(); ++i)
            {
                ImGuiTabItemFlags flags = (!m_landed && i == m_landing) ? ImGuiTabItemFlags_SetSelected : 0;
                if (ImGui::BeginTabItem(m_workspaces[i]->name(), nullptr, flags))
                {
                    m_active = i;
                    ImGui::Spacing();
                    m_workspaces[i]->on_ui(ctx);
                    ImGui::EndTabItem();
                }
            }
            ImGui::EndTabBar();
            m_landed = true;
        }
        ImGui::End();

        return m_workspaces[m_active]->view();
    }
}