// WorldGraph Node Editor
// Standalone Dear ImGui + imnodes tool for editing WorldGraph node graphs.
// Each node displays a 64×64 preview image of its output, ShaderGraph-style.
// The permanent "World Output" node collects generation passes and renders
// the final generated chunk using GenerateChunk().

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

#include "WorldGraph/WorldGraph.h"
#include "WorldGraph/WorldGraphNode.h"
#include "WorldGraph/WorldGraphTypes.h"
#include "WorldGraph/WorldGraphSerializer.h"
#include "WorldGraph/WorldGraphChunk.h"

#include <nlohmann/json.hpp>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <filesystem>
#include <fstream>
#include <memory>
#include <string>
#include <typeindex>
#include <unordered_map>
#include <vector>

using namespace WorldGraph;

// ─────────────────────────────────────────────────────────────────────────────
// Pin / link ID space
//
//  Regular node output pin  →  nodeID * 10 + 9
//  Regular node input slot S →  nodeID * 10 + S   (S = 0..8)
//
//  World Output input pin N  →  WORLD_OUTPUT_PIN_BASE + N
//  World Output static attr  →  WORLD_OUTPUT_PIN_BASE + 500
//
//  Regular link IDs          →  0 .. (num regular connections - 1)
//  World Output link N       →  WORLD_OUTPUT_LINK_BASE + N
// ─────────────────────────────────────────────────────────────────────────────

static constexpr int WORLD_OUTPUT_IMNODES_ID = 0;
static constexpr int WORLD_OUTPUT_PIN_BASE   = 10000;
static constexpr int WORLD_OUTPUT_LINK_BASE  = 20000;

static int OutputPin(Graph::NodeID n)           { return static_cast<int>(n) * 10 + 9; }
static int InputPin (Graph::NodeID n, int slot) { return static_cast<int>(n) * 10 + slot; }
static Graph::NodeID PinToNode(int attr)        { return static_cast<Graph::NodeID>(attr / 10); }
static int           PinToSlot(int attr)        { return attr % 10; }  // 9 = output

static bool IsWorldOutputPin(int attr)
{
    return attr >= WORLD_OUTPUT_PIN_BASE && attr < WORLD_OUTPUT_PIN_BASE + 500;
}

// ─────────────────────────────────────────────────────────────────────────────
// Tile registry – maps integer IDs to human-readable names
// ─────────────────────────────────────────────────────────────────────────────

struct TileEntry {
    int32_t     id       { 0 };
    std::string name     { "Tile" };
    float       color[3] { 1.0f, 1.0f, 1.0f };  // RGB in [0, 1]
};

// ─────────────────────────────────────────────────────────────────────────────
// Color helpers
// ─────────────────────────────────────────────────────────────────────────────

static void HsvToRgb(float h, float s, float v, uint8_t& r, uint8_t& g, uint8_t& b)
{
    float c = v * s;
    float x = c * (1.0f - std::fabs(std::fmod(h * 6.0f, 2.0f) - 1.0f));
    float m = v - c;
    float r1, g1, b1;
    int hi = static_cast<int>(h * 6.0f) % 6;
    switch (hi) {
        case 0: r1 = c; g1 = x; b1 = 0; break;
        case 1: r1 = x; g1 = c; b1 = 0; break;
        case 2: r1 = 0; g1 = c; b1 = x; break;
        case 3: r1 = 0; g1 = x; b1 = c; break;
        case 4: r1 = x; g1 = 0; b1 = c; break;
        default:r1 = c; g1 = 0; b1 = x; break;
    }
    r = static_cast<uint8_t>((r1 + m) * 255.0f);
    g = static_cast<uint8_t>((g1 + m) * 255.0f);
    b = static_cast<uint8_t>((b1 + m) * 255.0f);
}

static void ValueToRGBA(const Value& v,
                         const std::vector<TileEntry>* registry,
                         uint8_t& r, uint8_t& g, uint8_t& b, uint8_t& a)
{
    a = 255;
    switch (v.type) {
        case Type::Float: {
            uint8_t grey = static_cast<uint8_t>(std::clamp(v.AsFloat(), 0.0f, 1.0f) * 255.0f);
            r = g = b = grey;
            return;
        }
        case Type::Bool: {
            uint8_t c = v.AsBool() ? 255 : 0;
            r = g = b = c;
            return;
        }
        case Type::Int:
            break;
    }

    // Int path — look up registry first, fall back to hash.
    int32_t id = v.AsInt();
    if (registry) {
        for (const auto& t : *registry) {
            if (t.id == id) {
                r = static_cast<uint8_t>(t.color[0] * 255.0f);
                g = static_cast<uint8_t>(t.color[1] * 255.0f);
                b = static_cast<uint8_t>(t.color[2] * 255.0f);
                return;
            }
        }
    }
    // Fallback: tile 0 (AIR) = near-black, others get a hashed hue.
    if (id == 0) { r = g = b = 30; return; }
    uint32_t h = static_cast<uint32_t>(id) * 2654435761u;
    HsvToRgb((h & 0xFFu) / 255.0f, 0.7f, 0.85f, r, g, b);
}

// Pin colors by value type.
//   Float → blue   Int → orange   Bool → green
static ImU32 PinColor(Type t)
{
    switch (t) {
        case Type::Float: return IM_COL32(100, 150, 220, 255);
        case Type::Int:   return IM_COL32(220, 160,  60, 255);
        case Type::Bool:  return IM_COL32( 80, 200, 100, 255);
    }
    return IM_COL32(200, 200, 200, 255);
}
static ImU32 PinColorHovered(Type t)
{
    switch (t) {
        case Type::Float: return IM_COL32(130, 180, 255, 255);
        case Type::Int:   return IM_COL32(255, 195,  90, 255);
        case Type::Bool:  return IM_COL32(110, 235, 130, 255);
    }
    return IM_COL32(230, 230, 230, 255);
}

