/* osgEarth
 * Copyright 2025 Pelican Mapping
 * MIT License
 */
#pragma once

#include <osgEarthImGui/ImGuiPanel>
#include <osgEarth/Registry>
#include <osgEarth/VirtualProgram>
#include <osg/Uniform>
#include <osg/StateSet>
#include <chrono>
#include <list>
#include <vector>

namespace osgEarth
{
    using namespace osgEarth::Threading;
    using namespace osgEarth::Util;

    class ShaderGUI : public ImGuiPanel
    {
    private:
        struct UniformSpec {
            std::string _name;
            float _minval, _maxval, _value;
            osg::ref_ptr<osg::Uniform> _u;
        };
        std::vector<UniformSpec> _uniforms;

        struct DefineSpec {
            std::string _name;
            bool _checked;
        };
        std::vector<DefineSpec> _defines;

        ProgramRepo::ProgramMap _db;

    public:
        ShaderGUI(osg::ArgumentParser* args = nullptr) : ImGuiPanel("Shaders")
        {
            if (args)
            {
                while (args->find("--uniform") >= 0)
                {
                    UniformSpec spec;
                    if (args->read("--uniform", spec._name, spec._minval, spec._maxval))
                    {
                        spec._value = clamp(0.0f, spec._minval, spec._maxval);
                        spec._u = new osg::Uniform(spec._name.c_str(), spec._value);
                        _uniforms.emplace_back(spec);
                    }
                }

                while (args->find("--define") >= 0)
                {
                    DefineSpec spec;
                    if (args->read("--define", spec._name))
                    {
                        spec._checked = false;
                        _defines.push_back(spec);
                    }
                }
            }
        }

        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 (!_uniforms.empty())
                    ImGui::Text("Uniforms:");

                for (auto& def : _uniforms)
                {
                    if (ImGui::SliderFloat(def._name.c_str(), &def._value, def._minval, def._maxval))
                    {
                        def._u->set(def._value);
                        ri.getCurrentCamera()->getOrCreateStateSet()->addUniform(def._u.get(), osg::StateAttribute::OVERRIDE);
                    }
                }

                if (!_defines.empty())
                    ImGui::Text("Defines:");

                for (auto& def : _defines)
                {
                    if (ImGui::Checkbox(def._name.c_str(), &def._checked))
                    {
                        if (def._checked)
                            ri.getCurrentCamera()->getOrCreateStateSet()->setDefine(def._name, osg::StateAttribute::ON | osg::StateAttribute::OVERRIDE);
                        else
                            ri.getCurrentCamera()->getOrCreateStateSet()->setDefine(def._name, osg::StateAttribute::OFF | osg::StateAttribute::OVERRIDE);
                    }
                }

                ImGui::Separator();
                ImGui::TextColored(ImVec4(1, 1, 0, 1), "Active Programs:");

                if (_db.empty() || ImGui::Button("Refresh"))
                {
                    _db = Registry::programRepo().copy();
                }

                static void* open_shader = nullptr;
                int count = 0;
                int frame = ri.getState()->getFrameStamp()->getFrameNumber();
                for (auto& entry : _db)
                {
                    if (entry.second->_users.size() > 0 && (frame - entry.second->_frameLastUsed < 60))
                    {
                        auto program = entry.second->_program;
                        auto name = program->getName();
                        if (name.empty()) name = "[no name]";
                        std::string tag = name + "##" + std::to_string(count++);
                        if (ImGui::TreeNode(tag.c_str())) //name.c_str()))
                        {
                            auto num = program->getNumShaders();
                            for (unsigned i = 0; i < num; ++i)
                            {
                                auto shader = program->getShader(i);
                                auto shader_name = shader->getName();
                                if (shader_name.empty()) shader_name = "[no name]";

                                ImGui::PushID(count++);
                                std::string shader_tag = shader_name + "##" + std::to_string(count++);

                                if (ImGui::Button(shader_tag.c_str()))
                                {
                                    open_shader = shader;
                                }

                                if (open_shader == shader)
                                {
                                    ImGui::SetNextWindowSize(ImVec2(800, 600), ImGuiCond_FirstUseEver);
                                    ImGui::Begin(shader_name.c_str());

                                    if (ImGui::Button("Compile"))
                                        program->releaseGLObjects(nullptr);

                                    ImGui::SameLine();
                                    if (ImGui::Button("Close"))
                                        open_shader = nullptr;

                                    auto text = shader->getShaderSource();
                                    if (ImGuiEx::InputTextMultiline("##shader_source", &text, ImVec2(-1, -1)))
                                    {
                                        shader->setShaderSource(text);
                                        program->releaseGLObjects(nullptr);
                                    }

                                    ImGui::End();
                                }
                                ImGui::PopID();
                            }
                            ImGui::TreePop();
                        }
                    }
                }
            }
            ImGui::End();
        }
    };
}
