#include "gfx/dag_compile.h"
#include "gfx/dag_catalog.h"
#include <algorithm>
#include <regex>
#include <sstream>

namespace fn::gfx {

static constexpr int MAX_NODES       = 16;
static constexpr int MAX_PARAM_VEC4S = 64; // 256 floats — enough for 16 nodes × ~16 floats each

std::vector<int> dag_param_layout(const std::vector<DagStep>& pipeline) {
    std::vector<int> base(pipeline.size(), 0);
    int cursor = 0;
    for (size_t i = 0; i < pipeline.size(); ++i) {
        const DagNodeDef* def = dag_find(pipeline[i].name);
        int pc = def ? static_cast<int>(def->param_defaults.size()) : 0;
        base[i] = cursor;
        cursor += dag_vec4_count(pc);
    }
    return base;
}

std::string compile_dag_to_glsl(const std::vector<DagStep>& pipeline) {
    const int n = static_cast<int>(std::min(pipeline.size(), static_cast<size_t>(MAX_NODES)));
    std::ostringstream out;

    out << "uniform vec4 u_params[" << MAX_PARAM_VEC4S << "];\n";
    out << "uniform int u_preview_target; // -1 = real Output; >=0 = show out_<i>\n\n";

    if (n == 0) {
        out << "void main() {\n";
        out << "    vec2 uv = gl_FragCoord.xy / u_resolution;\n";
        out << "    (void)uv;\n";
        out << "    fragColor = vec4(0.04, 0.04, 0.06, 1.0);\n";
        out << "}\n";
        return out.str();
    }

    std::vector<int> base = dag_param_layout(pipeline);

    // Emit per-node functions (skip Output: it's a sink, no body)
    for (int i = 0; i < n; ++i) {
        const DagStep& step = pipeline[static_cast<size_t>(i)];
        const DagNodeDef* def = dag_find(step.name);
        if (!def) continue;
        if (def->kind == DagKind::Output) continue;

        int ni = def->num_inputs;
        out << "vec4 node_" << i << "(";
        if (ni >= 1) out << "vec4 a";
        if (ni >= 2) out << ", vec4 b";
        if (ni >= 3) out << ", vec4 c";
        if (ni >= 4) out << ", vec4 d";
        if (ni > 0)  out << ", ";
        out << "vec2 uv) {\n";
        out << def->body_glsl(base[static_cast<size_t>(i)]) << "\n";
        out << "}\n\n";
    }

    out << "void main() {\n";
    out << "    vec2 uv = gl_FragCoord.xy / u_resolution;\n";

    int last_valid_out = -1;
    int output_idx = -1;

    for (int i = 0; i < n; ++i) {
        const DagStep& step = pipeline[static_cast<size_t>(i)];
        const DagNodeDef* def = dag_find(step.name);
        if (!def) continue;
        if (def->kind == DagKind::Output) { output_idx = i; continue; }

        int ni = def->num_inputs;

        auto resolve = [&](int k) -> std::string {
            const std::string& sid = step.source_ids[static_cast<size_t>(k)];
            if (!sid.empty()) {
                for (int j = 0; j < i; ++j) {
                    if (pipeline[static_cast<size_t>(j)].id == sid) {
                        const DagNodeDef* jdef = dag_find(pipeline[static_cast<size_t>(j)].name);
                        if (jdef && jdef->kind == DagKind::Output) continue; // Output has no out_j
                        return "out_" + std::to_string(j);
                    }
                }
            }
            // Op/Blend with no source on this slot → black input (cannot fall back to
            // last_valid_out: that's how nodes "leak" into the canvas without being wired).
            return "vec4(0.0, 0.0, 0.0, 1.0)";
        };

        out << "    vec4 out_" << i << " = node_" << i << "(";
        for (int k = 0; k < ni; ++k) {
            if (k > 0) out << ", ";
            out << resolve(k);
        }
        if (ni > 0) out << ", ";
        out << "uv);\n";

        last_valid_out = i;
    }
    (void)last_valid_out;

    // Preview branch: if u_preview_target points to a valid out_<i>, emit it
    // and bail out before the Output-driven fragColor.
    for (int i = 0; i < n; ++i) {
        const DagStep& step = pipeline[static_cast<size_t>(i)];
        const DagNodeDef* def = dag_find(step.name);
        if (!def) continue;
        if (def->kind == DagKind::Output) continue;
        out << "    if (u_preview_target == " << i << ") { fragColor = out_" << i << "; return; }\n";
    }

    // Resolve fragColor: if there's an Output node with a connection, use that; else fallback.
    auto seed = [&]() { out << "    fragColor = vec4(0.04, 0.04, 0.06, 1.0);\n"; };

    // Strict policy: only emit what is wired into the Output node. With no
    // Output present, or with Output left disconnected, paint the seed color —
    // never silently fall back to the last evaluated node.
    if (output_idx >= 0) {
        const std::string& sid = pipeline[static_cast<size_t>(output_idx)].source_ids[0];
        int src = -1;
        if (!sid.empty()) {
            for (int j = 0; j < output_idx; ++j) {
                if (pipeline[static_cast<size_t>(j)].id == sid) { src = j; break; }
            }
        }
        if (src >= 0) out << "    fragColor = out_" << src << ";\n";
        else           seed();
    } else {
        seed();
    }

    out << "}\n";

    return out.str();
}

std::string compile_dag_to_glsl_baked(const std::vector<DagStep>& pipeline) {
    std::string s = compile_dag_to_glsl(pipeline);

    // Compute total vec4 slots actually used by the pipeline.
    auto base = dag_param_layout(pipeline);
    int total = 0;
    for (size_t i = 0; i < pipeline.size(); ++i) {
        const DagNodeDef* def = dag_find(pipeline[i].name);
        int pc = def ? static_cast<int>(def->param_defaults.size()) : 0;
        int v  = dag_vec4_count(pc);
        if (base[i] + v > total) total = base[i] + v;
    }
    if (total == 0) total = 1; // GLSL forbids zero-sized arrays

    // Pack current params into a flat float array (same layout as dag_uniforms_apply).
    std::vector<float> data(static_cast<size_t>(total * 4), 0.0f);
    for (size_t i = 0; i < pipeline.size(); ++i) {
        const DagNodeDef* def = dag_find(pipeline[i].name);
        if (!def) continue;
        int pc = static_cast<int>(def->param_defaults.size());
        int b  = base[i] * 4;
        for (int k = 0; k < pc && k < static_cast<int>(pipeline[i].params.size()); ++k) {
            data[static_cast<size_t>(b + k)] = pipeline[i].params[static_cast<size_t>(k)];
        }
    }

    // Build `const vec4 u_params[N] = vec4[N](vec4(...), ...);`
    std::ostringstream init;
    init << "const vec4 u_params[" << total << "] = vec4[" << total << "](";
    for (int i = 0; i < total; ++i) {
        if (i > 0) init << ", ";
        init << "vec4("
             << data[static_cast<size_t>(i * 4 + 0)] << ", "
             << data[static_cast<size_t>(i * 4 + 1)] << ", "
             << data[static_cast<size_t>(i * 4 + 2)] << ", "
             << data[static_cast<size_t>(i * 4 + 3)] << ")";
    }
    init << ");";

    // Replace the uniform u_params declaration with the const array.
    static const std::regex up_re(R"(uniform\s+vec4\s+u_params\[\d+\];)");
    s = std::regex_replace(s, up_re, init.str());

    // Replace the u_preview_target uniform with a const = -1 (kills the preview branches).
    static const std::regex pt_re(R"(uniform\s+int\s+u_preview_target;[^\n]*)");
    s = std::regex_replace(s, pt_re, "const int u_preview_target = -1;");

    return s;
}

} // namespace fn::gfx