// ─────────────────────────────────────────────────────────────────────────────
// VisualNode – wraps a graph node with its preview texture
// ─────────────────────────────────────────────────────────────────────────────

static constexpr int PREVIEW_SIZE = 64;

static GLuint MakePreviewTexture()
{
    GLuint tex = 0;
    glGenTextures(1, &tex);
    glBindTexture(GL_TEXTURE_2D, tex);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    std::array<uint8_t, PREVIEW_SIZE * PREVIEW_SIZE * 4> buf;
    buf.fill(40);
    for (int i = 3; i < (int)buf.size(); i += 4) buf[i] = 255;
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, PREVIEW_SIZE, PREVIEW_SIZE,
                 0, GL_RGBA, GL_UNSIGNED_BYTE, buf.data());
    return tex;
}

struct VisualNode {
    Graph::NodeID id    { Graph::INVALID_ID };
    GLuint        tex   { 0 };
    bool          dirty { true };

    void InitTex() { tex = MakePreviewTexture(); }
    void FreeTex() { if (tex) { glDeleteTextures(1, &tex); tex = 0; } }
};

// ─────────────────────────────────────────────────────────────────────────────
// Node catalogue – all creatable node types
// ─────────────────────────────────────────────────────────────────────────────

struct NodeMenuItem {
    const char* label;
    const char* category;
    std::function<std::unique_ptr<Node>()> create;
};

static const std::vector<NodeMenuItem> NODE_MENU = {
    { "Constant",      "Source",  [] { return std::make_unique<ConstantNode>(0.0f); } },
    { "TileID",        "Source",  [] { return std::make_unique<IDNode>(1); } },
    { "PositionX",     "Source",  [] { return std::make_unique<PositionXNode>(); } },
    { "PositionY",     "Source",  [] { return std::make_unique<PositionYNode>(); } },
    { "Add",           "Math",    [] { return std::make_unique<AddNode>(); } },
    { "Subtract",      "Math",    [] { return std::make_unique<SubtractNode>(); } },
    { "Multiply",      "Math",    [] { return std::make_unique<MultiplyNode>(); } },
    { "Divide",        "Math",    [] { return std::make_unique<DivideNode>(); } },
    { "Modulo",        "Math",    [] { return std::make_unique<ModuloNode>(); } },
    { "Sin",           "Math",    [] { return std::make_unique<SinNode>(); } },
    { "Cos",           "Math",    [] { return std::make_unique<CosNode>(); } },
    { "Less",          "Compare", [] { return std::make_unique<LessNode>(); } },
    { "Greater",       "Compare", [] { return std::make_unique<GreaterNode>(); } },
    { "LessEqual",     "Compare", [] { return std::make_unique<LessEqualNode>(); } },
    { "GreaterEqual",  "Compare", [] { return std::make_unique<GreaterEqualNode>(); } },
    { "Equal",         "Compare", [] { return std::make_unique<EqualNode>(); } },
    { "And",           "Logic",   [] { return std::make_unique<AndNode>(); } },
    { "Or",            "Logic",   [] { return std::make_unique<OrNode>(); } },
    { "Not",           "Logic",   [] { return std::make_unique<NotNode>(); } },
    { "Branch",        "Control", [] { return std::make_unique<BranchNode>(); } },
    { "QueryTile",     "Query",   [] { return std::make_unique<QueryTileNode>(0, -1, 1); } },
    { "QueryRange",    "Query",   [] { return std::make_unique<QueryRangeNode>(-1, -1, 1, 1, 1); } },
    { "QueryDistance", "Query",   [] { return std::make_unique<QueryDistanceNode>(1, 4); } },
};

// ─────────────────────────────────────────────────────────────────────────────
// NodeEditorApp
// ─────────────────────────────────────────────────────────────────────────────

class NodeEditorApp {
public:
    NodeEditorApp()
    {
        worldOutputTex = MakePreviewTexture();
    }

    ~NodeEditorApp()
    {
        for (auto& [id, vn] : visualNodes) vn.FreeTex();
        if (worldOutputTex) glDeleteTextures(1, &worldOutputTex);
    }

    // ── Preview settings ──────────────────────────────────────────────────────

    int      previewOriginX { 0 };
    int      previewOriginY { 0 };
    float    previewScale   { 1.0f };  // world units per pixel (per-node previews only)
    uint64_t previewSeed    { 0 };

    // ── Public render entry point ─────────────────────────────────────────────

    void Render()
    {
        // Position World Output node on first frame (or after NewGraph).
        if (pendingWorldOutputPos) {
            pendingWorldOutputPos = false;
            ImNodes::SetNodeGridSpacePos(WORLD_OUTPUT_IMNODES_ID, ImVec2(400.0f, 0.0f));
        }

        RenderDirtyPreviews();

        DrawMenuBar();

        ImGui::SetNextWindowPos(ImVec2(0, 18), ImGuiCond_Always);
        ImGui::SetNextWindowSize(ImVec2(200, ImGui::GetIO().DisplaySize.y - 18), ImGuiCond_Always);
        ImGui::Begin("Settings", nullptr,
                     ImGuiWindowFlags_NoResize | ImGuiWindowFlags_NoMove |
                     ImGuiWindowFlags_NoCollapse | ImGuiWindowFlags_NoBringToFrontOnFocus);
        DrawSettingsPanel();
        ImGui::End();

        ImGui::SetNextWindowPos(ImVec2(200, 18), ImGuiCond_Always);
        ImGui::SetNextWindowSize(
            ImVec2(ImGui::GetIO().DisplaySize.x - 200, ImGui::GetIO().DisplaySize.y - 18),
            ImGuiCond_Always);
        ImGui::Begin("Node Canvas", nullptr,
                     ImGuiWindowFlags_NoResize | ImGuiWindowFlags_NoMove |
                     ImGuiWindowFlags_NoCollapse | ImGuiWindowFlags_NoBringToFrontOnFocus |
                     ImGuiWindowFlags_NoScrollbar | ImGuiWindowFlags_NoScrollWithMouse);
        DrawNodeCanvas();
        ImGui::End();
    }

private:
    Graph graph;
    std::unordered_map<Graph::NodeID, VisualNode> visualNodes;
    std::string currentFile;
    std::vector<TileEntry> tileRegistry { TileEntry{0, "Empty", {30.0f/255.0f, 30.0f/255.0f, 30.0f/255.0f}} };

