/* -*-c++-*- */
/* osgEarth - Geospatial SDK for OpenSceneGraph
 * Copyright 2018 Pelican Mapping
 * http://osgearth.org
 *
 * osgEarth is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>
 */
#pragma once

#include <osgEarthImGui/ImGuiPanel>
#include <osgEarth/Threading>
#include <osgEarth/MemoryUtils>
#include <osgEarth/GLUtils>
#include <osgEarth/ShaderLoader>
#include <chrono>
#include <list>

namespace {
    const char* render_view_normals = R"(
#version 330
#pragma vp_function oeui_render_view_normals, fragment_lighting, last
in vec3 vp_Normal;
void oeui_render_view_normals(inout vec4 color) {
    color = vec4((vp_Normal+1.0)*0.5, 1);
}
)";

    const char* render_model_normals = R"(
#version 330
#pragma vp_function oeui_render_model_normals_vs, vertex_model, last
out vec3 vp_Normal;
out vec3 oeui_model_normal;
void oeui_render_model_normals_vs(inout vec4 vertex) {
    oeui_model_normal = vp_Normal;
}
[break]
#version 330
#pragma vp_function oeui_render_view_normals_fs, fragment_lighting, last
in vec3 oeui_model_normal;
void oeui_render_view_normals_fs(inout vec4 color) {
    color = vec4( (normalize(oeui_model_normal)+1.0)*0.5, 1);
}
)";

    const char* render_fb_normals = R"(
#version 330
#pragma vp_function oeui_render_fb_normals, fragment_lighting, last
in vec3 vp_Normal;
void oeui_render_fb_normals(inout vec4 color) {
    float a = step(0.5, color.a);
    float nz = normalize(vp_Normal).z;
    color.rgb = mix(vec3(0,0,0), vec3(1,1,1), (nz+1.0)*0.5);
    color = vec4(color.rgb, a);
}
)";

    const char* render_winding = R"(
#version 450
#extension GL_NV_fragment_shader_barycentric : enable
#pragma vp_function oeui_render_winding_fs, fragment_lighting, last
void oeui_render_winding_fs(inout vec4 color) {
    color.rgb = gl_FrontFacing ? vec3(0,0.75,0) : vec3(1,0,0);
    float b = min(gl_BaryCoordNV.x, min(gl_BaryCoordNV.y, gl_BaryCoordNV.z))*28.0;
    color = vec4(mix(vec3(1), color.rgb, clamp(b,0,1)), 1.0);
}
)";

    const char* render_outlines = R"(
#version 450
#extension GL_NV_fragment_shader_barycentric : enable
#pragma vp_function oeui_render_outlines, fragment_lighting, last
#define VP_STAGE_FRAGMENT
void oeui_render_outlines(inout vec4 color) {
    float b = min(gl_BaryCoordNV.x, min(gl_BaryCoordNV.y, gl_BaryCoordNV.z))*32.0;
    float mono = dot(color.rgb, vec3(0.299, 0.587, 0.114));
    mono = mod(mono + 0.25, 1.0); 
    color = vec4(mix(vec3(mono), color.rgb, clamp(b,0,1)), color.a);
}
)";

    const char* render_ao = R"(
#version 330
#pragma vp_function oeui_render_ao, fragment_lighting, last
struct OE_PBR { float displacement, roughness, ao, metal; } oe_pbr;
void oeui_render_ao(inout vec4 color) {
    color = vec4(oe_pbr.ao, oe_pbr.ao, oe_pbr.ao, 1);
}
)";

    const char* render_roughness = R"(
#version 330
#pragma vp_function oeui_render_roughness, fragment_lighting, last
struct OE_PBR { float displacement, roughness, ao, metal; } oe_pbr;
void oeui_render_roughness(inout vec4 color) {
    color = vec4(oe_pbr.roughness, oe_pbr.roughness, oe_pbr.roughness, 1);
}
)";

    const char* render_metal = R"(
#version 330
#pragma vp_function oeui_render_metal, fragment_lighting, last
struct OE_PBR { float displacement, roughness, ao, metal; } oe_pbr;
void oeui_render_metal(inout vec4 color) {
    color = vec4(oe_pbr.metal, oe_pbr.metal, oe_pbr.metal, 1);
}
)";

    const char* render_displacement = R"(
