/* -*-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/>
 */
#ifndef OSGEARTH_IMGUI_SCENE_GRAPH_GUI
#define OSGEARTH_IMGUI_SCENE_GRAPH_GUI

#include "ImGui"
#include <osgEarth/AnnotationUtils>
#include <osg/PolygonMode>
#include <osgEarth/TerrainEngineNode>

namespace osgEarth {
    namespace GUI
    {
        using namespace osgEarth;


        class SceneGraphGUI : public BaseGUI
        {
        public:

            class SceneHierarchyVisitor : public osg::NodeVisitor
            {
            public:
                SceneHierarchyVisitor() :
                    osg::NodeVisitor(osg::NodeVisitor::TRAVERSE_ALL_CHILDREN)
                {
                    setNodeMaskOverride(~0);
                }

                void apply(osg::Node& node)
                {
                    // Non groups act as leaf nodes
                    std::string label = getLabel(node);
                    ImGuiTreeNodeFlags node_flags = base_flags;
                    node_flags |= ImGuiTreeNodeFlags_Leaf | ImGuiTreeNodeFlags_NoTreePushOnOpen;
                    if (getSelectedNode() == &node)
                    {
                        node_flags |= ImGuiTreeNodeFlags_Selected;
                    }

                    ImGui::TreeNodeEx(&node, node_flags, label.c_str());
                    if (ImGui::IsItemClicked())
                    {                                         
                        setSelectedNodePath(getNodePath());
                    }
                }

                void apply(osg::Group& node)
                {
                    std::string label = getLabel(node);
                    ImGuiTreeNodeFlags node_flags = base_flags;
                    if (getSelectedNode() == &node)
                    {
                        node_flags |= ImGuiTreeNodeFlags_Selected;
                    }

                    if (isInSelectedNodePath(&node))
                    {
                        ImGui::SetNextItemOpen(true, ImGuiCond_Once);
                    }

                    bool node_open = ImGui::TreeNodeEx(&node, node_flags, label.c_str());
                    if (ImGui::IsItemClicked())
                    {                     
                        setSelectedNodePath(getNodePath());
                    }
                    if (node_open)
                    {
                        traverse(node);
                        ImGui::TreePop();
                    }
                }

                std::string getLabel(osg::Node& node)
                {
                    std::stringstream buf;
                    buf << node.getName() << " (" << typeid(node).name() << ")";

                    return buf.str();
                }

                bool isInSelectedNodePath(osg::Node* node)
                {
                    auto selectedNodePaths = getSelectedNodePath();                    

                    for (unsigned int i = 0; i < selectedNodePaths.size(); i++)
                    {
                        if (node == selectedNodePaths[i].get())
                        {
                            return true;
                        }
                    }
                    return false;
                }

                ImGuiTreeNodeFlags base_flags = ImGuiTreeNodeFlags_OpenOnArrow | ImGuiTreeNodeFlags_OpenOnDoubleClick | ImGuiTreeNodeFlags_SpanAvailWidth;
            };

            SceneGraphGUI() :
                BaseGUI("Scene Graph Inspector"),
                _installedSelectNodeHandler(false),
                _selectedBounds(0),
                _node(nullptr)
            {
            }

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

                if (ImGui::Begin(name(), visible()))
                {
                    if (_node == nullptr)
                        _node = ri.getCurrentCamera();

                    if (!_mapNode.valid())
                        _mapNode = osgEarth::findTopMostNodeOfType<MapNode>(ri.getCurrentCamera());

                    if (!_installedSelectNodeHandler)
                    {
                        auto view = dynamic_cast<osgViewer::View*>(ri.getView());
                        view->addEventHandler(new SelectNodeHandler());
                        _installedSelectNodeHandler = true;
                    }

                    auto size = ImGui::GetWindowSize();
                    
                    if (ImGui::CollapsingHeader("Scene"))
                    {
                        ImGui::BeginChild("Scene", ImVec2(0, 0.6f * size.y));
                        SceneHierarchyVisitor v;
                        _node->accept(v);

                        ImGui::EndChild();
                    }

                    if (getSelectedNode())
                    {
                        if (ImGui::CollapsingHeader("Properties"))
                        {
                            ImGui::BeginChild("Properties");
                            auto view = dynamic_cast<osgViewer::View*>(ri.getView());
                            auto manip = dynamic_cast<EarthManipulator*>(view->getCameraManipulator());
                            properties(getSelectedNode(), ri, manip, _mapNode.get());
                            ImGui::EndChild();
                        }
                    }
                    ImGui::End();
                }
                // If they closed the tool deselect the node.
                if (!isVisible() && getSelectedNode())
                {
                    // TODO:
                    //setSelectedNode(nullptr);
                    osg::NodePath empty;
                    setSelectedNodePath(empty);
                }