    // ── World Output node state ───────────────────────────────────────────────

    // Each entry is the graph node ID connected to that pass slot (INVALID_ID = empty).
    std::vector<Graph::NodeID> worldOutputPasses { Graph::INVALID_ID };
    GLuint worldOutputTex   { 0 };
    bool   worldOutputDirty { true };

    // Deferred first-frame positioning of the World Output node.
    bool pendingWorldOutputPos { true };

    // ── Preview rendering ─────────────────────────────────────────────────────

    void RenderDirtyPreviews()
    {
        std::array<uint8_t, PREVIEW_SIZE * PREVIEW_SIZE * 4> pixels;

        // Per-node previews (evaluate single node across a grid).
        for (auto& [id, vn] : visualNodes) {
            if (!vn.dirty) continue;
            vn.dirty = false;

            Node* node = graph.GetNode(id);
            if (!node) continue;

            for (int py = 0; py < PREVIEW_SIZE; ++py) {
                for (int px = 0; px < PREVIEW_SIZE; ++px) {
                    EvalContext ctx;
                    ctx.worldX = previewOriginX + static_cast<int>(px * previewScale);
                    ctx.worldY = previewOriginY + static_cast<int>(py * previewScale);
                    ctx.seed   = previewSeed;

                    Value v = graph.Evaluate(id, ctx);
                    int base = (py * PREVIEW_SIZE + px) * 4;
                    ValueToRGBA(v, &tileRegistry, pixels[base], pixels[base+1], pixels[base+2], pixels[base+3]);
                }
            }

            glBindTexture(GL_TEXTURE_2D, vn.tex);
            glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, PREVIEW_SIZE, PREVIEW_SIZE,
                         0, GL_RGBA, GL_UNSIGNED_BYTE, pixels.data());
        }