#version 330
#pragma vp_function oeui_render_metal, fragment_lighting, last
struct OE_PBR { float displacement, roughness, ao, metal; } oe_pbr;
void oeui_render_metal(inout vec4 color) {
    color = vec4(oe_pbr.displacement, oe_pbr.displacement, oe_pbr.displacement, 1);
}
)";

    const char* render_draw_id = R"(
#version 460
#pragma vp_function oeui_render_drawid_vs, vertex_model
flat out int oeui_drawid;
void oeui_render_drawid_vs(inout vec4 vertex) {
    oeui_drawid = gl_DrawID;
}
[break]
#version 460
#pragma vp_function oeui_render_drawid, fragment_lighting, last
flat in int oeui_drawid;
const vec3 colors[32] = { // thanks, GPT
    vec3(1.0, 0.0, 0.0),   // Red
    vec3(0.0, 1.0, 0.0),   // Green
    vec3(0.0, 0.0, 1.0),   // Blue
    vec3(1.0, 1.0, 0.0),   // Yellow
    vec3(1.0, 0.0, 1.0),   // Magenta
    vec3(0.0, 1.0, 1.0),   // Cyan
    vec3(1.0, 0.5, 0.0),   // Orange
    vec3(0.5, 1.0, 0.0),   // Lime
    vec3(0.0, 0.5, 1.0),   // Sky Blue
    vec3(0.5, 0.0, 1.0),   // Purple
    vec3(1.0, 0.5, 0.5),   // Light Red
    vec3(0.5, 1.0, 0.5),   // Light Green
    vec3(0.5, 0.5, 1.0),   // Light Blue
    vec3(1.0, 1.0, 0.5),   // Light Yellow
    vec3(1.0, 0.5, 1.0),   // Light Magenta
    vec3(0.5, 1.0, 1.0),   // Light Cyan
    vec3(0.8, 0.2, 0.2),   // Dark Red
    vec3(0.2, 0.8, 0.2),   // Dark Green
    vec3(0.2, 0.2, 0.8),   // Dark Blue
    vec3(0.8, 0.8, 0.2),   // Dark Yellow
    vec3(0.8, 0.2, 0.8),   // Dark Magenta
    vec3(0.2, 0.8, 0.8),   // Dark Cyan
    vec3(0.8, 0.5, 0.2),   // Brown
    vec3(0.5, 0.8, 0.2),   // Olive Green
    vec3(0.2, 0.5, 0.8),   // Steel Blue
    vec3(0.5, 0.2, 0.8),   // Indigo
    vec3(0.8, 0.5, 0.5),   // Salmon
    vec3(0.5, 0.8, 0.5),   // Light Olive Green
    vec3(0.5, 0.5, 0.8),   // Cornflower Blue
    vec3(0.8, 0.8, 0.5),   // Light Khaki
    vec3(0.8, 0.5, 0.8),   // Orchid
    vec3(0.5, 0.8, 0.8)    // Light Slate Gray
};
void oeui_render_drawid(inout vec4 color) {
    color.rgb = colors[oeui_drawid % 32];
}
)";

    const char* render_elevation_marker = R"(
#version 330
#pragma vp_function oeui_render_elevation_marker_vs, vertex_view
out vec4 oeui_color;
flat out int oe_terrain_vertexMarker;
void oeui_render_elevation_marker_vs(inout vec4 vertex) {
    oeui_color = vec4(0);
    int marker = oe_terrain_vertexMarker;
    if ((marker & 4) > 0) // HAS_ELEVATION
        oeui_color = vec4(1,0,0,1);
    else if ((marker & 16) > 0) // CONSTRAINT
        oeui_color = vec4(1,1,0,1);
}
[break]
#pragma vp_function oeui_render_elevation_marker_fs, fragment_lighting, last
in vec4 oeui_color;
void oeui_render_elevation_marker_fs(inout vec4 color) {
    color.rgb = mix(color.rgb, oeui_color.rgb, oeui_color.a);
}
)";
}

namespace osgEarth
{
    using namespace osgEarth::Threading;

    class RenderingGUI : public ImGuiPanel
    {
    private:
        osg::observer_ptr<MapNode> _mapNode;
        using time_point = std::chrono::time_point<std::chrono::steady_clock>;
        time_point _lastFrame;
        std::queue<int> _times;
        int _time_accum;
        int _frameCounter;
        int _fps;
        std::string _renderMode;
        bool _renderViewNormals;
        bool _renderModelNormals;
        bool _renderWinding;
        bool _renderOutlines;

