#include "pxl8_graph.h"

#include <math.h>
#include <stdlib.h>
#include <string.h>

#include "pxl8_math.h"

static inline u32 hash2d(i32 x, i32 y, u32 seed) {
    return pxl8_hash32((u32)x ^ ((u32)y * 2654435769u) ^ seed);
}

static inline f32 hash2d_f(i32 x, i32 y, u32 seed) {
    return (f32)hash2d(x, y, seed) / (f32)0xFFFFFFFF;
}

static f32 gradient2d(i32 ix, i32 iy, f32 fx, f32 fy, u32 seed) {
    u32 h = hash2d(ix, iy, seed);
    f32 angle = (f32)h / (f32)0xFFFFFFFF * 6.28318530718f;
    f32 gx = cosf(angle);
    f32 gy = sinf(angle);
    return gx * fx + gy * fy;
}

static inline f32 smoothstep(f32 t) {
    return t * t * (3.0f - 2.0f * t);
}

static inline f32 lerp(f32 a, f32 b, f32 t) {
    return a + t * (b - a);
}

static f32 noise_value(f32 x, f32 y, f32 scale, u32 seed) {
    f32 sx = x * scale;
    f32 sy = y * scale;
    i32 ix = (i32)floorf(sx);
    i32 iy = (i32)floorf(sy);
    f32 fx = sx - (f32)ix;
    f32 fy = sy - (f32)iy;
    f32 u = smoothstep(fx);
    f32 v = smoothstep(fy);

    f32 n00 = hash2d_f(ix, iy, seed);
    f32 n10 = hash2d_f(ix + 1, iy, seed);
    f32 n01 = hash2d_f(ix, iy + 1, seed);
    f32 n11 = hash2d_f(ix + 1, iy + 1, seed);

    return lerp(lerp(n00, n10, u), lerp(n01, n11, u), v);
}

static f32 noise_perlin(f32 x, f32 y, f32 scale, u32 seed) {
    f32 sx = x * scale;
    f32 sy = y * scale;
    i32 ix = (i32)floorf(sx);
    i32 iy = (i32)floorf(sy);
    f32 fx = sx - (f32)ix;
    f32 fy = sy - (f32)iy;
    f32 u = smoothstep(fx);
    f32 v = smoothstep(fy);

    f32 n00 = gradient2d(ix, iy, fx, fy, seed);
    f32 n10 = gradient2d(ix + 1, iy, fx - 1.0f, fy, seed);
    f32 n01 = gradient2d(ix, iy + 1, fx, fy - 1.0f, seed);
    f32 n11 = gradient2d(ix + 1, iy + 1, fx - 1.0f, fy - 1.0f, seed);

    f32 result = lerp(lerp(n00, n10, u), lerp(n01, n11, u), v);
    return result * 0.5f + 0.5f;
}

static f32 noise_fbm(f32 x, f32 y, i32 octaves, f32 scale, f32 persistence, u32 seed) {
    f32 value = 0.0f;
    f32 amplitude = 1.0f;
    f32 frequency = scale;
    f32 max_value = 0.0f;

    for (i32 i = 0; i < octaves; i++) {
        value += amplitude * noise_perlin(x, y, frequency, seed + (u32)i * 1337);
        max_value += amplitude;
        amplitude *= persistence;
        frequency *= 2.0f;
    }

    return value / max_value;
}

static f32 noise_ridged(f32 x, f32 y, i32 octaves, f32 scale, f32 persistence, u32 seed) {
    f32 value = 0.0f;
    f32 amplitude = 1.0f;
    f32 frequency = scale;
    f32 max_value = 0.0f;

    for (i32 i = 0; i < octaves; i++) {
        f32 n = noise_perlin(x, y, frequency, seed + (u32)i * 1337);
        n = 1.0f - fabsf(n * 2.0f - 1.0f);
        value += amplitude * n;
        max_value += amplitude;
        amplitude *= persistence;
        frequency *= 2.0f;
    }

    return value / max_value;
}

static f32 noise_turbulence(f32 x, f32 y, i32 octaves, f32 scale, f32 persistence, u32 seed) {
    f32 value = 0.0f;
    f32 amplitude = 1.0f;
    f32 frequency = scale;
    f32 max_value = 0.0f;

    for (i32 i = 0; i < octaves; i++) {
        f32 n = noise_perlin(x, y, frequency, seed + (u32)i * 1337);
        value += amplitude * fabsf(n * 2.0f - 1.0f);
        max_value += amplitude;
        amplitude *= persistence;
        frequency *= 2.0f;
    }

    return value / max_value;
}