        // World Output preview — runs GenerateChunk with all connected passes.
        if (worldOutputDirty) {
            worldOutputDirty = false;

            std::vector<GenerationPass> passes;
            for (auto nodeID : worldOutputPasses) {
                if (nodeID != Graph::INVALID_ID && graph.GetNode(nodeID))
                    passes.push_back({ graph, nodeID });
            }

            if (passes.empty()) {
                pixels.fill(30);
                for (int i = 3; i < (int)pixels.size(); i += 4) pixels[i] = 255;
            } else {
                // Each pixel = 1 world tile starting at previewOriginX/Y.
                TileGrid chunk = GenerateChunk(passes,
                    previewOriginX, previewOriginY,
                    PREVIEW_SIZE, PREVIEW_SIZE,
                    previewSeed);

                for (int py = 0; py < PREVIEW_SIZE; ++py) {
                    for (int px = 0; px < PREVIEW_SIZE; ++px) {
                        int32_t tileID = chunk.Get(previewOriginX + px, previewOriginY + py);
                        int base = (py * PREVIEW_SIZE + px) * 4;
                        Value v = Value::MakeInt(tileID);
                        ValueToRGBA(v, &tileRegistry, pixels[base], pixels[base+1], pixels[base+2], pixels[base+3]);
                    }
                }
            }

            glBindTexture(GL_TEXTURE_2D, worldOutputTex);
            glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, PREVIEW_SIZE, PREVIEW_SIZE,
                         0, GL_RGBA, GL_UNSIGNED_BYTE, pixels.data());
        }
    }

    void MarkAllDirty()
    {
        for (auto& [id, vn] : visualNodes) vn.dirty = true;
        worldOutputDirty = true;
    }

    // ── UI drawing ────────────────────────────────────────────────────────────

    void DrawMenuBar()
    {
        if (!ImGui::BeginMainMenuBar()) return;
        if (ImGui::BeginMenu("File")) {
            if (ImGui::MenuItem("New"))               { NewGraph(); }
            if (ImGui::MenuItem("Open...", "Ctrl+O")) { OpenFileDialog(); }
            if (ImGui::MenuItem("Save",    "Ctrl+S")) { SaveCurrent(); }
            if (ImGui::MenuItem("Save As..."))        { SaveAsDialog(); }
            ImGui::EndMenu();
        }
        if (!currentFile.empty())
            ImGui::TextDisabled("  %s", currentFile.c_str());
        ImGui::EndMainMenuBar();
    }

    void DrawSettingsPanel()
    {
        ImGui::SeparatorText("Preview");

        bool changed = false;
        changed |= ImGui::DragInt("Origin X", &previewOriginX, 1.0f);
        changed |= ImGui::DragInt("Origin Y", &previewOriginY, 1.0f);
        changed |= ImGui::DragFloat("Scale",  &previewScale, 0.1f, 0.1f, 64.0f, "%.2f wp/px");

        int seed32 = static_cast<int>(previewSeed & 0xFFFFFFFFu);
        if (ImGui::DragInt("Seed", &seed32)) {
            previewSeed = static_cast<uint64_t>(static_cast<uint32_t>(seed32));
            changed = true;
        }

        if (changed) MarkAllDirty();

        ImGui::Spacing();
        ImGui::SeparatorText("Tile IDs");

        int toRemove = -1;
        for (int i = 0; i < static_cast<int>(tileRegistry.size()); ++i) {
            ImGui::PushID(i);
            TileEntry& entry = tileRegistry[i];

            // Color button — opens a popup picker.
            char cpopup[32];
            snprintf(cpopup, sizeof(cpopup), "##cpick%d", i);
            if (ImGui::ColorButton(cpopup,
                    ImVec4(entry.color[0], entry.color[1], entry.color[2], 1.0f),
                    ImGuiColorEditFlags_NoTooltip | ImGuiColorEditFlags_NoBorder,
                    ImVec2(18, 18))) {
                ImGui::OpenPopup(cpopup);
            }
            if (ImGui::BeginPopup(cpopup)) {
                if (ImGui::ColorPicker3("##cp", entry.color,
                        ImGuiColorEditFlags_NoLabel | ImGuiColorEditFlags_NoSidePreview))
                    MarkAllDirty();
                ImGui::EndPopup();
            }
            ImGui::SameLine(0, 3);

            // ID field (narrow) — locked for the reserved Empty entry
            ImGui::BeginDisabled(entry.id == 0);
            ImGui::SetNextItemWidth(32);
            ImGui::DragInt("##tid", &entry.id, 1, 0, 9999);
            ImGui::EndDisabled();
            ImGui::SameLine(0, 3);

            // Name field (fills remaining width minus remove button)
            ImGui::SetNextItemWidth(ImGui::GetContentRegionAvail().x - 20);
            char buf[64];
            strncpy(buf, entry.name.c_str(), sizeof(buf) - 1);
            buf[sizeof(buf) - 1] = '\0';
            if (ImGui::InputText("##tname", buf, sizeof(buf)))
                entry.name = buf;
            ImGui::SameLine(0, 3);

            // ID 0 "Empty" is permanent — show a disabled remove button
            ImGui::BeginDisabled(entry.id == 0);
            if (ImGui::SmallButton("x"))
                toRemove = i;
            ImGui::EndDisabled();

            ImGui::PopID();
        }
        if (toRemove >= 0)
            tileRegistry.erase(tileRegistry.begin() + toRemove);

        if (ImGui::SmallButton("+ Add Tile")) {
            int32_t nextId = tileRegistry.empty() ? 1
                           : tileRegistry.back().id + 1;
            TileEntry entry;
            entry.id   = nextId;
            entry.name = "Tile";
            // Seed the default color from the hash so new tiles start distinct.
            uint8_t hr, hg, hb;
            uint32_t h = static_cast<uint32_t>(nextId) * 2654435761u;
            HsvToRgb((h & 0xFFu) / 255.0f, 0.7f, 0.85f, hr, hg, hb);
            entry.color[0] = hr / 255.0f;
            entry.color[1] = hg / 255.0f;
            entry.color[2] = hb / 255.0f;
            tileRegistry.push_back(entry);
        }

        ImGui::Spacing();
        ImGui::SeparatorText("World Output");

        int connected = 0;
        for (auto id : worldOutputPasses)
            if (id != Graph::INVALID_ID) ++connected;
        ImGui::Text("%d / %zu pass(es) connected", connected, worldOutputPasses.size());

        ImGui::Spacing();
        ImGui::SeparatorText("Graph");
        ImGui::Text("%zu nodes", graph.NodeCount());

        ImGui::Spacing();
        ImGui::SeparatorText("Help");
        ImGui::TextWrapped(
            "Right-click canvas to add nodes.\n"
            "Drag an output pin into a World Output pass slot to register a generation pass.\n"
            "Delete key removes selected nodes.\n"
            "Scale controls per-node previews only; World Output always shows 1 tile/pixel.\n"
            "\nQuery nodes preview blank (no prev-pass data in per-node preview)."
        );
    }

    void DrawNodeCanvas()
    {
        ImNodes::BeginNodeEditor();

        // Right-click blank canvas → add node menu
        if (ImGui::IsWindowHovered(ImGuiFocusedFlags_RootAndChildWindows) &&
            ImNodes::IsEditorHovered() &&
            ImGui::IsMouseReleased(ImGuiMouseButton_Right) &&
            !ImGui::IsAnyItemHovered())
        {
            ImGui::OpenPopup("##add_node_menu");
        }

        ImGui::SetNextWindowSize(ImVec2(220, 300), ImGuiCond_Always);
        if (ImGui::BeginPopup("##add_node_menu")) {
            ImVec2 spawnPos = ImGui::GetMousePosOnOpeningCurrentPopup();

            static char filterBuf[128] = "";
            if (ImGui::IsWindowAppearing()) {
                ImGui::SetKeyboardFocusHere();
                filterBuf[0] = '\0';
            }
            ImGui::SetNextItemWidth(-1);
            ImGui::InputText("##filter", filterBuf, sizeof(filterBuf));
            ImGui::Separator();

            const bool filtering = filterBuf[0] != '\0';

            auto spawnItem = [&](const NodeMenuItem& item) {
                if (ImGui::MenuItem(item.label)) {
                    Graph::NodeID newId = AddNode(item.create());
                    ImNodes::SetNodeScreenSpacePos(static_cast<int>(newId), spawnPos);
                }
            };

            // Scrollable list region — fills the remaining space below the filter.
            ImGui::BeginChild("##node_list", ImVec2(0, 0), false);

            if (filtering) {
                // Case-insensitive flat list — matches label or category name.
                std::string needle = filterBuf;
                for (char& c : needle) c = static_cast<char>(std::tolower(static_cast<unsigned char>(c)));

                for (const auto& item : NODE_MENU) {
                    auto contains = [&](const char* s) {
                        std::string hay = s;
                        for (char& c : hay) c = static_cast<char>(std::tolower(static_cast<unsigned char>(c)));
                        return hay.find(needle) != std::string::npos;
                    };
                    if (contains(item.label) || contains(item.category))
                        spawnItem(item);
                }
            } else {
                // Tree view grouped by category, all expanded by default.
                const char* curCat = nullptr;
                bool catOpen = false;
                for (const auto& item : NODE_MENU) {
                    if (!curCat || strcmp(curCat, item.category) != 0) {
                        if (curCat && catOpen) ImGui::TreePop();
                        catOpen = ImGui::TreeNodeEx(item.category, ImGuiTreeNodeFlags_DefaultOpen);
                        curCat  = item.category;
                    }
                    if (catOpen)
                        spawnItem(item);
                }
                if (curCat && catOpen) ImGui::TreePop();
            }

            ImGui::EndChild();
            ImGui::EndPopup();
        }

        // Draw the permanent World Output node first (so it's rendered behind others).
        DrawWorldOutputNode();

        // Draw all regular nodes.
        for (auto& [id, vn] : visualNodes)
            DrawNode(vn);

        // Draw regular graph connections.
        int linkId = 0;
        for (auto& [id, vn] : visualNodes) {
            Node* node = graph.GetNode(id);
            if (!node) continue;
            for (int slot = 0; slot < static_cast<int>(node->GetInputCount()); ++slot) {
                auto src = graph.GetInput(id, slot);
                if (src.has_value())
                    ImNodes::Link(linkId++, OutputPin(*src), InputPin(id, slot));
            }
        }

        // Draw World Output pass connections (fixed link IDs avoid re-enumeration issues).
        for (int i = 0; i < static_cast<int>(worldOutputPasses.size()); ++i) {
            if (worldOutputPasses[i] != Graph::INVALID_ID)
                ImNodes::Link(WORLD_OUTPUT_LINK_BASE + i,
                              OutputPin(worldOutputPasses[i]),
                              WORLD_OUTPUT_PIN_BASE + i);
        }

        ImNodes::EndNodeEditor();

        // ── Handle new connections ────────────────────────────────────────────

        int fromAttr, toAttr;
        if (ImNodes::IsLinkCreated(&fromAttr, &toAttr)) {
            bool fromIsWO = IsWorldOutputPin(fromAttr);
            bool toIsWO   = IsWorldOutputPin(toAttr);

            if (toIsWO || fromIsWO) {
                // One end is a World Output pass slot; the other must be a regular output.
                int woPin      = toIsWO ? toAttr  : fromAttr;
                int regularPin = toIsWO ? fromAttr : toAttr;

                if (PinToSlot(regularPin) == 9) {  // must be an output pin
                    int passIdx = woPin - WORLD_OUTPUT_PIN_BASE;
                    if (passIdx >= 0 && passIdx < static_cast<int>(worldOutputPasses.size())) {
                        // World Output only accepts Int (tile ID) outputs.
                        Node* srcNode = graph.GetNode(PinToNode(regularPin));
                        if (srcNode && srcNode->GetOutputType() == Type::Int) {
                            worldOutputPasses[passIdx] = PinToNode(regularPin);
                            worldOutputDirty = true;
                        }
                    }
                }
            } else {
                // Regular node-to-node connection.
                Graph::NodeID fromNode, toNode;
                int toSlot;
                if (PinToSlot(fromAttr) == 9 && PinToSlot(toAttr) != 9) {
                    fromNode = PinToNode(fromAttr);
                    toNode   = PinToNode(toAttr);
                    toSlot   = PinToSlot(toAttr);
                } else if (PinToSlot(toAttr) == 9 && PinToSlot(fromAttr) != 9) {
                    fromNode = PinToNode(toAttr);
                    toNode   = PinToNode(fromAttr);
                    toSlot   = PinToSlot(fromAttr);
                } else {
                    goto done_link;
                }
                // Only connect when output type matches the destination input type.
                Node* srcNode = graph.GetNode(fromNode);
                Node* dstNode = graph.GetNode(toNode);
                if (srcNode && dstNode) {
                    auto dstInputTypes = dstNode->GetInputTypes();
                    if (toSlot < static_cast<int>(dstInputTypes.size()) &&
                        srcNode->GetOutputType() == dstInputTypes[toSlot])
                    {
                        if (graph.Connect(fromNode, toNode, toSlot))
                            MarkAllDirty();
                    }
                }
            }
        }
        done_link:;

        // ── Handle link deletion ──────────────────────────────────────────────

        int destroyedLink;
        if (ImNodes::IsLinkDestroyed(&destroyedLink)) {
            if (destroyedLink >= WORLD_OUTPUT_LINK_BASE) {
                // World Output link
                int passIdx = destroyedLink - WORLD_OUTPUT_LINK_BASE;
                if (passIdx >= 0 && passIdx < static_cast<int>(worldOutputPasses.size())) {
                    worldOutputPasses[passIdx] = Graph::INVALID_ID;
                    worldOutputDirty = true;
                }
            } else {
                // Regular link — re-enumerate to find which one.
                int idx = 0;
                bool found = false;
                for (auto& [id, vn] : visualNodes) {
                    if (found) break;
                    Node* node = graph.GetNode(id);
                    if (!node) continue;
                    for (int slot = 0; slot < static_cast<int>(node->GetInputCount()); ++slot) {
                        if (graph.GetInput(id, slot).has_value()) {
                            if (idx == destroyedLink) {
                                graph.Disconnect(id, slot);
                                MarkAllDirty();
                                found = true;
                                break;
                            }
                            ++idx;
                        }
                    }
                }
            }
        }

        // ── Delete selected nodes ─────────────────────────────────────────────

        if (ImGui::IsKeyPressed(ImGuiKey_Delete)) {
            int count = ImNodes::NumSelectedNodes();
            if (count > 0) {
                std::vector<int> selected(count);
                ImNodes::GetSelectedNodes(selected.data());
                for (int sid : selected) {
                    if (sid == WORLD_OUTPUT_IMNODES_ID) continue;  // permanent node

                    Graph::NodeID gid = static_cast<Graph::NodeID>(sid);

                    // Clear any World Output pass that referenced this node.
                    for (auto& pass : worldOutputPasses)
                        if (pass == gid) { pass = Graph::INVALID_ID; worldOutputDirty = true; }

                    graph.RemoveNode(gid);
                    auto it = visualNodes.find(gid);
                    if (it != visualNodes.end()) {
                        it->second.FreeTex();
                        visualNodes.erase(it);
                    }
                }
                MarkAllDirty();
            }
        }
    }

    // ── World Output node ─────────────────────────────────────────────────────

    void DrawWorldOutputNode()
    {
        ImNodes::PushColorStyle(ImNodesCol_TitleBar,         IM_COL32( 30,  80, 160, 255));
        ImNodes::PushColorStyle(ImNodesCol_TitleBarHovered,  IM_COL32( 45, 100, 190, 255));
        ImNodes::PushColorStyle(ImNodesCol_TitleBarSelected, IM_COL32( 60, 120, 210, 255));

        ImNodes::BeginNode(WORLD_OUTPUT_IMNODES_ID);

        ImNodes::BeginNodeTitleBar();
        ImGui::Text("World Output");
        ImNodes::EndNodeTitleBar();

        // One input pin per generation pass — Int only.
        for (int i = 0; i < static_cast<int>(worldOutputPasses.size()); ++i) {
            ImNodes::PushColorStyle(ImNodesCol_Pin,        PinColor(Type::Int));
            ImNodes::PushColorStyle(ImNodesCol_PinHovered, PinColorHovered(Type::Int));
            ImNodes::BeginInputAttribute(WORLD_OUTPUT_PIN_BASE + i);
            Graph::NodeID connected = worldOutputPasses[i];
            if (connected != Graph::INVALID_ID) {
                Node* n = graph.GetNode(connected);
                ImGui::Text("Pass %d  (%s)", i, n ? n->GetName().c_str() : "?");
            } else {
                ImGui::TextDisabled("Pass %d  (empty)", i);
            }
            ImNodes::EndInputAttribute();
            ImNodes::PopColorStyle();
            ImNodes::PopColorStyle();
        }

        // Add / remove pass buttons.
        ImNodes::BeginStaticAttribute(WORLD_OUTPUT_PIN_BASE + 500);
        if (ImGui::SmallButton("+ Pass")) {
            worldOutputPasses.push_back(Graph::INVALID_ID);
        }
        if (worldOutputPasses.size() > 1) {
            ImGui::SameLine();
            if (ImGui::SmallButton("- Pass")) {
                worldOutputPasses.pop_back();
                worldOutputDirty = true;
            }
        }
        ImNodes::EndStaticAttribute();

        // Generated chunk preview — always 1 tile per pixel.
        ImGui::Image(
            static_cast<ImTextureID>(static_cast<uint64_t>(worldOutputTex)),
            ImVec2(PREVIEW_SIZE, PREVIEW_SIZE));

        ImNodes::EndNode();

        ImNodes::PopColorStyle();
        ImNodes::PopColorStyle();
        ImNodes::PopColorStyle();
    }

    // ── Regular node ─────────────────────────────────────────────────────────

    void DrawNode(VisualNode& vn)
    {
        Node* node = graph.GetNode(vn.id);
        if (!node) return;

        ImNodes::BeginNode(static_cast<int>(vn.id));

        ImNodes::BeginNodeTitleBar();
        ImGui::Text("%s", node->GetName().c_str());
        ImNodes::EndNodeTitleBar();

        // Input pins.
        auto inputTypes = node->GetInputTypes();
        for (int slot = 0; slot < static_cast<int>(inputTypes.size()); ++slot) {
            ImNodes::PushColorStyle(ImNodesCol_Pin,        PinColor(inputTypes[slot]));
            ImNodes::PushColorStyle(ImNodesCol_PinHovered, PinColorHovered(inputTypes[slot]));
            ImNodes::BeginInputAttribute(InputPin(vn.id, slot));
            const char* t = inputTypes[slot] == Type::Float ? "f"
                          : inputTypes[slot] == Type::Int   ? "i" : "b";
            ImGui::TextDisabled("[%s] in%d", t, slot);
            ImNodes::EndInputAttribute();
            ImNodes::PopColorStyle();
            ImNodes::PopColorStyle();
        }

        DrawNodeParams(vn.id, node);

        ImGui::Image(
            static_cast<ImTextureID>(static_cast<uint64_t>(vn.tex)),
            ImVec2(PREVIEW_SIZE, PREVIEW_SIZE));

        ImNodes::PushColorStyle(ImNodesCol_Pin,        PinColor(node->GetOutputType()));
        ImNodes::PushColorStyle(ImNodesCol_PinHovered, PinColorHovered(node->GetOutputType()));
        ImNodes::BeginOutputAttribute(OutputPin(vn.id));
        const char* outType = node->GetOutputType() == Type::Float ? "Float"
                            : node->GetOutputType() == Type::Int   ? "Int" : "Bool";
        ImGui::Text("out (%s) →", outType);
        ImNodes::EndOutputAttribute();
        ImNodes::PopColorStyle();
        ImNodes::PopColorStyle();

        ImNodes::EndNode();
    }

    // ── Per-type param drawers ────────────────────────────────────────────────
    // Each function receives the app (for instance state like tileRegistry) and
    // the node, and returns true when any value changed.
    // Static member functions can access private members through the app ref.
    //
    // To support a new node type:
    //   1. Add a Draw*Params function here.
    //   2. Register it in the kDrawers table in DrawNodeParams below.

    static bool DrawConstantParams(NodeEditorApp& /*app*/, Node* node)
    {
        auto* n = static_cast<ConstantNode*>(node);
        return ImGui::DragFloat("##val", &n->value, 0.01f);
    }

    static bool DrawIDParams(NodeEditorApp& app, Node* node)
    {
        auto* n = static_cast<IDNode*>(node);
        bool changed = false;
        if (app.tileRegistry.empty()) {
            changed = ImGui::DragInt("##id", &n->tileID, 1, 0, 9999);
        } else {
            int selIdx = -1;
            for (int i = 0; i < static_cast<int>(app.tileRegistry.size()); ++i)
                if (app.tileRegistry[i].id == n->tileID) { selIdx = i; break; }
            const char* preview = selIdx >= 0 ? app.tileRegistry[selIdx].name.c_str() : "???";
            ImGui::SetNextItemWidth(130);
            if (ImGui::BeginCombo("##tile", preview)) {
                for (int i = 0; i < static_cast<int>(app.tileRegistry.size()); ++i) {
                    bool sel = (i == selIdx);
                    char label[80];
                    snprintf(label, sizeof(label), "%s  (%d)",
                             app.tileRegistry[i].name.c_str(), app.tileRegistry[i].id);
                    if (ImGui::Selectable(label, sel)) {
                        n->tileID = app.tileRegistry[i].id;
                        changed   = true;
                    }
                    if (sel) ImGui::SetItemDefaultFocus();
                }
                ImGui::EndCombo();
            }
        }
        return changed;
    }

    static bool DrawQueryTileParams(NodeEditorApp& /*app*/, Node* node)
    {
        auto* n = static_cast<QueryTileNode*>(node);
        bool changed = false;
        changed |= ImGui::DragInt("##ox",  &n->offsetX,    1);
        ImGui::SameLine();
        changed |= ImGui::DragInt("##oy",  &n->offsetY,    1);
        changed |= ImGui::DragInt("##eid", &n->expectedID, 1, 0, 1024);
        return changed;
    }

    static bool DrawQueryRangeParams(NodeEditorApp& /*app*/, Node* node)
    {
        auto* n = static_cast<QueryRangeNode*>(node);
        bool changed = false;
        changed |= ImGui::DragInt("##mnx",  &n->minX,   1);
        ImGui::SameLine();
        changed |= ImGui::DragInt("##mxx",  &n->maxX,   1);
        changed |= ImGui::DragInt("##mny",  &n->minY,   1);
        ImGui::SameLine();
        changed |= ImGui::DragInt("##mxy",  &n->maxY,   1);
        changed |= ImGui::DragInt("##rtid", &n->tileID, 1, 0, 1024);
        return changed;
    }

    static bool DrawQueryDistanceParams(NodeEditorApp& /*app*/, Node* node)
    {
        auto* n = static_cast<QueryDistanceNode*>(node);
        bool changed = false;
        changed |= ImGui::DragInt("##dtid", &n->tileID,      1, 0, 1024);
        changed |= ImGui::DragInt("##md",   &n->maxDistance, 1, 1,   32);
        return changed;
    }

    // ── Param dispatcher ──────────────────────────────────────────────────────

    void DrawNodeParams(Graph::NodeID /*id*/, Node* node)
    {
        using DrawFn = bool(*)(NodeEditorApp&, Node*);
        static const std::unordered_map<std::type_index, DrawFn> kDrawers = {
            { typeid(ConstantNode),      DrawConstantParams      },
            { typeid(IDNode),            DrawIDParams            },
            { typeid(QueryTileNode),     DrawQueryTileParams     },
            { typeid(QueryRangeNode),    DrawQueryRangeParams    },
            { typeid(QueryDistanceNode), DrawQueryDistanceParams },
        };

        auto it = kDrawers.find(typeid(*node));
        if (it != kDrawers.end() && it->second(*this, node))
            MarkAllDirty();
    }

    // ── Graph management ──────────────────────────────────────────────────────

    Graph::NodeID AddNode(std::unique_ptr<Node> node)
    {
        Graph::NodeID id = graph.AddNode(std::move(node));
        VisualNode vn;
        vn.id    = id;
        vn.dirty = true;
        vn.InitTex();
        visualNodes.emplace(id, std::move(vn));
        return id;
    }

    void NewGraph()
    {
        for (auto& [id, vn] : visualNodes) vn.FreeTex();
        visualNodes.clear();
        graph = Graph{};
        worldOutputPasses = { Graph::INVALID_ID };
        worldOutputDirty  = true;
        pendingWorldOutputPos = true;
        currentFile = "";
    }

    // ── Serialization ─────────────────────────────────────────────────────────

    void Save(const std::string& path)
    {
        nlohmann::json root;
        root["graph"] = GraphSerializer::ToJson(graph);

        // Node positions (including World Output).
        nlohmann::json positions = nlohmann::json::object();
        for (auto& [id, vn] : visualNodes) {
            ImVec2 pos = ImNodes::GetNodeGridSpacePos(static_cast<int>(id));
            positions[std::to_string(id)] = { pos.x, pos.y };
        }
        {
            ImVec2 pos = ImNodes::GetNodeGridSpacePos(WORLD_OUTPUT_IMNODES_ID);
            positions["__worldOutput"] = { pos.x, pos.y };
        }

        root["editor"]["nodePositions"] = positions;
        root["editor"]["seed"]           = previewSeed;
        root["editor"]["previewOriginX"] = previewOriginX;
        root["editor"]["previewOriginY"] = previewOriginY;
        root["editor"]["previewScale"]   = previewScale;

        nlohmann::json passes = nlohmann::json::array();
        for (auto id : worldOutputPasses)
            passes.push_back(id);
        root["editor"]["worldOutputPasses"] = passes;

        nlohmann::json tiles = nlohmann::json::array();
        for (auto& t : tileRegistry)
            tiles.push_back({ {"id", t.id}, {"name", t.name},
                              {"color", { t.color[0], t.color[1], t.color[2] }} });
        root["editor"]["tileRegistry"] = tiles;

        std::ofstream f(path);
        if (f) { f << root.dump(2); currentFile = path; }
        else   fprintf(stderr, "Failed to write %s\n", path.c_str());
    }

    void Load(const std::string& path)
    {
        std::ifstream f(path);
        if (!f) { fprintf(stderr, "Cannot open %s\n", path.c_str()); return; }

        nlohmann::json root;
        try { root = nlohmann::json::parse(f); }
        catch (...) { fprintf(stderr, "JSON parse error in %s\n", path.c_str()); return; }

        auto newGraph = GraphSerializer::FromJson(root.value("graph", nlohmann::json::object()));
        if (!newGraph) { fprintf(stderr, "Invalid graph in %s\n", path.c_str()); return; }

        NewGraph();
        graph = std::move(*newGraph);

        uint32_t maxId = root["graph"].value("nextId", 1u);
        for (uint32_t id = 1; id < maxId; ++id) {
            if (graph.GetNode(id)) {
                VisualNode vn;
                vn.id    = id;
                vn.dirty = true;
                vn.InitTex();
                visualNodes.emplace(id, std::move(vn));
            }
        }

        if (root.contains("editor")) {
            const auto& ed = root["editor"];
            previewSeed    = ed.value("seed",           uint64_t{0});
            previewOriginX = ed.value("previewOriginX", 0);
            previewOriginY = ed.value("previewOriginY", 0);
            previewScale   = ed.value("previewScale",   1.0f);

            if (ed.contains("tileRegistry")) {
                tileRegistry.clear();
                for (auto& t : ed["tileRegistry"]) {
                    TileEntry entry;
                    entry.id   = t.value("id",   0);
                    entry.name = t.value("name", std::string("Tile"));
                    if (t.contains("color") && t["color"].size() == 3) {
                        entry.color[0] = t["color"][0].get<float>();
                        entry.color[1] = t["color"][1].get<float>();
                        entry.color[2] = t["color"][2].get<float>();
                    }
                    tileRegistry.push_back(entry);
                }
                // Always guarantee the Empty (ID 0) sentinel is present.
                bool hasEmpty = std::any_of(tileRegistry.begin(), tileRegistry.end(),
                    [](const TileEntry& e) { return e.id == 0; });
                if (!hasEmpty)
                    tileRegistry.insert(tileRegistry.begin(),
                        TileEntry{0, "Empty", {30.0f/255.0f, 30.0f/255.0f, 30.0f/255.0f}});
            }

            if (ed.contains("worldOutputPasses")) {
                worldOutputPasses.clear();
                for (auto& v : ed["worldOutputPasses"])
                    worldOutputPasses.push_back(v.get<Graph::NodeID>());
                if (worldOutputPasses.empty())
                    worldOutputPasses.push_back(Graph::INVALID_ID);
            }

            if (ed.contains("nodePositions")) {
                for (auto& [key, val] : ed["nodePositions"].items()) {
                    float x = val[0].get<float>();
                    float y = val[1].get<float>();
                    if (key == "__worldOutput") {
                        ImNodes::SetNodeGridSpacePos(WORLD_OUTPUT_IMNODES_ID, ImVec2(x, y));
                        pendingWorldOutputPos = false;
                    } else {
                        ImNodes::SetNodeGridSpacePos(std::stoi(key), ImVec2(x, y));
                    }
                }
            }
        }

        currentFile = path;
    }

    // ── File dialogs ──────────────────────────────────────────────────────────

    void OpenFileDialog()
    {
        IGFD::FileDialogConfig cfg;
        cfg.path = currentFile.empty() ? "."
                 : std::filesystem::path(currentFile).parent_path().string();
        ImGuiFileDialog::Instance()->OpenDialog("OpenFileDlg", "Open Graph", ".json", cfg);
    }

    void SaveAsDialog()
    {
        IGFD::FileDialogConfig cfg;
        if (!currentFile.empty()) {
            std::filesystem::path p(currentFile);
            cfg.path     = p.parent_path().string();
            cfg.fileName = p.filename().string();
        } else {
            cfg.path     = ".";
            cfg.fileName = "graph.json";
        }
        ImGuiFileDialog::Instance()->OpenDialog("SaveFileDlg", "Save Graph", ".json", cfg);
    }

    void SaveCurrent() { if (currentFile.empty()) SaveAsDialog(); else Save(currentFile); }