                updateBoundsDebug(_mapNode.get());
            }

            void updateBoundsDebug(MapNode* mapNode)
            {
                if (_selectedBounds)
                {
                    mapNode->removeChild(_selectedBounds);
                    _selectedBounds = nullptr;
                }

                auto selectedNode = getSelectedNode();

                if (selectedNode)
                {
                    osg::RefNodePath refNodePath = getSelectedNodePath();
                    osg::NodePath nodePath;
                    for (unsigned int i = 0; i < refNodePath.size(); ++i)
                    {
                        nodePath.push_back(refNodePath[i].get());
                    }
                    
                    if (nodePath.back()->asTransform())
                    {
                        nodePath.pop_back();
                    }

                    osg::Matrixd localToWorld = osg::computeLocalToWorld(nodePath);

                    osg::Drawable* drawable = selectedNode->asDrawable();
                    if (drawable)
                    {
                        osg::BoundingBoxd bb = drawable->getBoundingBox();
                        if (bb.valid())
                        {
                            _selectedBounds = new osg::MatrixTransform;
                            _selectedBounds->setName("Bounds");
                            _selectedBounds->setMatrix(localToWorld);

                            osg::MatrixTransform* center = new osg::MatrixTransform;
                            center->addChild(AnnotationUtils::createBoundingBox(bb, Color::Yellow));
                            _selectedBounds->addChild(center);                           

                            _selectedBounds->setNodeMask(1 << 2);
                            mapNode->addChild(_selectedBounds);
                        }
                    }
                    else
                    {
                        osg::BoundingSphered bs = osg::BoundingSphered(selectedNode->getBound().center(), selectedNode->getBound().radius());
                        if (bs.radius() > 0)
                        {
                            _selectedBounds = new osg::MatrixTransform;
                            _selectedBounds->setName("Bounds");
                            _selectedBounds->setMatrix(localToWorld);

                            osg::MatrixTransform* center = new osg::MatrixTransform;
                            center->setMatrix(osg::Matrix::translate(bs.center()));
                            center->addChild(AnnotationUtils::createSphere(selectedNode->getBound().radius(), Color::Yellow));
                            _selectedBounds->addChild(center);

                            _selectedBounds->getOrCreateStateSet()->setAttributeAndModes(new osg::PolygonMode(osg::PolygonMode::FRONT_AND_BACK, osg::PolygonMode::LINE), 1);
                            _selectedBounds->setNodeMask(1 << 2);
                            mapNode->addChild(_selectedBounds);
                        }
                    }
                }
            }


