// World Graph Node Editor
// A standalone Dear ImGui + imnodes tool for editing WorldGraph node graphs.
// Graphs are saved/loaded as Flecs JSON (graph.json in the working directory).

#include <GLFW/glfw3.h>
#include "imgui.h"
#include "imgui_impl_glfw.h"
#include "imgui_impl_opengl3.h"
#include "imnodes.h"

#include "flecs.h"
#include "Components/WorldGraph.hpp"
#include "Types/WorldGraph/WorldGraph.h"

#include <cstdio>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <unordered_map>
#include <algorithm>

// ─── Editor state ─────────────────────────────────────────────────────────────

struct EditorNode
{
    int           nodeId;       // imnodes node ID
    flecs::entity entity;
    VisualNodeKind kind;
    bool          isOutput;
};

struct EditorLink
{
    int linkId;
    int srcAttrId;  // output attribute of source node
    int dstAttrId;  // input attribute (pin slot) of destination node
};

// Attribute ID encoding:
//   node's output pin  = nodeId * 10 + 9
//   node's input pin N = nodeId * 10 + N   (N = 0,1,2)

static int OutputAttr(int nodeId) { return nodeId * 10 + 9; }
static int InputAttr (int nodeId, int slot) { return nodeId * 10 + slot; }
static int AttrToNode(int attr) { return attr / 10; }
static int AttrToSlot(int attr) { return attr % 10; }  // 9 = output

// ─── Kind metadata ────────────────────────────────────────────────────────────

struct KindInfo { const char* name; int inputCount; bool hasParamConstant; bool hasParamNoise; bool hasParamIsTile; bool hasParamTileDistance; };

static const KindInfo KindInfoTable[] = {
    // name              ins  const noise isTile tileDist
    {"Add",              2,   0, 0, 0, 0},
    {"Subtract",         2,   0, 0, 0, 0},
    {"Multiply",         2,   0, 0, 0, 0},
    {"Divide",           2,   0, 0, 0, 0},
    {"Modulo",           2,   0, 0, 0, 0},
    {"Pow",              2,   0, 0, 0, 0},
    {"Max",              2,   0, 0, 0, 0},
    {"Min",              2,   0, 0, 0, 0},
    {"Negate",           1,   0, 0, 0, 0},
    {"Abs",              1,   0, 0, 0, 0},
    {"Ceil",             1,   0, 0, 0, 0},
    {"Floor",            1,   0, 0, 0, 0},
    {"Sin",              1,   0, 0, 0, 0},
    {"Cos",              1,   0, 0, 0, 0},
    {"Tan",              1,   0, 0, 0, 0},
    {"Exp",              1,   0, 0, 0, 0},
    {"Log",              1,   0, 0, 0, 0},
    {"Square",           1,   0, 0, 0, 0},
    {"Round",            1,   0, 0, 0, 0},
    {"OneMinus",         1,   0, 0, 0, 0},
    {"Lerp",             3,   0, 0, 0, 0},
    {"Clamp",            3,   0, 0, 0, 0},
    {"Equal",            2,   0, 0, 0, 0},
    {"Smaller",          2,   0, 0, 0, 0},
    {"Greater",          2,   0, 0, 0, 0},
    {"SmallerEqual",     2,   0, 0, 0, 0},
    {"GreaterEqual",     2,   0, 0, 0, 0},
    {"And",              2,   0, 0, 0, 0},
    {"Or",               2,   0, 0, 0, 0},
    {"Branch",           3,   0, 0, 0, 0},
    {"Simplex",          0,   0, 1, 0, 0},
    {"OpenSimplex",      0,   0, 1, 0, 0},
    {"Perlin",           0,   0, 1, 0, 0},
    {"Value",            0,   0, 1, 0, 0},
    {"ValueCubic",       0,   0, 1, 0, 0},
    {"Constant",         0,   1, 0, 0, 0},
    {"IsTile",           0,   0, 0, 1, 0},
    {"TileDistance",     0,   0, 0, 0, 1},
};

static const KindInfo& GetKindInfo(VisualNodeKind k)
{
    return KindInfoTable[static_cast<int>(k)];
}

// ─── Application ─────────────────────────────────────────────────────────────

class NodeEditorApp
{
public:
    NodeEditorApp()
    {
        Flecs_WorldGraph(World);
        GraphRoot = World.entity("Graph");
    }

    // ── Graph I/O ─────────────────────────────────────────────────────────────

    void Save(const char* path)
    {
        // Sync editor positions back to components before saving
        for (auto& n : Nodes)
        {
            ImVec2 pos = ImNodes::GetNodeGridSpacePos(n.nodeId);
            n.entity.set<VisualNodePos>({pos.x, pos.y});
        }
        char* json = ecs_world_to_json(World.c_ptr(), nullptr);
        std::ofstream f(path);
        f << json;
        ecs_os_free(json);
        printf("Saved to %s\n", path);
    }

