#include "gfx/code_to_generator.h"
#include "gfx/uniform_parser.h"
#include <regex>
#include <sstream>
#include <string>

namespace fn::gfx {

// Replace every occurrence of `from` with `to` in `s`.
static void replace_all(std::string& s, const std::string& from, const std::string& to) {
    if (from.empty()) return;
    size_t p = 0;
    while ((p = s.find(from, p)) != std::string::npos) {
        s.replace(p, from.size(), to);
        p += to.size();
    }
}

// Extract the body of `void main() { ... }` (content between the outermost braces).
// Returns empty string if no balanced body is found.
static std::string extract_main_body(const std::string& source) {
    static const std::regex main_re(R"(\bvoid\s+main\s*\(\s*\)\s*\{)");
    std::smatch m;
    if (!std::regex_search(source, m, main_re)) return "";

    size_t start = static_cast<size_t>(m.position(0)) + m.length(0); // after the '{'
    int depth = 1;
    size_t end = start;
    for (; end < source.size() && depth > 0; ++end) {
        if (source[end] == '{') ++depth;
        else if (source[end] == '}') --depth;
        if (depth == 0) break;
    }
    if (depth != 0) return "";
    return source.substr(start, end - start);
}

// Strip lines that match `vec2 uv = <expr>;` (the function already provides uv).
static std::string strip_uv_decl(const std::string& body) {
    static const std::regex uv_re(R"(^\s*vec2\s+uv\s*=\s*[^;]+;\s*$)");
    std::ostringstream out;
    std::istringstream in(body);
    std::string line;
    bool first = true;
    while (std::getline(in, line)) {
        if (std::regex_match(line, uv_re)) continue;
        if (!first) out << '\n';
        out << line;
        first = false;
    }
    return out.str();
}

// Replace `fragColor = X;` with `return X;`. Tolerant to whitespace; one-shot regex.
static std::string fragcolor_to_return(const std::string& body) {
    static const std::regex fc_re(R"(\bfragColor\s*=\s*([^;]+);)");
    return std::regex_replace(body, fc_re, "return $1;");
}

// GLSL type info for a uniform: alias declaration, control kind, slot count.
struct TypeInfo {
    bool        ok = true;
    const char* alias_type   = "float"; // GLSL keyword
    const char* alias_swizzle = ".x";   // suffix on u_params[k]
    DagControl::Kind ctrl_kind = DagControl::Kind::Slider;
    int components = 1;                  // floats consumed
};

static TypeInfo type_info_for(GLSLType t, WidgetKind w) {
    TypeInfo ti;
    switch (t) {
        case GLSLType::Float:
            ti.alias_type    = "float";
            ti.alias_swizzle = ".x";
            ti.ctrl_kind     = DagControl::Kind::Slider;
            ti.components    = 1;
            break;
        case GLSLType::Int:
            ti.alias_type    = "int";
            ti.alias_swizzle = ".x";
            ti.ctrl_kind     = DagControl::Kind::Slider;
            ti.components    = 1;
            break;
        case GLSLType::Bool:
            ti.alias_type    = "bool";
            ti.alias_swizzle = ".x";
            ti.ctrl_kind     = DagControl::Kind::Slider; // 0/1 slider
            ti.components    = 1;
            break;
        case GLSLType::Vec2:
            ti.alias_type    = "vec2";
            ti.alias_swizzle = ".xy";
            ti.ctrl_kind     = DagControl::Kind::XY;
            ti.components    = 2;
            break;
        case GLSLType::Vec3:
            ti.alias_type    = "vec3";
            ti.alias_swizzle = ".xyz";
            ti.ctrl_kind     = DagControl::Kind::Color;
            ti.components    = 3;
            break;
        case GLSLType::Vec4:
            // Treat as Color if widget says so, otherwise unsupported (KISS).
            if (w == WidgetKind::Color) {
                ti.alias_type    = "vec4";
                ti.alias_swizzle = "";
                ti.ctrl_kind     = DagControl::Kind::Color;
                ti.components    = 3; // we render only RGB; alpha sits unused
            } else {
                ti.ok = false;
            }
            break;
    }
    // Widget overrides (when annotated)
    if (ti.ok) {
        if (w == WidgetKind::Color && t != GLSLType::Vec3 && t != GLSLType::Vec4) {
            ti.ctrl_kind = DagControl::Kind::Slider; // can't color a non-vec3/4
        }
        if (w == WidgetKind::Toggle) {
            ti.ctrl_kind = DagControl::Kind::Slider;
        }
        if (w == WidgetKind::XY && t == GLSLType::Vec2) {
            ti.ctrl_kind = DagControl::Kind::XY;
        }
    }
    return ti;
}

// Special-case: a vec4 typed as Color has alias `.xyz` for the local declared as vec3?
// We keep the alias as the full vec4 to avoid changing the user's body, and treat the
// control as Color over the first 3 components.
static const char* alias_swizzle_for(GLSLType t) {
    switch (t) {
        case GLSLType::Float:
        case GLSLType::Int:
        case GLSLType::Bool: return ".x";
        case GLSLType::Vec2: return ".xy";
        case GLSLType::Vec3: return ".xyz";
        case GLSLType::Vec4: return "";
    }
    return "";
}

CodeToGeneratorResult code_to_generator(const std::string& source) {
    CodeToGeneratorResult r;

    auto descs = parse_uniforms(source);

    // Reject sampler2D and other unsupported forms by inspecting raw source.
    // parse_uniforms already drops sampler2D silently; warn the user explicitly.
    if (source.find("sampler2D") != std::string::npos) {
        r.err = "sampler2D uniforms are not supported in saved generators";
        return r;
    }

    if (descs.empty() && source.find("uniform") != std::string::npos) {
        // Has uniforms but parser found none → unsupported types.
        // (Could be more granular; for now just warn.)
    }

    if (descs.size() > 16) {
        r.err = "too many uniforms (max 16 per generator)";
        return r;
    }

    std::string body_main = extract_main_body(source);
    if (body_main.empty()) {
        r.err = "could not find `void main() { ... }` body";
        return r;
    }
    if (body_main.find("fragColor") == std::string::npos) {
        r.err = "body must assign `fragColor` at least once";
        return r;
    }

    // Build aliases (preamble) and DagControls.
    std::ostringstream pre;
    int slot = 0;
    for (const auto& d : descs) {
        TypeInfo ti = type_info_for(d.glsl_type, d.widget);
        if (!ti.ok) {
            r.err = "uniform '" + d.name + "': vec4 only supported with @color";
            return r;
        }

        // Local alias: float u_speed = u_params[__BASE__+slot].x;
        pre << "    " << ti.alias_type << ' ' << d.name
            << " = " << ti.alias_type << "(u_params[__BASE__+" << slot << "]"
            << alias_swizzle_for(d.glsl_type) << ");\n";

        // Defaults: pack components into the 4 floats of this slot
        const int base_float = slot * 4;
        for (int c = 0; c < ti.components && c < 4; ++c) {
            r.param_defaults.push_back(d.defaults[c]);
        }
        for (int c = ti.components; c < 4; ++c) {
            r.param_defaults.push_back(0.0f);
        }
        // Param names: NAME, NAME.y, NAME.z... (for debug only)
        r.param_names.push_back(d.name);
        for (int c = 1; c < 4; ++c) r.param_names.push_back("");

        // Build DagControl
        DagControl ctrl;
        ctrl.kind  = ti.ctrl_kind;
        ctrl.label = d.name;
        ctrl.min   = d.min[0];
        ctrl.max   = d.max[0];
        ctrl.step  = d.step;
        if (ctrl.kind == DagControl::Kind::Slider) {
            ctrl.param_idx = { base_float, -1, -1 };
        } else if (ctrl.kind == DagControl::Kind::XY) {
            ctrl.param_idx = { base_float, base_float + 1, -1 };
        } else if (ctrl.kind == DagControl::Kind::Color) {
            ctrl.param_idx = { base_float, base_float + 1, base_float + 2 };
            ctrl.min = 0.0f;
            ctrl.max = 1.0f;
        }
        r.controls.push_back(std::move(ctrl));

        ++slot;
    }

    r.param_count = static_cast<int>(r.param_defaults.size());

    // Transform body: strip uv decl + fragColor → return.
    std::string body = strip_uv_decl(body_main);
    body = fragcolor_to_return(body);

    // Final body_template: preamble (aliases) + transformed body.
    std::ostringstream tpl;
    tpl << pre.str() << body;
    r.body_template = tpl.str();
    r.ok = true;
    return r;
}

DagNodeDef make_generator_def(const std::string& name,
                              const std::string& label,
                              const std::string& desc,
                              const CodeToGeneratorResult& tr) {
    DagNodeDef d;
    d.name           = name;
    d.label          = label.empty() ? name : label;
    d.desc           = desc;
    d.kind           = DagKind::Gen;
    d.num_inputs     = 0;
    d.param_names    = tr.param_names;
    d.param_defaults = tr.param_defaults;
    d.controls       = tr.controls;
    d.is_builtin     = false;

    const std::string body_template = tr.body_template;
    d.body_glsl = [body_template](int base) -> std::string {
        std::string s = body_template;
        replace_all(s, "__BASE__", std::to_string(base));
        return s;
    };
    return d;
}

} // namespace fn::gfx