    public:
        RenderingGUI() : ImGuiPanel("Rendering"),
            _frameCounter(0), _time_accum(0),
            _renderViewNormals(false), _renderModelNormals(false),
            _renderWinding(false), _renderOutlines(false) { }

        void load(const Config& conf) override
        {
        }

        void save(Config& conf) override
        {
        }

        void setRenderMode(const std::string& mode, osg::RenderInfo& ri)
        {
            auto* vp = VirtualProgram::getOrCreate(stateset(ri));
            if (!_renderMode.empty())
                ShaderLoader::unload(vp, _renderMode);
            _renderMode = mode;
            if (!_renderMode.empty())
                ShaderLoader::load(vp, _renderMode);
        }

        void draw(osg::RenderInfo& ri) override
        {
            if (!isVisible())
                return;

            if (!findNodeOrHide(_mapNode, ri))
                return;

            ImGui::Begin(name(), visible());
            {
                if (ImGuiLTable::Begin("LOD"))
                {
                    float sse = _mapNode->getScreenSpaceError();
                    if (ImGuiLTable::SliderFloat("SSE", &sse, 1.0f, 200.0f))
                        _mapNode->setScreenSpaceError(sse);

                    float lod_scale = camera(ri)->getLODScale();
                    if (ImGuiLTable::SliderFloat("LOD Scale", &lod_scale, 0.1f, 4.0f))
                        camera(ri)->setLODScale(lod_scale);

                    ImGuiLTable::End();
                }

                ImGui::Separator();

                ImGui::TextColored(ImVec4(1, 1, 0, 1), "Overlays");
                static int s_renderMode = 0;
                int m = 0;

                if (ImGui::RadioButton("Off", &s_renderMode, m++)) {
                    setRenderMode("", ri);
                }
                if (ImGui::RadioButton("Wireframe overlay", &s_renderMode, m++)) {
                    setRenderMode(render_outlines, ri);
                }
                if (ImGui::RadioButton("Front/backfacing triangles", &s_renderMode, m++)) {
                    setRenderMode(render_winding, ri);
                }
                if (ImGui::RadioButton("Normals (front/back)", &s_renderMode, m++)) {
                    setRenderMode(render_fb_normals, ri);
                }
                if (ImGui::RadioButton("Normals (view space)", &s_renderMode, m++)) {
                    setRenderMode(render_view_normals, ri);
                }
                if (ImGui::RadioButton("Normals (model space)", &s_renderMode, m++)) {
                    setRenderMode(render_model_normals, ri);
                }
                if (ImGui::RadioButton("Metal (PBR)", &s_renderMode, m++)) {
                    setRenderMode(render_metal, ri);
                }
                if (ImGui::RadioButton("Roughness (PBR)", &s_renderMode, m++)) {
                    setRenderMode(render_roughness, ri);
                }
                if (ImGui::RadioButton("AO (PBR)", &s_renderMode, m++)) {
                    setRenderMode(render_ao, ri);
                }
                if (ImGui::RadioButton("Displacement (PBR)", &s_renderMode, m++)) {
                    setRenderMode(render_displacement, ri);
                }
                if (ImGui::RadioButton("Draw ID", &s_renderMode, m++)) {
                    setRenderMode(render_draw_id, ri);
                }
                if (ImGui::RadioButton("Elevation markers", &s_renderMode, m++)) {
                    setRenderMode(render_elevation_marker, ri);
                }

                if (GLUtils::useNVGL())
                {
                    static bool s_gpuculldebug = false;
                    if (ImGui::Checkbox("GPU cull debug view", &s_gpuculldebug)) {
                        stateset(ri)->removeDefine("OE_GPUCULL_DEBUG");
                        if (s_gpuculldebug)
                            stateset(ri)->setDefine("OE_GPUCULL_DEBUG", "1");
                        else
                            stateset(ri)->setDefine("OE_GPUCULL_DEBUG", "0");
                    }
                }

                ImGui::Separator();

                const osgViewer::ViewerBase::ThreadingModel models[] = {
                    osgViewer::ViewerBase::SingleThreaded,
                    osgViewer::ViewerBase::DrawThreadPerContext,
                    osgViewer::ViewerBase::CullDrawThreadPerContext,
                    osgViewer::ViewerBase::CullThreadPerCameraDrawThreadPerContext
                };
                const std::string modelNames[] = {
                    "SingleThreaded",
                    "DrawThreadPerContext",
                    "CullDrawThreadPerContext",
                    "CullThreadPerCameraDrawThreadPerContext"
                };

                auto vb = view(ri)->getViewerBase();
                int tmi;
                for (tmi = 0; tmi < 4; ++tmi)
                    if (models[tmi] == vb->getThreadingModel())
                        break;

                ImGui::Text("OSG Threading Model: ");
                ImGui::SameLine();
                if (ImGui::Button(modelNames[tmi].c_str())) {
                    auto new_tm = models[(tmi + 1) % 4];
                    vb->addUpdateOperation(new OneTimer([vb, new_tm]() {
                        vb->setThreadingModel(new_tm); }));
                }
            }
            ImGui::End();
        }
    };