    void Load(const char* path)
    {
        std::ifstream f(path);
        if (!f) { printf("Could not open %s\n", path); return; }
        std::ostringstream ss; ss << f.rdbuf();
        std::string json = ss.str();
        ecs_world_from_json(World.c_ptr(), json.c_str(), nullptr);
        RebuildEditorState();
        printf("Loaded from %s\n", path);
    }

    // ── Rebuild editor state from Flecs world ─────────────────────────────────

    void RebuildEditorState()
    {
        Nodes.clear();
        Links.clear();
        NextId = 1;

        // Map entity → nodeId
        std::unordered_map<uint64_t, int> entityToNode;

        World.query_builder<VisualNodeType>()
            .with(flecs::ChildOf, GraphRoot)
            .build()
            .each([&](flecs::entity e, VisualNodeType& t)
            {
                int nid = NextId++;
                entityToNode[e.id()] = nid;
                Nodes.push_back({ nid, e, t.kind, e.has<VisualNodeOutput>() });

                // Restore canvas position
                const VisualNodePos* pos = e.try_get<VisualNodePos>();
                if (pos) ImNodes::SetNodeGridSpacePos(nid, ImVec2{pos->x, pos->y});
            });

        // Rebuild links
        for (auto& n : Nodes)
        {
            auto addLink = [&](int slot, flecs::entity src)
            {
                if (!src) return;
                auto it = entityToNode.find(src.id());
                if (it == entityToNode.end()) return;
                int srcNodeId = it->second;
                Links.push_back({ NextId++,
                    OutputAttr(srcNodeId),
                    InputAttr(n.nodeId, slot) });
            };
            addLink(0, n.entity.target<InputPin0>());
            addLink(1, n.entity.target<InputPin1>());
            addLink(2, n.entity.target<InputPin2>());
        }
    }

    // ── Add a new node ────────────────────────────────────────────────────────

    void AddNode(VisualNodeKind kind)
    {
        static int nameCounter = 0;
        char name[64];
        snprintf(name, sizeof(name), "%s_%d", GetKindInfo(kind).name, nameCounter++);

        flecs::entity e = World.entity(name)
            .child_of(GraphRoot)
            .set<VisualNodeType>({ kind })
            .set<VisualNodePos>({ 100.f, 100.f });

        // Add default parameters
        if (GetKindInfo(kind).hasParamConstant)    e.set<NodeParam_Constant>({ 0.f });
        if (GetKindInfo(kind).hasParamNoise)       e.set<NodeParam_Noise>({ 0.01f });
        if (GetKindInfo(kind).hasParamIsTile)      e.set<NodeParam_IsTile>({ 0, 0, TileType::Air });
        if (GetKindInfo(kind).hasParamTileDistance)e.set<NodeParam_TileDistance>({ 1, TileType::Air });

        int nid = NextId++;
        Nodes.push_back({ nid, e, kind, false });
        ImNodes::SetNodeGridSpacePos(nid, { 100.f, 100.f });
    }

    // ── Render one frame ──────────────────────────────────────────────────────