static void voronoi(f32 x, f32 y, f32 scale, u32 seed, f32* cell_dist, f32* edge_dist, i32* cell_id) {
    f32 sx = x * scale;
    f32 sy = y * scale;
    i32 cx = (i32)floorf(sx);
    i32 cy = (i32)floorf(sy);
    f32 fx = sx - (f32)cx;
    f32 fy = sy - (f32)cy;

    f32 min_dist = 1e30f;
    f32 second_dist = 1e30f;
    i32 closest_id = 0;

    for (i32 dy = -1; dy <= 1; dy++) {
        for (i32 dx = -1; dx <= 1; dx++) {
            i32 nx = cx + dx;
            i32 ny = cy + dy;
            u32 h = hash2d(nx, ny, seed);
            f32 px = (f32)dx + (f32)(h & 0xFF) / 255.0f - 0.5f - fx;
            f32 py = (f32)dy + (f32)((h >> 8) & 0xFF) / 255.0f - 0.5f - fy;
            f32 dist = px * px + py * py;

            if (dist < min_dist) {
                second_dist = min_dist;
                min_dist = dist;
                closest_id = (i32)h;
            } else if (dist < second_dist) {
                second_dist = dist;
            }
        }
    }

    *cell_dist = sqrtf(min_dist);
    *edge_dist = sqrtf(second_dist) - sqrtf(min_dist);
    *cell_id = closest_id;
}

static f32 gradient_linear(f32 x, f32 y, f32 angle) {
    f32 dx = cosf(angle);
    f32 dy = sinf(angle);
    f32 result = x * dx + y * dy;
    return fmaxf(0.0f, fminf(1.0f, result));
}

static f32 gradient_radial(f32 x, f32 y, f32 cx, f32 cy) {
    f32 dx = x - cx;
    f32 dy = y - cy;
    return fmaxf(0.0f, fminf(1.0f, sqrtf(dx * dx + dy * dy)));
}

pxl8_graph* pxl8_graph_create(u32 capacity) {
    pxl8_graph* graph = calloc(1, sizeof(pxl8_graph));
    if (!graph) return NULL;

    graph->nodes = calloc(capacity, sizeof(pxl8_node));
    if (!graph->nodes) {
        free(graph);
        return NULL;
    }

    graph->capacity = capacity;
    graph->count = 0;
    graph->seed = 0;
    graph->output_reg = 0;
    return graph;
}

void pxl8_graph_destroy(pxl8_graph* graph) {
    if (!graph) return;
    free(graph->nodes);
    free(graph);
}

void pxl8_graph_clear(pxl8_graph* graph) {
    if (!graph) return;
    graph->count = 0;
    graph->output_reg = 0;
}

u8 pxl8_graph_add_node(pxl8_graph* graph, pxl8_graph_op op, u8 in0, u8 in1, u8 in2, u8 in3, f32 param) {
    if (!graph || graph->count >= graph->capacity) return 0;

    u8 out = (u8)(graph->count + 4);
    pxl8_node* node = &graph->nodes[graph->count++];
    node->op = (u8)op;
    node->out = out;
    node->in[0] = in0;
    node->in[1] = in1;
    node->in[2] = in2;
    node->in[3] = in3;
    node->param = param;

    return out;
}

void pxl8_graph_set_output(pxl8_graph* graph, u8 reg) {
    if (graph) graph->output_reg = reg;
}

void pxl8_graph_set_seed(pxl8_graph* graph, u32 seed) {
    if (graph) graph->seed = seed;
}