            void properties(osg::Node* node, osg::RenderInfo& ri, EarthManipulator* manip, MapNode* mapNode)
            {
                if (ImGui::Button("Zoom to"))
                {
                    osg::NodePath nodePath = node->getParentalNodePaths()[0];
                    osg::Matrixd localToWorld = osg::computeLocalToWorld(nodePath);

                    osg::BoundingSphered bs(node->getBound().center(), node->getBound().radius());
                    if (bs.valid())
                    {
                        bs.center() += localToWorld.getTrans();
                        osg::Vec3d c = bs.center();
                        double r = bs.radius();
                        const SpatialReference* mapSRS = mapNode->getMap()->getSRS();

                        std::vector<GeoPoint> points;
                        GeoPoint p;
                        p.fromWorld(mapSRS, osg::Vec3d(c.x() + r, c.y(), c.z())); points.push_back(p);
                        p.fromWorld(mapSRS, osg::Vec3d(c.x() - r, c.y(), c.z())); points.push_back(p);
                        p.fromWorld(mapSRS, osg::Vec3d(c.x(), c.y() + r, c.z())); points.push_back(p);
                        p.fromWorld(mapSRS, osg::Vec3d(c.x(), c.y() - r, c.z())); points.push_back(p);
                        p.fromWorld(mapSRS, osg::Vec3d(c.x(), c.y(), c.z() + r)); points.push_back(p);
                        p.fromWorld(mapSRS, osg::Vec3d(c.x(), c.y(), c.z() - r)); points.push_back(p);

                        ViewFitter fitter(mapNode->getMap()->getSRS(), ri.getCurrentCamera());
                        Viewpoint vp;
                        if (fitter.createViewpoint(points, vp))
                        {
                            manip->setViewpoint(vp, 2.0);
                        }
                    }
                }

                bool visible = node->getNodeMask() != 0;
                if (ImGui::Checkbox("Visible", &visible))
                {
                    node->setNodeMask(visible);
                }

                if (ImGui::TreeNode("General"))
                {
                    ImGui::Text("Type %s", typeid(*node).name());
                    ImGui::Text("Name %s", node->getName().c_str());

                    osg::Drawable* drawable = node->asDrawable();
                    if (drawable)
                    {                        
                        osg::BoundingBox bbox = drawable->getBoundingBox();
                        double width = bbox.xMax() - bbox.xMin();
                        double depth = bbox.yMax() - bbox.yMin();
                        double height = bbox.zMax() - bbox.zMin();

                        ImGui::Text("Bounding box (local) center=(%.3f, %.3f, %.3f) dimensions=%.3f x %.3f %.3f", bbox.center().x(), bbox.center().y(), bbox.center().z(), width, depth, height);
                    }
                    else
                    {
                        osg::BoundingSphere bs = node->getBound();
                        ImGui::Text("Bounding sphere (local) center=(%.3f, %.3f, %.3f) radius=%.3f", bs.center().x(), bs.center().y(), bs.center().z(), bs.radius());
                    }

                    ImGui::TreePop();
                }

                osg::MatrixTransform* matrixTransform = dynamic_cast<osg::MatrixTransform*>(node);
                if (matrixTransform)
                {
                    if (ImGui::TreeNode("Transform"))
                    {
                        bool dirty = false;
                        osg::Matrix matrix = matrixTransform->getMatrix();
                        osg::Vec3d translate = matrixTransform->getMatrix().getTrans();
                        if (ImGui::InputScalarN("Translation", ImGuiDataType_Double, translate._v, 3))
                        {
                            dirty = true;
                        }

                        osg::Vec3d scale = matrixTransform->getMatrix().getScale();
                        if (ImGui::InputScalarN("Scale", ImGuiDataType_Double, scale._v, 3))
                        {
                            dirty = true;
                        }
                        osg::Quat rot = matrixTransform->getMatrix().getRotate();
                        if (ImGui::InputScalarN("Rotation", ImGuiDataType_Double, rot._v, 4))
                        {
                            dirty = true;
                        }

                        if (dirty)
                        {
                            osg::Matrix newMatrix = osg::Matrix::translate(translate) * osg::Matrix::rotate(rot) * osg::Matrixd::scale(scale);
                            matrixTransform->setMatrix(newMatrix);
                        }
                        ImGui::TreePop();
                    }
                }

                osg::StateSet* stateset = node->getStateSet();
                if (stateset)
                {
                    if (ImGui::TreeNode("State Set"))
                    {
                        if (!stateset->getTextureAttributeList().empty() && ImGui::TreeNode("Textures"))
                        {
                            unsigned int textureWidth = 50;
                            for (unsigned int i = 0; i < stateset->getTextureAttributeList().size(); ++i)
                            {
                                osg::Texture2D* texture = dynamic_cast<osg::Texture2D*>(stateset->getTextureAttribute(i, osg::StateAttribute::TEXTURE));

                                if (texture)
                                {
                                    ImGui::PushID(texture);
                                    ImGui::BeginGroup();
                                    ImGui::Text("Unit %d", i);
                                    ImGui::Text("Name %s", texture->getName().c_str());
                                    ImGuiUtil::Texture(texture, ri, textureWidth);
                                    ImGui::EndGroup();
                                }
                            }
                            ImGui::TreePop();
                        }

                        VirtualProgram* virtualProgram = VirtualProgram::get(stateset);
                        if (virtualProgram && ImGui::TreeNode("Shaders"))
                        {
                            VirtualProgram::ShaderMap shaderMap;
                            virtualProgram->getShaderMap(shaderMap);
                            for (auto& s : shaderMap)
                            {
                                ImGui::Text("Name %s", s.second._shader->getName().c_str());
                                ImGui::Text("Location %s", FunctionLocationToString(s.second._shader->getLocation()));
                                ImGui::TextWrapped(s.second._shader->getShaderSource().c_str());
                                ImGui::Separator();
                            }
                            ImGui::TreePop();
                        }

                        ImGui::TreePop();
                    }
                }

                osg::Geometry* geometry = dynamic_cast<osg::Geometry*>(node);
                if (geometry && ImGui::TreeNode("Geometry Properties"))
                {
                    if (geometry->getVertexArray())
                    {
                        ImGui::Text("Vertices %d", geometry->getVertexArray()->getNumElements());
                    }

                    if (geometry->getColorArray())
                    {
                        ImGui::Text("Colors %d", geometry->getColorArray()->getNumElements());
                    }

                    if (geometry->getNormalArray())
                    {
                        ImGui::Text("Normals %d", geometry->getNormalArray()->getNumElements());
                    }

                    for (auto &p : geometry->getPrimitiveSetList())
                    {
                        ImGui::Text("%s Mode=%s Primitives=%d Instances=%d", typeid(*p.get()).name(), GLModeToString(p->getMode()), p->getNumPrimitives(), p->getNumInstances());
                    }
                }
            }