    void Render()
    {
        // ── Menu bar ─────────────────────────────────────────────────────────
        if (ImGui::BeginMainMenuBar())
        {
            if (ImGui::BeginMenu("File"))
            {
                if (ImGui::MenuItem("Save", "Ctrl+S")) Save("graph.json");
                if (ImGui::MenuItem("Load", "Ctrl+O")) Load("graph.json");
                ImGui::EndMenu();
            }
            if (ImGui::BeginMenu("Add Node"))
            {
                for (int i = 0; i <= (int)VisualNodeKind::TileDistance; ++i)
                {
                    if (ImGui::MenuItem(KindInfoTable[i].name))
                        AddNode(static_cast<VisualNodeKind>(i));
                }
                ImGui::EndMenu();
            }
            ImGui::EndMainMenuBar();
        }

        // ── Full-screen canvas ────────────────────────────────────────────────
        ImGuiIO& io = ImGui::GetIO();
        ImGui::SetNextWindowPos({0, ImGui::GetFrameHeight()});
        ImGui::SetNextWindowSize({io.DisplaySize.x, io.DisplaySize.y - ImGui::GetFrameHeight()});
        ImGui::Begin("##canvas", nullptr,
            ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoResize |
            ImGuiWindowFlags_NoScrollbar | ImGuiWindowFlags_NoMove);

        ImNodes::BeginNodeEditor();

        // ── Draw nodes ────────────────────────────────────────────────────────
        for (auto& n : Nodes)
        {
            const KindInfo& info = GetKindInfo(n.kind);

            if (n.isOutput)
                ImNodes::PushColorStyle(ImNodesCol_TitleBar, IM_COL32(150, 80, 20, 255));

            ImNodes::BeginNode(n.nodeId);

            ImNodes::BeginNodeTitleBar();
            ImGui::TextUnformatted(info.name);
            if (n.isOutput) ImGui::SameLine(); if (n.isOutput) ImGui::TextUnformatted(" [OUT]");
            ImNodes::EndNodeTitleBar();

            // Input pins
            static const char* slotLabels[] = { "In 0", "In 1", "In 2" };
            for (int s = 0; s < info.inputCount; ++s)
            {
                ImNodes::BeginInputAttribute(InputAttr(n.nodeId, s));
                ImGui::TextUnformatted(slotLabels[s]);
                ImNodes::EndInputAttribute();
            }

            // Inline parameters
            ImGui::PushItemWidth(120.f);
            if (info.hasParamConstant)
            {
                auto* p = n.entity.try_get_mut<NodeParam_Constant>();
                ImGui::DragFloat("##val", &p->value, 0.01f);
            }
            if (info.hasParamNoise)
            {
                auto* p = n.entity.try_get_mut<NodeParam_Noise>();
                ImGui::DragFloat("Freq##", &p->frequency, 0.001f, 0.0001f, 10.f);
            }
            if (info.hasParamIsTile)
            {
                auto* p = n.entity.try_get_mut<NodeParam_IsTile>();
                ImGui::DragInt("RelX##", reinterpret_cast<int*>(&p->relativeX), 1.f, -4, 4);
                ImGui::DragInt("RelY##", reinterpret_cast<int*>(&p->relativeY), 1.f, -4, 4);
                int t = static_cast<int>(p->tileType);
                if (ImGui::Combo("Type##", &t, "Air\0Filler\0Liquid\0Ore\0NPC\0Plant\0"))
                    p->tileType = static_cast<TileType>(t);
            }
            if (info.hasParamTileDistance)
            {
                auto* p = n.entity.try_get_mut<NodeParam_TileDistance>();
                int r = p->range;
                if (ImGui::SliderInt("Range##", &r, 1, 3)) p->range = static_cast<int8_t>(r);
                int t = static_cast<int>(p->tileType);
                if (ImGui::Combo("Type##td", &t, "Air\0Filler\0Liquid\0Ore\0NPC\0Plant\0"))
                    p->tileType = static_cast<TileType>(t);
            }
            ImGui::PopItemWidth();

            // Output pin
            ImNodes::BeginOutputAttribute(OutputAttr(n.nodeId));
            ImGui::TextUnformatted("Out");
            ImNodes::EndOutputAttribute();

            ImNodes::EndNode();

            if (n.isOutput)
                ImNodes::PopColorStyle();
        }

        // ── Draw links ────────────────────────────────────────────────────────
        for (auto& l : Links)
            ImNodes::Link(l.linkId, l.srcAttrId, l.dstAttrId);

        ImNodes::EndNodeEditor();

        // ── Handle new link ───────────────────────────────────────────────────
        {
            int startAttr, endAttr;
            if (ImNodes::IsLinkCreated(&startAttr, &endAttr))
            {
                // Normalise: output attr should be startAttr
                if (AttrToSlot(startAttr) != 9) std::swap(startAttr, endAttr);
                if (AttrToSlot(startAttr) == 9 && AttrToSlot(endAttr) != 9)
                {
                    int srcNodeId = AttrToNode(startAttr);
                    int dstNodeId = AttrToNode(endAttr);
                    int slot      = AttrToSlot(endAttr);

                    auto srcIt = std::find_if(Nodes.begin(), Nodes.end(), [&](auto& n){ return n.nodeId == srcNodeId; });
                    auto dstIt = std::find_if(Nodes.begin(), Nodes.end(), [&](auto& n){ return n.nodeId == dstNodeId; });

                    if (srcIt != Nodes.end() && dstIt != Nodes.end())
                    {
                        flecs::entity src = srcIt->entity;
                        flecs::entity dst = dstIt->entity;

                        // Remove existing connection on this slot
                        if (slot == 0) dst.remove<InputPin0>(flecs::Wildcard);
                        if (slot == 1) dst.remove<InputPin1>(flecs::Wildcard);
                        if (slot == 2) dst.remove<InputPin2>(flecs::Wildcard);

                        // Add new connection
                        if (slot == 0) dst.add<InputPin0>(src);
                        if (slot == 1) dst.add<InputPin1>(src);
                        if (slot == 2) dst.add<InputPin2>(src);

                        Links.push_back({ NextId++, startAttr, endAttr });
                    }
                }
            }
        }

        // ── Handle link deletion ──────────────────────────────────────────────
        {
            int linkId;
            if (ImNodes::IsLinkDestroyed(&linkId))
            {
                auto it = std::find_if(Links.begin(), Links.end(), [&](auto& l){ return l.linkId == linkId; });
                if (it != Links.end())
                {
                    int dstNodeId = AttrToNode(it->dstAttrId);
                    int slot      = AttrToSlot(it->dstAttrId);
                    auto dstIt = std::find_if(Nodes.begin(), Nodes.end(), [&](auto& n){ return n.nodeId == dstNodeId; });
                    if (dstIt != Nodes.end())
                    {
                        if (slot == 0) dstIt->entity.remove<InputPin0>(flecs::Wildcard);
                        if (slot == 1) dstIt->entity.remove<InputPin1>(flecs::Wildcard);
                        if (slot == 2) dstIt->entity.remove<InputPin2>(flecs::Wildcard);
                    }
                    Links.erase(it);
                }
            }
        }

        // ── Context menu: mark output / delete ───────────────────────────────
        {
            int hoveredNode = -1;
            if (ImNodes::IsNodeHovered(&hoveredNode) && ImGui::IsMouseClicked(ImGuiMouseButton_Right))
            {
                ContextMenuNodeId = hoveredNode;
                ImGui::OpenPopup("node_ctx");
            }

            if (ImGui::BeginPopup("node_ctx"))
            {
                int nodeId = ContextMenuNodeId;
                auto it = std::find_if(Nodes.begin(), Nodes.end(), [&](auto& n){ return n.nodeId == nodeId; });
                if (it != Nodes.end())
                {
                    if (ImGui::MenuItem("Set as Output"))
                    {
                        for (auto& n : Nodes) { n.isOutput = false; n.entity.remove<VisualNodeOutput>(); }
                        it->isOutput = true;
                        it->entity.add<VisualNodeOutput>();
                    }
                    if (ImGui::MenuItem("Delete Node"))
                    {
                        // Remove all links connected to this node
                        Links.erase(std::remove_if(Links.begin(), Links.end(), [&](auto& l)
                        {
                            return AttrToNode(l.srcAttrId) == nodeId || AttrToNode(l.dstAttrId) == nodeId;
                        }), Links.end());
                        it->entity.destruct();
                        Nodes.erase(it);
                    }
                }
                ImGui::EndPopup();
            }
        }

        ImGui::End();
    }