    class NVGLInspectorGUI : public ImGuiPanel
    {
    public:
        NVGLInspectorGUI() : ImGuiPanel("NVGL Inspector")
        {
        }

        void load(const Config& conf) override
        {
        }

        void save(Config& conf) override
        {
        }

        void draw(osg::RenderInfo& ri) override
        {
            if (!isVisible())
                return;

            ImGui::Begin(name(), visible());
            {
                if (!GLUtils::useNVGL())
                {
                    ImGui::TextColored(ImVec4(1, 0, 0, 1), "NVGL not enabled");
                }
                else
                {
                    auto pools = GLObjectPool::getAll();
                    for (auto& iter : pools)
                    {
                        auto cxid = iter.first;
                        auto glpool = iter.second;
                        auto globjects = glpool->objects();

                        if (pools.size() > 1)
                            ImGui::Text("Context %d", cxid);

                        double glmem = (double)glpool->totalBytes() / 1048576.;
                        ImGui::TextColored(ImVec4(1, 1, 0, 1), "NVGL Memory: %.1lf MB", glmem);
                        ImGui::SameLine();
                        static bool sort_by_size = false;
                        ImGui::Checkbox("Sort by size", &sort_by_size);

                        std::map<std::string, std::list<GLObject::Ptr>> categories;
                        for (auto& obj : globjects) {
                            categories[obj->category()].emplace_back(obj);
                        }

                        for (auto& cat : categories)
                        {
                            unsigned total = 0;
                            for (auto& obj : cat.second)
                                total += obj->size();

                            char header[128];
                            sprintf(header, "%s (%d @ %.1lf MB)###%s", cat.first.c_str(), (int)cat.second.size(), (double)total / 1048576., cat.first.c_str());

                            if (sort_by_size)
                            {
                                cat.second.sort([](const auto& lhs, const auto& rhs) {
                                    return lhs->size() > rhs->size();
                                    });
                            }

                            if (ImGui::TreeNode(header))
                            {
                                ImGuiTableFlags flags = ImGuiTableFlags_SizingFixedFit | ImGuiTableFlags_Borders;
                                flags &= ~ImGuiTableFlags_BordersOuter;
                                if (ImGui::BeginTable("globj", 4, flags))
                                {
                                    for (auto& obj : cat.second)
                                    {
                                        ImGui::TableNextColumn();
                                        ImGui::Text("%6.1lf KB", (double)obj->size() / (double)1024.);
                                        ImGui::TableNextColumn();
                                        ImGui::Text("%d", obj->recycles());
                                        ImGui::TableNextColumn();
                                        ImGui::Text("%s", obj->uid().c_str());
                                        ImGui::TableNextColumn();
                                        ImGui::Text("(%d)", obj->name());
                                    }
                                    ImGui::EndTable();
                                }

                                ImGui::TreePop();
                            }
                        }
                        if (pools.size() > 1)
                            ImGui::Separator();

                        ImGui::Separator();

                        auto recycle_attempts = glpool->recycleHits() + glpool->recycleMisses();
                        if (recycle_attempts > 0)
                        {
                            float re = (float)glpool->recycleHits() / (float)recycle_attempts;
                            ImGui::Text("Recycle Efficiency: %.0f%%", 100.0f * re);
                        }
                    }

                    ImGui::Separator();
                    int kb_per_frame = GLObjectPool::getBytesToDeletePerFrame() / 1024;
                    if (ImGuiLTable::Begin("Settings"))
                    {
                        if (ImGuiLTable::SliderInt("Rel KB per frame", &kb_per_frame, 128, 1024))
                            GLObjectPool::setBytesToDeletePerFrame(kb_per_frame * 1024);

                        ImGuiLTable::End();
                    }
                }
            }
            ImGui::End();
        }
    };
}