#ifdef DAG_COMPILE_TEST
#include <cassert>
#include <cstdio>

static bool contains(const std::string& hay, const std::string& needle) {
    return hay.find(needle) != std::string::npos;
}

int main() {
    using namespace fn::gfx;

    // 1. Empty pipeline → seed color
    {
        std::vector<DagStep> p;
        auto s = compile_dag_to_glsl(p);
        assert(contains(s, "fragColor = vec4(0.04"));
    }

    // 2. Single Gen + Output wired → fragColor = out_0
    {
        std::vector<DagStep> p;
        DagStep g; g.id = "a"; g.name = "plasma"; p.push_back(g);
        DagStep o; o.id = "out"; o.name = "output"; o.source_ids[0] = "a"; p.push_back(o);
        auto s = compile_dag_to_glsl(p);
        assert(contains(s, "vec4 node_0"));
        assert(contains(s, "vec4 out_0 = node_0("));
        assert(contains(s, "fragColor = out_0"));
    }

    // 3. Gen + Op + Output → Op uses out_0 as input, fragColor = out_1
    {
        std::vector<DagStep> p;
        DagStep g; g.id = "a"; g.name = "plasma"; p.push_back(g);
        DagStep o; o.id = "b"; o.name = "invert"; o.source_ids[0] = "a"; p.push_back(o);
        DagStep f; f.id = "out"; f.name = "output"; f.source_ids[0] = "b"; p.push_back(f);
        auto s = compile_dag_to_glsl(p);
        assert(contains(s, "out_1 = node_1(out_0, uv)"));
        assert(contains(s, "fragColor = out_1"));
    }

    // 4. Blend with multi-source → both inputs resolved
    {
        std::vector<DagStep> p;
        DagStep a; a.id = "a"; a.name = "plasma"; p.push_back(a);
        DagStep b; b.id = "b"; b.name = "solid";  p.push_back(b);
        DagStep m; m.id = "m"; m.name = "blend_mix";
        m.source_ids[0] = "a"; m.source_ids[1] = "b";
        p.push_back(m);
        DagStep o; o.id = "out"; o.name = "output"; o.source_ids[0] = "m"; p.push_back(o);
        auto s = compile_dag_to_glsl(p);
        assert(contains(s, "out_2 = node_2(out_0, out_1, uv)"));
        assert(contains(s, "fragColor = out_2"));
    }

    // 4b. Strict mode: nodes without Output → seed (never leaks last node).
    // Note: the preview branch emits `if (u_preview_target == i) fragColor = out_i;`
    // which we don't penalise; what matters is the *final* fragColor (after the
    // preview ifs) — that must be the seed, not a node output.
    {
        std::vector<DagStep> p;
        DagStep g; g.id = "a"; g.name = "plasma"; p.push_back(g);
        auto s = compile_dag_to_glsl(p);
        assert(contains(s, "vec4 out_0 = node_0("));     // node still emitted
        // The seed line must appear *after* the last preview branch
        size_t seed_pos    = s.rfind("fragColor = vec4(0.04");
        size_t preview_pos = s.rfind("u_preview_target ==");
        assert(seed_pos != std::string::npos);
        assert(preview_pos == std::string::npos || seed_pos > preview_pos);
    }

    // 5. Output node drives fragColor from its source, not from last index
    {
        std::vector<DagStep> p;
        DagStep g1; g1.id = "g1"; g1.name = "plasma"; p.push_back(g1);
        DagStep g2; g2.id = "g2"; g2.name = "solid";  p.push_back(g2);
        DagStep o;  o.id  = "o";  o.name  = "output";
        o.source_ids[0] = "g1";  // connect Output to the plasma (first gen), not last
        p.push_back(o);
        auto s = compile_dag_to_glsl(p);
        // Output must NOT emit a node_2 function
        assert(!contains(s, "vec4 node_2("));
        // fragColor must come from out_0 (plasma), not out_1 (solid)
        assert(contains(s, "fragColor = out_0"));
    }

    // 6. Output with no connection → seed fallback
    {
        std::vector<DagStep> p;
        DagStep g; g.id = "g"; g.name = "plasma"; p.push_back(g);
        DagStep o; o.id = "o"; o.name = "output"; p.push_back(o); // no source
        auto s = compile_dag_to_glsl(p);
        assert(contains(s, "fragColor = vec4(0.04"));
    }

    // 7. Baked variant: const arrays, no uniforms u_params / u_preview_target
    {
        std::vector<DagStep> p;
        DagStep g; g.id = "a"; g.name = "plasma"; g.params = {2.0f, 3.0f}; p.push_back(g);
        DagStep o; o.id = "out"; o.name = "output"; o.source_ids[0] = "a"; p.push_back(o);
        auto s = compile_dag_to_glsl_baked(p);
        assert(!contains(s, "uniform vec4 u_params"));
        assert(!contains(s, "uniform int u_preview_target"));
        assert(contains(s, "const vec4 u_params["));
        assert(contains(s, "vec4(2"));               // baked first param
        assert(contains(s, "const int u_preview_target = -1"));
        assert(contains(s, "fragColor = out_0"));
    }

    std::printf("dag_compile: 8/8 asserts passed\n");
    return 0;
}
#endif