public:
    void RenderFileDialogs()
    {
        ImVec2 dlgSize(600, 400);

        if (ImGuiFileDialog::Instance()->Display("OpenFileDlg",
                ImGuiWindowFlags_NoCollapse, dlgSize)) {
            if (ImGuiFileDialog::Instance()->IsOk())
                Load(ImGuiFileDialog::Instance()->GetFilePathName());
            ImGuiFileDialog::Instance()->Close();
        }

        if (ImGuiFileDialog::Instance()->Display("SaveFileDlg",
                ImGuiWindowFlags_NoCollapse, dlgSize)) {
            if (ImGuiFileDialog::Instance()->IsOk())
                Save(ImGuiFileDialog::Instance()->GetFilePathName());
            ImGuiFileDialog::Instance()->Close();
        }
    }

    void HandleKeyboardShortcuts()
    {
        ImGuiIO& io = ImGui::GetIO();
        if (io.KeyCtrl && ImGui::IsKeyPressed(ImGuiKey_S)) SaveCurrent();
        if (io.KeyCtrl && ImGui::IsKeyPressed(ImGuiKey_O)) ImGui::OpenPopup("##open_path");
    }
};

// ─────────────────────────────────────────────────────────────────────────────
// Main / render loop
// ─────────────────────────────────────────────────────────────────────────────

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

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

    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
    glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
#ifdef __APPLE__
    glfwWindowHint(GLFW_OPENGL_FORWARD_COMPAT, GL_TRUE);
#endif

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

    glfwMakeContextCurrent(window);
    glfwSwapInterval(1);

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

    ImGuiIO& io = ImGui::GetIO();
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;

    ImGui::StyleColorsDark();
    ImNodes::StyleColorsDark();
    ImNodes::GetStyle().NodePadding = ImVec2(8, 8);

    ImGui_ImplGlfw_InitForOpenGL(window, true);
    ImGui_ImplOpenGL3_Init("#version 330");

    NodeEditorApp app;

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

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

        app.HandleKeyboardShortcuts();
        app.Render();
        app.RenderFileDialogs();

        ImGui::Render();

        int dispW, dispH;
        glfwGetFramebufferSize(window, &dispW, &dispH);
        glViewport(0, 0, dispW, dispH);
        glClearColor(0.15f, 0.15f, 0.15f, 1.0f);
        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;
}