    flecs::world   World;
    flecs::entity  GraphRoot;
    std::vector<EditorNode> Nodes;
    std::vector<EditorLink> Links;
    int NextId = 1;
    int ContextMenuNodeId = -1;
};

// ─── GLFW error callback ──────────────────────────────────────────────────────

static void GLFWErrorCallback(int error, const char* desc)
{
    fprintf(stderr, "GLFW error %d: %s\n", error, desc);
}

// ─── main ─────────────────────────────────────────────────────────────────────

int main()
{
    glfwSetErrorCallback(GLFWErrorCallback);
    if (!glfwInit()) return 1;

    const char* glsl_version = "#version 330";
    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);

    GLFWwindow* window = glfwCreateWindow(1280, 720, "World Graph Editor", nullptr, nullptr);
    if (!window) { glfwTerminate(); return 1; }

    glfwMakeContextCurrent(window);
    glfwSwapInterval(1);

    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImNodes::CreateContext();

    ImGui::StyleColorsDark();
    ImNodes::StyleColorsDark();

    ImGui_ImplGlfw_InitForOpenGL(window, true);
    ImGui_ImplOpenGL3_Init(glsl_version);

    NodeEditorApp app;

    while (!glfwWindowShouldClose(window))
    {
        glfwPollEvents();

        ImGui_ImplOpenGL3_NewFrame();
        ImGui_ImplGlfw_NewFrame();
        ImGui::NewFrame();

        // Ctrl+S / Ctrl+O shortcuts
        if (ImGui::IsKeyDown(ImGuiKey_LeftCtrl) || ImGui::IsKeyDown(ImGuiKey_RightCtrl))
        {
            if (ImGui::IsKeyPressed(ImGuiKey_S)) app.Save("graph.json");
            if (ImGui::IsKeyPressed(ImGuiKey_O)) app.Load("graph.json");
        }

        app.Render();

        ImGui::Render();
        int w, h;
        glfwGetFramebufferSize(window, &w, &h);
        glViewport(0, 0, w, h);
        glClearColor(0.15f, 0.15f, 0.15f, 1.f);
        glClear(GL_COLOR_BUFFER_BIT);
        ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());
        glfwSwapBuffers(window);
    }

    ImGui_ImplOpenGL3_Shutdown();
    ImGui_ImplGlfw_Shutdown();
    ImNodes::DestroyContext();
    ImGui::DestroyContext();
    glfwDestroyWindow(window);
    glfwTerminate();
    return 0;
}