class EntityNodeUI {
    Entity entity;
    array<EntityNodeUI@> children;
    TreePanel@ rootPanel;

    string displayName;
    vec3 networkColor;
    EntityNetworkBehaviour initialNetworkBehaviour;

    EntityNodeUI(Entity _entity, TreePanel@ _root) {
        entity = _entity;
        @rootPanel = _root;

        updateInfo();
        buildChildren();
    }

    // Updates cached info
    void updateInfo() {
        displayName = entity.name;
        if (displayName == "") displayName = "-"; else displayName = "\uf1b2 " + entity.name;
        initialNetworkBehaviour = entity.networkBehaviour;

        if (!entity.isValidNetworkBehaviour()) {
            networkColor = vec3(1, 0.4f, 0.4f);
        } else {
            EntityNetworkBehaviour net = entity.getForcedNetworkBehaviour();
            if (net == EntityNetworkBehaviour::Client) networkColor = vec3(0.6, 1, 0.7); else if (net == EntityNetworkBehaviour::Server) networkColor = vec3(0.6, 0.7, 1); else networkColor = vec3(1, 1, 1);
        }
    }

    // Build persistent children nodes
    void buildChildren() {
        children.resize(0);
        array<Entity>@ childEntities = entity.getChildrens();
        for (uint i = 0; i < childEntities.length(); i++) {
            if (rootPanel.networkBehaviourFilter != EntityNetworkBehaviour::Parent) {
                EntityNetworkBehaviour entBeh = childEntities[i].getForcedNetworkBehaviour();
                if (entBeh != EntityNetworkBehaviour::Parent && entBeh != rootPanel.networkBehaviourFilter) 
                    continue;
            }

            EntityNodeUI@ child = @EntityNodeUI(childEntities[i], rootPanel);
            children.insertLast(child);
        }
    }

    bool isActive() {
        return entity == ActiveEntity::getActiveEntity();
    }

    // Call each frame to render the node tree
    void render() {
        bool opened = false;
        if (children.length() == 0) {
            ImGui::treeNodeLeaf("##" + entity.id, isActive());
            renderInteraction();
        } else {
            opened = ImGui::treeNode("##" + entity.id, isActive(), SimpleFunction(this.renderChildren));
            if (!opened) renderInteraction();
        }
    }

    void renderChildren() {
        renderInteraction();
        if (!entity.exists) return;
        for (uint i = 0; i < children.length(); i++) {
            if (children[i].entity.exists) children[i].render();
        }
    }

    void renderInteraction() {
        // Drag and drop
        ImGui::dragDropSource("ENTITY", any(entity), entity.name);
        ImGui::dragDropTarget("ENTITY", ReciverFunction(this.entityDrop));

        // Selection
        if (ImGui::isItemClicked(0))
            ActiveEntity::setActiveEntity(entity);

        // Context menu
        ImGui::contextItemPopup("POP_ENTITY_" + entity.id, SimpleFunction(this.renderContextMenu));
        ImGui::simplePopup("RENAME_ENTITY_" + entity.id, SimpleFunction(this.renameEntity));
        
        if (!entity.exists) return;

        // Network color and name
        ImGui::sameline();
        ImGui::textColor(networkColor.x, networkColor.y, networkColor.z, displayName);

        if (initialNetworkBehaviour != entity.networkBehaviour) rootPanel.onInit();

        if (initialNetworkBehaviour != EntityNetworkBehaviour::Parent) {
            ImGui::sameline();
            if (initialNetworkBehaviour != EntityNetworkBehaviour::Server)
                ImGui::textEnd("\uf233 Client "); else
                ImGui::textEnd("\uf233 Server ");
        }

    }

    void entityDrop(any@ data) {
        Entity dropped;
        data.retrieve(dropped);
        if (dropped.isRoot || entity.isDescendantOf(dropped)) return;
        dropped.parent = entity;
        rootPanel.onInit();
    }