            const char* GLModeToString(GLenum mode)
            {
                switch (mode)
                {
                case(GL_POINTS): return "GL_POINTS";
                case(GL_LINES): return "GL_LINES";
                case(GL_TRIANGLES): return "GL_TRIANGLES";
                case(GL_QUADS): return "GL_QUADS";
                case(GL_LINE_STRIP): return "GL_LINE_STRIP";
                case(GL_LINE_LOOP): return "GL_LINE_LOOP";
                case(GL_TRIANGLE_STRIP): return "GL_TRIANGLE_STRIP";
                case(GL_TRIANGLE_FAN): return "GL_TRIANGLE_FAN";
                case(GL_QUAD_STRIP): return "GL_QUAD_STRIP";
                case(GL_PATCHES): return "GL_PATCHES";
                case(GL_POLYGON): return "GL_POLYGON";
                default: return "Unknown";
                }
            }

            const char* FunctionLocationToString(ShaderComp::FunctionLocation location)
            {
                using namespace osgEarth::ShaderComp;

                switch (location)
                {
                case LOCATION_VERTEX_MODEL: return "LOCATION_VERTEX_MODEL";
                case LOCATION_VERTEX_VIEW: return "LOCATION_VERTEX_VIEW";
                case LOCATION_VERTEX_CLIP: return "LOCATION_VERTEX_CLIP";
                case LOCATION_TESS_CONTROL: return "LOCATION_TESS_CONTROL";
                case LOCATION_TESS_EVALUATION: return "LOCATION_TESS_EVALUATION";
                case LOCATION_GEOMETRY: return "LOCATION_GEOMETRY";
                case LOCATION_FRAGMENT_COLORING: return "LOCATION_FRAGMENT_COLORING";
                case LOCATION_FRAGMENT_LIGHTING: return "LOCATION_FRAGMENT_LIGHTING";
                case LOCATION_FRAGMENT_OUTPUT: return "LOCATION_FRAGMENT_OUTPUT";
                case LOCATION_UNDEFINED: return "LOCATION_UNDEFINED";
                default: return "Undefined";
                }
            }

            bool _installedSelectNodeHandler;
            osg::MatrixTransform* _selectedBounds;
            osg::observer_ptr<osg::Node> _node;
            osg::observer_ptr<MapNode> _mapNode;
        };
    }
}

#endif // OSGEARTH_IMGUI_SCENE_GRAPH_GUI