f32 pxl8_graph_eval(const pxl8_graph* graph, pxl8_graph_context* ctx) {
    if (!graph || !ctx) return 0.0f;

    for (u32 i = 0; i < graph->count; i++) {
        const pxl8_node* n = &graph->nodes[i];
        f32 a = ctx->regs[n->in[0]];
        f32 b = ctx->regs[n->in[1]];
        f32 c = ctx->regs[n->in[2]];
        f32 d = ctx->regs[n->in[3]];
        f32 result = 0.0f;

        switch (n->op) {
            case PXL8_OP_CONST:           result = n->param; break;
            case PXL8_OP_INPUT_AGE:       result = ctx->regs[3]; break;
            case PXL8_OP_INPUT_SEED:      result = (f32)ctx->seed; break;
            case PXL8_OP_INPUT_TIME:      result = ctx->regs[2]; break;
            case PXL8_OP_INPUT_X:         result = ctx->regs[0]; break;
            case PXL8_OP_INPUT_Y:         result = ctx->regs[1]; break;

            case PXL8_OP_ABS:             result = fabsf(a); break;
            case PXL8_OP_CEIL:            result = ceilf(a); break;
            case PXL8_OP_COS:             result = cosf(a); break;
            case PXL8_OP_FLOOR:           result = floorf(a); break;
            case PXL8_OP_FRACT:           result = a - floorf(a); break;
            case PXL8_OP_NEGATE:          result = -a; break;
            case PXL8_OP_SIN:             result = sinf(a); break;
            case PXL8_OP_SQRT:            result = sqrtf(a); break;

            case PXL8_OP_ADD:             result = a + b; break;
            case PXL8_OP_DIV:             result = b != 0.0f ? a / b : 0.0f; break;
            case PXL8_OP_MAX:             result = fmaxf(a, b); break;
            case PXL8_OP_MIN:             result = fminf(a, b); break;
            case PXL8_OP_MOD:             result = b != 0.0f ? fmodf(a, b) : 0.0f; break;
            case PXL8_OP_MUL:             result = a * b; break;
            case PXL8_OP_POW:             result = powf(a, b); break;
            case PXL8_OP_SUB:             result = a - b; break;

            case PXL8_OP_CLAMP:           result = fmaxf(b, fminf(c, a)); break;
            case PXL8_OP_LERP:            result = a + c * (b - a); break;
            case PXL8_OP_SELECT:          result = c > 0.0f ? a : b; break;
            case PXL8_OP_SMOOTHSTEP: {
                f32 t = fmaxf(0.0f, fminf(1.0f, (c - a) / (b - a)));
                result = t * t * (3.0f - 2.0f * t);
                break;
            }

            case PXL8_OP_NOISE_FBM:       result = noise_fbm(a, b, (i32)n->param, c, d, ctx->seed); break;
            case PXL8_OP_NOISE_PERLIN:    result = noise_perlin(a, b, c, ctx->seed); break;
            case PXL8_OP_NOISE_RIDGED:    result = noise_ridged(a, b, (i32)n->param, c, d, ctx->seed); break;
            case PXL8_OP_NOISE_TURBULENCE:result = noise_turbulence(a, b, (i32)n->param, c, d, ctx->seed); break;
            case PXL8_OP_NOISE_VALUE:     result = noise_value(a, b, c, ctx->seed); break;

            case PXL8_OP_VORONOI_CELL: {
                f32 cell, edge; i32 id;
                voronoi(a, b, c, ctx->seed, &cell, &edge, &id);
                result = cell;
                break;
            }
            case PXL8_OP_VORONOI_EDGE: {
                f32 cell, edge; i32 id;
                voronoi(a, b, c, ctx->seed, &cell, &edge, &id);
                result = edge;
                break;
            }
            case PXL8_OP_VORONOI_ID: {
                f32 cell, edge; i32 id;
                voronoi(a, b, c, ctx->seed, &cell, &edge, &id);
                result = (f32)(id & 0xFF) / 255.0f;
                break;
            }

            case PXL8_OP_GRADIENT_LINEAR: result = gradient_linear(a, b, c); break;
            case PXL8_OP_GRADIENT_RADIAL: result = gradient_radial(a, b, c, d); break;

            case PXL8_OP_QUANTIZE: {
                u8 base = (u8)n->param;
                f32 range = b;
                f32 clamped = fmaxf(0.0f, fminf(1.0f, a));
                result = (f32)(base + (u8)fminf(clamped * range, range - 1.0f));
                break;
            }

            default: break;
        }

        ctx->regs[n->out] = result;
    }

    return ctx->regs[graph->output_reg];
}

void pxl8_graph_eval_texture(const pxl8_graph* graph, u8* buffer, i32 width, i32 height) {
    if (!graph || !buffer) return;

    pxl8_graph_context ctx = {0};
    ctx.seed = graph->seed;

    for (i32 y = 0; y < height; y++) {
        for (i32 x = 0; x < width; x++) {
            ctx.regs[0] = (f32)x / (f32)width;
            ctx.regs[1] = (f32)y / (f32)height;
            ctx.regs[2] = 0.0f;
            ctx.regs[3] = 0.0f;

            f32 result = pxl8_graph_eval(graph, &ctx);
            buffer[y * width + x] = (u8)fmaxf(0.0f, fminf(255.0f, result));
        }
    }
}