    void renderContextMenu() {
        if (entity.isRoot) {
            if (ImGui::menuItem("New Entity")) {
                entity.createChild("node");
                rootPanel.onInit();
            }

            ImGui::separator();
            ImGui::subMenu("3d Object", SimpleFunction(this.render3dObjectMenu));
        } else {
            if (ImGui::menuItem("Add child")) {
                entity.createChild("node");
                rootPanel.onInit();
            }
            if (ImGui::menuItem("Rename")) {
                ImGui::openPopup("RENAME_ENTITY_" + entity.id, any(entity));
                rootPanel.onInit();
            }
            if (ImGui::menuItem("Destroy")) {
                if (ActiveEntity::entity == entity) ActiveEntity::setActiveEntity(World::getRoot());

                entity.destroy();
                rootPanel.onInit();
            }

            ImGui::separator();
            ImGui::subMenu("3d Object", SimpleFunction(this.render3dObjectMenu));
        }
    }

    void render3dObjectMenu() {
        if (ImGui::menuItem("Cube")) {
            Entity child = entity.createChild("node");
            MeshComponent childMesh = child.addComponent<MeshComponent>();
            childMesh.meshResource = Builtin::cube();

            ShaderComponent childShader = child.addComponent<ShaderComponent>();
            childShader.shader = Builtin::simpleShader();

            rootPanel.onInit();
        }

        if (ImGui::menuItem("Sphere")) {
            Entity child = entity.createChild("node");
            MeshComponent childMesh = child.addComponent<MeshComponent>();
            childMesh.meshResource = Builtin::sphere();

            ShaderComponent childShader = child.addComponent<ShaderComponent>();
            childShader.shader = Builtin::simpleShader();

            rootPanel.onInit();
        }
    }

    void renameEntity() {
        string outName;
        if (ImGui::inputText("Entity name ", entity.name, outName)) entity.name = outName;
        rootPanel.onInit();
    }

    // Call when hierarchy changes
    void refresh() {
        updateInfo();
        buildChildren();
    }
}

class TreePanel : Panel {
    EntityNodeUI@ rootNode;
    bool infoChanged = false;
    EntityNetworkBehaviour networkBehaviourFilter = EntityNetworkBehaviour::Parent;

    void onInit() {
        infoChanged = true;
    }

    void onMenuBar() {
        if (ImGui::menuItem("Refresh")) {
            infoChanged = true;
            networkBehaviourFilter = EntityNetworkBehaviour::Parent;
        }

        string networkBehaviourFilterName;
        if (networkBehaviourFilter == EntityNetworkBehaviour::Parent) networkBehaviourFilterName = "Server&Client"; else if (networkBehaviourFilter == EntityNetworkBehaviour::Server) networkBehaviourFilterName = "Server"; else networkBehaviourFilterName = "Client";
        ImGui::drawComboEnd("##Filter", networkBehaviourFilterName, SimpleFunction(this.networkFilter));
    }

    void networkFilter() {
        if (ImGui::drawComboItem("Server&Client", networkBehaviourFilter == EntityNetworkBehaviour::Parent)){
            networkBehaviourFilter = EntityNetworkBehaviour::Parent;
            infoChanged = true;
        }
        if (ImGui::drawComboItem("Server", networkBehaviourFilter == EntityNetworkBehaviour::Server)){
            networkBehaviourFilter = EntityNetworkBehaviour::Server;
            infoChanged = true;
        }
        if (ImGui::drawComboItem("Client", networkBehaviourFilter == EntityNetworkBehaviour::Client)){
            networkBehaviourFilter = EntityNetworkBehaviour::Client;
            infoChanged = true;
        }
    }

    void onImGui() {
        if (infoChanged) {
            Entity root = World::getRoot();
            @rootNode = @EntityNodeUI(root, this);

            infoChanged = false;
        }

        ImGui::contextMenuPopup("POP_ROOT", SimpleFunction(rootNode.renderContextMenu));

        for (uint i = 0; i < rootNode.children.length(); i++) {
            if (rootNode.children[i].entity.exists) rootNode.children[i].render();
        }
    }
}