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

namespace fn::gfx {

static constexpr int MAX_NODES = 16;

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[16];\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();
    }

    // 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(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);
                    }
                }
            }
            if (last_valid_out < 0) return "vec4(0.0, 0.0, 0.0, 1.0)";
            return "out_" + std::to_string(last_valid_out);
        };

        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;
    }

    // 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"; };

    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 if (last_valid_out >= 0) {
        out << "    fragColor = out_" << last_valid_out << ";\n";
    } else {
        seed();
    }

    out << "}\n";

    return out.str();
}

} // 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 → fragColor = out_0
    {
        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 node_0"));
        assert(contains(s, "vec4 out_0 = node_0("));
        assert(contains(s, "fragColor = out_0"));
    }

    // 3. Gen + Op → Op uses out_0 as input a
    {
        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);
        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);
        auto s = compile_dag_to_glsl(p);
        assert(contains(s, "out_2 = node_2(out_0, out_1, uv)"));
    }

    // 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"));
    }

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