#include "gfx/dag_node_editor.h"
#include "gfx/dag_catalog.h"
#include "imgui.h"
#include "imgui_node_editor.h"
#include <algorithm>
#include <cstdio>
#include <functional>
#include <queue>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace ed = ax::NodeEditor;

namespace fn::gfx {

static constexpr int MAX_NODES = 16;

static ed::EditorContext*         s_ctx = nullptr;
static uint32_t                   s_next_uid = 1;
static std::unordered_set<uint32_t> s_positioned;

// ── ID encoding ──────────────────────────────────────────────────────────────
// node id   = editor_uid
// output pin id = (editor_uid << 8) | 0
// input pin id  = (editor_uid << 8) | (slot + 1)   slot 0..3
// link id       = (from_node_uid << 20) | (to_node_uid << 8) | slot

static uintptr_t node_id(uint32_t uid) {
    return static_cast<uintptr_t>(uid);
}
static uintptr_t output_pin_id(uint32_t uid) {
    return (static_cast<uintptr_t>(uid) << 8) | 0u;
}
static uintptr_t input_pin_id(uint32_t uid, int slot) {
    return (static_cast<uintptr_t>(uid) << 8) | static_cast<uintptr_t>(slot + 1);
}
static uintptr_t link_id(uint32_t from_uid, uint32_t to_uid, int slot) {
    return (static_cast<uintptr_t>(from_uid) << 20)
         | (static_cast<uintptr_t>(to_uid)   <<  8)
         | static_cast<uintptr_t>(slot);
}

// Decode pin id back
static bool is_output_pin(uintptr_t id) { return (id & 0xFF) == 0; }
static uint32_t uid_from_pin(uintptr_t id) { return static_cast<uint32_t>(id >> 8); }
static int slot_from_input_pin(uintptr_t id) { return static_cast<int>(id & 0xFF) - 1; }

static int find_by_uid(const std::vector<DagStep>& p, uint32_t uid) {
    for (int i = 0; i < static_cast<int>(p.size()); ++i) {
        if (p[static_cast<size_t>(i)].editor_uid == uid) return i;
    }
    return -1;
}
static int find_by_id(const std::vector<DagStep>& p, const std::string& id) {
    for (int i = 0; i < static_cast<int>(p.size()); ++i) {
        if (p[static_cast<size_t>(i)].id == id) return i;
    }
    return -1;
}

static ImVec4 kind_color(DagKind kind) {
    switch (kind) {
        case DagKind::Gen:    return ImVec4(0.25f, 0.55f, 0.90f, 1.0f);
        case DagKind::Op:     return ImVec4(0.65f, 0.40f, 0.90f, 1.0f);
        case DagKind::Blend:  return ImVec4(0.90f, 0.65f, 0.15f, 1.0f);
        case DagKind::Output: return ImVec4(0.85f, 0.25f, 0.25f, 1.0f);
    }
    return ImVec4(1, 1, 1, 1);
}

// Topological sort using Kahn's algorithm.
// Returns false if a cycle is detected (should not happen — BeginCreate rejects cycles).
static bool topo_sort(std::vector<DagStep>& pipeline) {
    int n = static_cast<int>(pipeline.size());
    if (n == 0) return true;

    // Build adjacency: edge from src -> dst if dst.source_ids[k] == src.id
    std::unordered_map<std::string, int> id_to_idx;
    for (int i = 0; i < n; ++i) id_to_idx[pipeline[static_cast<size_t>(i)].id] = i;

    std::vector<int> in_degree(static_cast<size_t>(n), 0);
    std::vector<std::vector<int>> adj(static_cast<size_t>(n));

    for (int i = 0; i < n; ++i) {
        const DagStep& s = pipeline[static_cast<size_t>(i)];
        const DagNodeDef* def = dag_find(s.name);
        int ni = def ? def->num_inputs : 0;
        for (int k = 0; k < ni; ++k) {
            const std::string& sid = s.source_ids[static_cast<size_t>(k)];
            if (sid.empty()) continue;
            auto it = id_to_idx.find(sid);
            if (it == id_to_idx.end()) continue;
            int src = it->second;
            adj[static_cast<size_t>(src)].push_back(i);
            in_degree[static_cast<size_t>(i)]++;
        }
    }

    std::queue<int> q;
    for (int i = 0; i < n; ++i) {
        if (in_degree[static_cast<size_t>(i)] == 0) q.push(i);
    }

    std::vector<int> order;
    order.reserve(static_cast<size_t>(n));
    while (!q.empty()) {
        int u = q.front(); q.pop();
        order.push_back(u);
        for (int v : adj[static_cast<size_t>(u)]) {
            if (--in_degree[static_cast<size_t>(v)] == 0) q.push(v);
        }
    }

    if (static_cast<int>(order.size()) != n) return false; // cycle

    std::vector<DagStep> sorted;
    sorted.reserve(static_cast<size_t>(n));
    for (int idx : order) sorted.push_back(pipeline[static_cast<size_t>(idx)]);

    // Force Output nodes to the back so all their dependencies live at smaller
    // indices (compiler and cycle validator search strictly before the target).
    std::stable_partition(sorted.begin(), sorted.end(), [](const DagStep& s) {
        const DagNodeDef* d = dag_find(s.name);
        return !(d && d->kind == DagKind::Output);
    });

    pipeline = std::move(sorted);
    return true;
}

bool dag_node_editor(std::vector<DagStep>& pipeline) {
    bool changed = false;

    // Ensure UIDs are assigned (e.g. for nodes created before this editor was active)
    for (auto& step : pipeline) {
        if (step.editor_uid == 0) {
            step.editor_uid = s_next_uid++;
        }
    }

    // Create context on first call
    if (!s_ctx) {
        ed::Config cfg;
        cfg.SettingsFile = nullptr; // no disk persistence
        s_ctx = ed::CreateEditor(&cfg);
    }

    // Palette drop: detect without a capturing target (an InvisibleButton would
    // steal clicks from the node editor). Observe the active drag-drop payload
    // and, if the mouse is over this window and the user releases LMB, queue an
    // add at that canvas position.
    static std::string s_pending_add_name;
    static ImVec2      s_pending_add_pos(0, 0);
    static bool        s_pending_add = false;

    const bool window_hovered = ImGui::IsWindowHovered(
        ImGuiHoveredFlags_ChildWindows | ImGuiHoveredFlags_AllowWhenBlockedByActiveItem);
    if (window_hovered) {
        if (const ImGuiPayload* p = ImGui::GetDragDropPayload()) {
            if (p->IsDataType("DAG_NODE_TYPE") && ImGui::IsMouseReleased(ImGuiMouseButton_Left)) {
                s_pending_add_name.assign(static_cast<const char*>(p->Data),
                                          static_cast<size_t>(p->DataSize));
                s_pending_add_pos = ImGui::GetMousePos();
                s_pending_add = true;
            }
        }
    }

    ed::SetCurrentEditor(s_ctx);
    ed::Begin("dag_editor", ImVec2(0.0f, 0.0f));

    if (s_pending_add) {
        const DagNodeDef* def = dag_find(s_pending_add_name);
        if (def && static_cast<int>(pipeline.size()) < MAX_NODES) {
            uint32_t uid = s_next_uid++;
            DagStep step;
            step.id     = "n" + std::to_string(uid);
            step.name   = def->name;
            step.params = def->param_defaults;
            step.editor_uid = uid;
            ImVec2 canvas_pos = ed::ScreenToCanvas(s_pending_add_pos);
            step.editor_pos_x = canvas_pos.x;
            step.editor_pos_y = canvas_pos.y;
            // Insert before the Output node so the Output stays at the back;
            // otherwise new nodes can never be wired into it (compiler and
            // cycle check only search indices strictly before the target).
            auto insert_it = pipeline.end();
            for (auto it = pipeline.begin(); it != pipeline.end(); ++it) {
                const DagNodeDef* d = dag_find(it->name);
                if (d && d->kind == DagKind::Output) { insert_it = it; break; }
            }
            pipeline.insert(insert_it, step);
            changed = true;
        }
        s_pending_add = false;
    }

    // ── Draw nodes ───────────────────────────────────────────────────────────
    for (int i = 0; i < static_cast<int>(pipeline.size()); ++i) {
        DagStep& step = pipeline[static_cast<size_t>(i)];
        const DagNodeDef* def = dag_find(step.name);
        if (!def) continue;

        ImVec4 col = kind_color(def->kind);

        // Initial position only — after that, the editor owns the node's position
        // and user drags must not be overwritten.
        if (s_positioned.find(step.editor_uid) == s_positioned.end()) {
            if (step.editor_pos_x == 0.0f && step.editor_pos_y == 0.0f) {
                step.editor_pos_x = 50.0f + static_cast<float>(i) * 220.0f;
                step.editor_pos_y = 100.0f;
            }
            ed::SetNodePosition(ed::NodeId(node_id(step.editor_uid)),
                                ImVec2(step.editor_pos_x, step.editor_pos_y));
            s_positioned.insert(step.editor_uid);
        }

        ed::BeginNode(ed::NodeId(node_id(step.editor_uid)));

        // Header
        ImGui::PushStyleColor(ImGuiCol_Text, col);
        ImGui::TextUnformatted(def->label.c_str());
        ImGui::PopStyleColor();
        ImGui::Dummy(ImVec2(0, 2));

        // Input pins (left column) + controls (middle) + output pin (right)
        // We use a simple layout: input pins vertically, then controls, then output pin
        int ni = def->num_inputs;

        // Input pins
        for (int k = 0; k < ni; ++k) {
            ed::BeginPin(ed::PinId(input_pin_id(step.editor_uid, k)), ed::PinKind::Input);
            char lbl[16];
            std::snprintf(lbl, sizeof(lbl), "-> in%d", k);
            ImGui::TextUnformatted(lbl);
            ed::EndPin();
        }

        if (ni == 0) {
            // Gen: no inputs, push a small placeholder so layout is consistent
            ImGui::Dummy(ImVec2(4, 4));
        }

        // Controls
        ImGui::PushID(static_cast<int>(step.editor_uid));
        for (size_t ci = 0; ci < def->controls.size(); ++ci) {
            const DagControl& ctrl = def->controls[ci];
            ImGui::SetNextItemWidth(150.0f);
            char uid_lbl[64];
            std::snprintf(uid_lbl, sizeof(uid_lbl), "%s##%u%zu", ctrl.label.c_str(), step.editor_uid, ci);
            if (ctrl.kind == DagControl::Kind::Slider) {
                int pidx = ctrl.param_idx[0];
                if (pidx >= 0 && pidx < 4) {
                    ImGui::SliderFloat(uid_lbl, &step.params[static_cast<size_t>(pidx)], ctrl.min, ctrl.max);
                }
            } else if (ctrl.kind == DagControl::Kind::XY) {
                int px = ctrl.param_idx[0], py = ctrl.param_idx[1];
                if (px >= 0 && px < 4 && py >= 0 && py < 4 && py == px + 1) {
                    ImGui::SliderFloat2(uid_lbl, &step.params[static_cast<size_t>(px)], ctrl.min, ctrl.max);
                }
            } else if (ctrl.kind == DagControl::Kind::Color) {
                int pr = ctrl.param_idx[0];
                if (pr >= 0 && pr + 2 < 4) {
                    // Suspend the node editor while the color picker popup is
                    // open so its clicks don't reach the canvas (and so the
                    // popup isn't clipped to the node bounds).
                    ed::Suspend();
                    ImGui::ColorEdit3(uid_lbl, &step.params[static_cast<size_t>(pr)],
                                      ImGuiColorEditFlags_NoInputs |
                                      ImGuiColorEditFlags_NoLabel   |
                                      ImGuiColorEditFlags_AlphaBar);
                    ed::Resume();
                }
            }
        }
        ImGui::PopID();

        // Output pin (skip for the terminal Output node — it has no output)
        if (def->kind != DagKind::Output) {
            ImGui::Dummy(ImVec2(0, 2));
            ed::BeginPin(ed::PinId(output_pin_id(step.editor_uid)), ed::PinKind::Output);
            ImGui::Text("out ->");
            ed::EndPin();
        }

        ed::EndNode();
    }

    // ── Draw existing links ──────────────────────────────────────────────────
    for (int i = 0; i < static_cast<int>(pipeline.size()); ++i) {
        const DagStep& step = pipeline[static_cast<size_t>(i)];
        const DagNodeDef* def = dag_find(step.name);
        if (!def) continue;
        for (int k = 0; k < def->num_inputs; ++k) {
            const std::string& sid = step.source_ids[static_cast<size_t>(k)];
            if (sid.empty()) continue;
            int src_idx = find_by_id(pipeline, sid);
            if (src_idx < 0) continue;
            uint32_t src_uid = pipeline[static_cast<size_t>(src_idx)].editor_uid;
            ed::Link(ed::LinkId(link_id(src_uid, step.editor_uid, k)),
                     ed::PinId(output_pin_id(src_uid)),
                     ed::PinId(input_pin_id(step.editor_uid, k)));
        }
    }

    // ── Handle link creation ─────────────────────────────────────────────────
    if (ed::BeginCreate()) {
        ed::PinId start_pin, end_pin;
        if (ed::QueryNewLink(&start_pin, &end_pin)) {
            uintptr_t sp = start_pin.Get();
            uintptr_t ep = end_pin.Get();

            // Normalise: start must be output, end must be input
            if (!is_output_pin(sp) && is_output_pin(ep)) {
                std::swap(sp, ep);
            }

            bool valid = is_output_pin(sp) && !is_output_pin(ep);
            if (valid) {
                uint32_t from_uid = uid_from_pin(sp);
                uint32_t to_uid   = uid_from_pin(ep);
                int slot          = slot_from_input_pin(ep);

                int from_idx = find_by_uid(pipeline, from_uid);
                int to_idx   = find_by_uid(pipeline, to_uid);

                // Real cycle check: would the new edge from->to introduce a path
                // from `from` back to itself? It does iff `from` already (transitively)
                // depends on `to`. We walk source_ids of `from` and reject if we
                // ever hit `to`. Vector index order is irrelevant — topo_sort runs
                // at end-of-frame and reorders the pipeline.
                bool cycle = false;
                if (from_uid == to_uid) {
                    cycle = true;
                } else if (from_idx >= 0 && to_idx >= 0) {
                    const std::string to_id = pipeline[static_cast<size_t>(to_idx)].id;
                    std::function<bool(const std::string&)> depends_on_to;
                    depends_on_to = [&](const std::string& node_id) -> bool {
                        if (node_id == to_id) return true;
                        int idx = -1;
                        for (int i = 0; i < static_cast<int>(pipeline.size()); ++i) {
                            if (pipeline[static_cast<size_t>(i)].id == node_id) { idx = i; break; }
                        }
                        if (idx < 0) return false;
                        const DagStep& s = pipeline[static_cast<size_t>(idx)];
                        const DagNodeDef* d = dag_find(s.name);
                        int n = d ? d->num_inputs : 0;
                        for (int k = 0; k < n; ++k) {
                            const std::string& sid = s.source_ids[static_cast<size_t>(k)];
                            if (!sid.empty() && depends_on_to(sid)) return true;
                        }
                        return false;
                    };
                    cycle = depends_on_to(pipeline[static_cast<size_t>(from_idx)].id);
                }

                if (cycle) {
                    ed::RejectNewItem(ImVec4(1, 0.3f, 0.3f, 1));
                    valid = false;
                }

                if (valid && from_idx >= 0 && to_idx >= 0 && slot >= 0) {
                    if (ed::AcceptNewItem()) {
                        pipeline[static_cast<size_t>(to_idx)].source_ids[static_cast<size_t>(slot)] =
                            pipeline[static_cast<size_t>(from_idx)].id;
                        changed = true;
                    }
                }
            } else {
                ed::RejectNewItem(ImVec4(0.5f, 0.5f, 0.5f, 1));
            }
        }
    }
    ed::EndCreate();

    // ── Handle deletion ──────────────────────────────────────────────────────
    if (ed::BeginDelete()) {
        ed::LinkId del_lid;
        while (ed::QueryDeletedLink(&del_lid)) {
            if (ed::AcceptDeletedItem()) {
                uintptr_t raw = del_lid.Get();
                uint32_t to_uid = static_cast<uint32_t>((raw >> 8) & 0xFFF);
                int      slot   = static_cast<int>(raw & 0xFF);
                int to_idx = find_by_uid(pipeline, to_uid);
                if (to_idx >= 0 && slot < 4) {
                    pipeline[static_cast<size_t>(to_idx)].source_ids[static_cast<size_t>(slot)].clear();
                    changed = true;
                }
            }
        }

        ed::NodeId del_nid;
        while (ed::QueryDeletedNode(&del_nid)) {
            uint32_t uid = static_cast<uint32_t>(del_nid.Get());
            int idx = find_by_uid(pipeline, uid);
            // Refuse to delete the Output node (it's the sink)
            const DagNodeDef* ddef = (idx >= 0) ? dag_find(pipeline[static_cast<size_t>(idx)].name) : nullptr;
            if (ddef && ddef->kind == DagKind::Output) {
                ed::RejectDeletedItem();
                continue;
            }
            if (ed::AcceptDeletedItem()) {
                if (idx >= 0) {
                    const std::string& del_step_id = pipeline[static_cast<size_t>(idx)].id;
                    for (auto& step : pipeline) {
                        for (auto& sid : step.source_ids) {
                            if (sid == del_step_id) sid.clear();
                        }
                    }
                    pipeline.erase(pipeline.begin() + idx);
                    changed = true;
                }
            }
        }
    }
    ed::EndDelete();

    // ── Save node positions back to steps ────────────────────────────────────
    for (auto& step : pipeline) {
        auto pos = ed::GetNodePosition(ed::NodeId(node_id(step.editor_uid)));
        step.editor_pos_x = pos.x;
        step.editor_pos_y = pos.y;
    }

    ed::End();
    ed::SetCurrentEditor(nullptr);

    // ── Topological sort after topology change ───────────────────────────────
    if (changed) {
        if (!topo_sort(pipeline)) {
            // Cycle detected — should not happen given BeginCreate validation
            // but log a warning and leave order as-is
            std::fprintf(stderr, "dag_node_editor: cycle detected, skipping topo sort\n");
        }
    }

    return changed;
}

void dag_node_editor_destroy() {
    if (s_ctx) {
        ed::DestroyEditor(s_ctx);
        s_ctx = nullptr;
    }
}

} // namespace fn::gfx