#ifdef CODE_TO_GENERATOR_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. Plasma-like Code: float speed + vec3 color → 2 controls (slider + color)
    {
        const char* src = R"glsl(
uniform float u_speed; // @slider min=0.1 max=5 default=1
uniform vec3  u_color; // @color default=0.5,0.2,0.8

void main() {
    vec2 uv = gl_FragCoord.xy / u_resolution;
    float t = u_time * u_speed;
    vec3 c = u_color * (0.5 + 0.5 * cos(t + uv.xyx + vec3(0.0, 2.0, 4.0)));
    fragColor = vec4(c, 1.0);
}
)glsl";
        auto r = code_to_generator(src);
        assert(r.ok);
        assert(r.controls.size() == 2);
        assert(r.controls[0].kind == DagControl::Kind::Slider);
        assert(r.controls[0].label == "u_speed");
        assert(r.controls[0].param_idx[0] == 0);
        assert(r.controls[0].max == 5.0f);
        assert(r.controls[1].kind == DagControl::Kind::Color);
        assert(r.controls[1].label == "u_color");
        assert(r.controls[1].param_idx[0] == 4);
        assert(r.controls[1].param_idx[2] == 6);
        assert(r.param_count == 8); // 2 uniforms × 4 floats
        assert(r.param_defaults.size() == 8);
        assert(r.param_defaults[0] == 1.0f);   // u_speed default
        assert(r.param_defaults[4] == 0.5f);   // u_color.r
        assert(r.param_defaults[5] == 0.2f);
        assert(r.param_defaults[6] == 0.8f);

        // body_template: should declare aliases referencing __BASE__
        assert(contains(r.body_template, "float u_speed"));
        assert(contains(r.body_template, "u_params[__BASE__+0]"));
        assert(contains(r.body_template, "vec3 u_color"));
        assert(contains(r.body_template, "u_params[__BASE__+1]"));
        // uv = ... line stripped
        assert(!contains(r.body_template, "vec2 uv = gl_FragCoord"));
        // fragColor → return
        assert(contains(r.body_template, "return vec4(c, 1.0);"));
        assert(!contains(r.body_template, "fragColor ="));
    }

    // 2. make_generator_def + lambda substitution
    {
        const char* src = R"glsl(
uniform float u_x; // @slider min=0 max=1 default=0.5
void main() {
    fragColor = vec4(u_x, 0.0, 0.0, 1.0);
}
)glsl";
        auto tr = code_to_generator(src);
        assert(tr.ok);
        DagNodeDef def = make_generator_def("test_gen", "Test", "desc", tr);
        assert(!def.is_builtin);
        assert(def.kind == DagKind::Gen);
        assert(def.num_inputs == 0);
        std::string b3 = def.body_glsl(3);
        assert(contains(b3, "u_params[3+0]"));
        assert(!contains(b3, "__BASE__"));
        std::string b0 = def.body_glsl(0);
        assert(contains(b0, "u_params[0+0]"));
    }

    // 3. Missing void main → error
    {
        auto r = code_to_generator("uniform float u_x;\n");
        assert(!r.ok);
        assert(r.err.find("void main") != std::string::npos);
    }

    // 4. No fragColor → error
    {
        auto r = code_to_generator("void main() { }\n");
        assert(!r.ok);
        assert(r.err.find("fragColor") != std::string::npos);
    }

    // 5. sampler2D rejected
    {
        const char* src = R"glsl(
uniform sampler2D u_tex;
void main() { fragColor = vec4(1.0); }
)glsl";
        auto r = code_to_generator(src);
        assert(!r.ok);
        assert(r.err.find("sampler2D") != std::string::npos);
    }

    // 6. vec2 with @xy → XY control
    {
        const char* src = R"glsl(
uniform vec2 u_pos; // @xy min=-1 max=1 default=0,0
void main() { fragColor = vec4(u_pos, 0.0, 1.0); }
)glsl";
        auto r = code_to_generator(src);
        assert(r.ok);
        assert(r.controls.size() == 1);
        assert(r.controls[0].kind == DagControl::Kind::XY);
        assert(r.controls[0].param_idx[0] == 0);
        assert(r.controls[0].param_idx[1] == 1);
        assert(r.controls[0].min == -1.0f);
        assert(r.controls[0].max == 1.0f);
        assert(contains(r.body_template, "vec2 u_pos"));
        assert(contains(r.body_template, "u_params[__BASE__+0].xy"));
    }

    // 7. No uniforms at all → empty controls, body_template still ok
    {
        const char* src = R"glsl(
void main() {
    vec2 uv = gl_FragCoord.xy / u_resolution;
    fragColor = vec4(uv, 0.0, 1.0);
}
)glsl";
        auto r = code_to_generator(src);
        assert(r.ok);
        assert(r.controls.empty());
        assert(r.param_count == 0);
        assert(contains(r.body_template, "return vec4(uv, 0.0, 1.0);"));
    }

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