// sand_sim.c
#include "sand_sim.h"

#include <stdio.h>
#include <stdlib.h>
#include <time.h>

#include "gl_utils.h"  // for create_compute_program_from_file
void sand_paint_gpu(SandSim* sim,
                    float brushX,
                    float brushY,
                    float brushRadius,
                    int mode) {
    if (!sim || !sim->prog_paint) return;
    if (mode == 0) return;  // nothing to do

    glUseProgram(sim->prog_paint);

    GLint locGrid = glGetUniformLocation(sim->prog_paint, "u_gridSize");
    GLint locPos = glGetUniformLocation(sim->prog_paint, "u_brushPos");
    GLint locRad = glGetUniformLocation(sim->prog_paint, "u_brushRadius");
    GLint locMode = glGetUniformLocation(sim->prog_paint, "u_mode");

    if (locGrid >= 0) glUniform2i(locGrid, sim->gridW, sim->gridH);
    if (locPos >= 0) glUniform2f(locPos, brushX, brushY);
    if (locRad >= 0) glUniform1f(locRad, brushRadius);
    if (locMode >= 0) glUniform1i(locMode, mode);

    // Bind current state as read-write image
    glBindImageTexture(0, sim->tex_curr,
                       0, GL_FALSE, 0,
                       GL_READ_WRITE, GL_R8UI);

    // Dispatch over entire grid; 16x16 workgroup like your sim shader
    GLuint groupsX = (sim->gridW + 16 - 1) / 16;
    GLuint groupsY = (sim->gridH + 16 - 1) / 16;
    glDispatchCompute(groupsX, groupsY, 1);

    // Ensure writes visible to subsequent rendering/compute
    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);
}

static void init_textures(SandSim* sim) {
    // tex_curr
    glGenTextures(1, &sim->tex_curr);
    glBindTexture(GL_TEXTURE_2D, sim->tex_curr);
    glTexStorage2D(GL_TEXTURE_2D, 1, GL_R8UI, sim->gridW, sim->gridH);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);

    // tex_next
    glGenTextures(1, &sim->tex_next);
    glBindTexture(GL_TEXTURE_2D, sim->tex_next);
    glTexStorage2D(GL_TEXTURE_2D, 1, GL_R8UI, sim->gridW, sim->gridH);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);

    // NEW: tex_claim for the atomic “ownership” buffer
    glGenTextures(1, &sim->tex_claim);
    glBindTexture(GL_TEXTURE_2D, sim->tex_claim);
    glTexStorage2D(GL_TEXTURE_2D, 1, GL_R32UI, sim->gridW, sim->gridH);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);

    // Unbind
    glBindTexture(GL_TEXTURE_2D, 0);
}

static void upload_initial_state(SandSim* sim) {
    size_t count = (size_t)sim->gridW * (size_t)sim->gridH;
    uint8_t* data = calloc(count, 1);
    if (!data) {
        fprintf(stderr, "[sand_init] Out of memory for initial grid\n");
        exit(EXIT_FAILURE);
    }

    // Seed RNG once (safe enough for now)
    static bool seeded = false;
    if (!seeded) {
        seeded = true;
        srand((unsigned int)time(NULL));
    }

    // Fill top 1/3 of grid with random sand vs air
    for (int y = 0; y < sim->gridH / 3; ++y) {
        for (int x = 0; x < sim->gridW; ++x) {
            data[(size_t)y * sim->gridW + x] = (rand() & 1) ? 1 : 0;  // 0=air,1=sand
        }
    }

    // Upload into tex_curr and tex_next
    glBindTexture(GL_TEXTURE_2D, sim->tex_curr);
    glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, sim->gridW, sim->gridH,
                    GL_RED_INTEGER, GL_UNSIGNED_BYTE, data);

    glBindTexture(GL_TEXTURE_2D, sim->tex_next);
    glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, sim->gridW, sim->gridH,
                    GL_RED_INTEGER, GL_UNSIGNED_BYTE, data);

    glBindTexture(GL_TEXTURE_2D, 0);

    free(data);
}

bool sand_init(SandSim* sim, int gridW, int gridH, const char* computeShaderPath) {
    if (!sim || gridW <= 0 || gridH <= 0 || !computeShaderPath) {
        fprintf(stderr, "[sand_init] Invalid arguments\n");
        return false;
    }

    sim->gridW = gridW;
    sim->gridH = gridH;
    sim->tex_curr = 0;
    sim->tex_next = 0;
    sim->tex_claim = 0;
    sim->prog_sim = 0;
    sim->prog_paint = 0;

    init_textures(sim);
    upload_initial_state(sim);

    sim->prog_sim = create_compute_program_from_file(computeShaderPath);
    if (!sim->prog_sim) {
        fprintf(stderr, "[sand_init] Failed to create compute cell tick program\n");
        return false;
    }
    sim->prog_paint = create_compute_program_from_file("./shaders/sand_paint.comp");
    if (!sim->prog_paint) {
        fprintf(stderr, "[sand_init] Failed to create compute paint program\n");
        return false;
    }
    sim->prog_relax = create_compute_program_from_file("./shaders/sand_relax.comp");
    if (!sim->prog_relax) {
        fprintf(stderr, "[sand_init] Failed to create compute relax program\n");
        return false;
    }

    return true;
}

void sand_step_gpu(SandSim* sim) {
    glUseProgram(sim->prog_sim);
    static unsigned int frameCounter = 0;
    frameCounter++;

    GLint locFrame = glGetUniformLocation(sim->prog_sim, "u_frame");
    if (locFrame >= 0) {
        glUniform1ui(locFrame, frameCounter);
    }
    GLint locGrid = glGetUniformLocation(sim->prog_sim, "u_gridSize");
    if (locGrid >= 0) {
        glUniform2i(locGrid, sim->gridW, sim->gridH);
    }

    // Clear next state to AIR (0)

    const GLubyte zeroByte[1] = {0};
    glClearTexImage(sim->tex_next,
                    0,
                    GL_RED_INTEGER,
                    GL_UNSIGNED_BYTE,
                    zeroByte);

    // Clear claim buffer to 0
    const GLuint zero1[1] = {0};
    glClearTexImage(sim->tex_claim,
                    0,
                    GL_RED_INTEGER,  // format doesn't really matter; data is uint
                    GL_UNSIGNED_INT,
                    zero1);

    // Bind images
    glBindImageTexture(0, sim->tex_curr,
                       0, GL_FALSE, 0,
                       GL_READ_ONLY, GL_R8UI);
    glBindImageTexture(1, sim->tex_next,
                       0, GL_FALSE, 0,
                       GL_WRITE_ONLY, GL_R8UI);
    glBindImageTexture(2, sim->tex_claim,
                       0, GL_FALSE, 0,
                       GL_READ_WRITE, GL_R32UI);

    GLuint groupsX = (sim->gridW + 16 - 1) / 16;
    GLuint groupsY = (sim->gridH + 16 - 1) / 16;
    glDispatchCompute(groupsX, groupsY, 1);

    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

    // ping-pong: curr <-> next
    GLuint tmp = sim->tex_curr;
    sim->tex_curr = sim->tex_next;
    sim->tex_next = tmp;
}
void sand_relax_gpu(SandSim* sim) {
    if (!sim || !sim->prog_relax) return;

    glUseProgram(sim->prog_relax);

    static unsigned int relaxFrame = 0;
    relaxFrame++;

    GLint locFrame = glGetUniformLocation(sim->prog_relax, "u_frame");
    if (locFrame >= 0) {
        glUniform1ui(locFrame, relaxFrame);
    }

    GLint locGrid = glGetUniformLocation(sim->prog_relax, "u_gridSize");
    if (locGrid >= 0) {
        glUniform2i(locGrid, sim->gridW, sim->gridH);
    }

    // Clear next state to AIR (0) – tex_next is GL_R8UI
    const GLubyte zeroByte[1] = {0};
    glClearTexImage(sim->tex_next,
                    0,
                    GL_RED_INTEGER,
                    GL_UNSIGNED_BYTE,
                    zeroByte);

    // Clear claim buffer (GL_R32UI) to 0
    const GLuint zero1[1] = {0};
    glClearTexImage(sim->tex_claim,
                    0,
                    GL_RED_INTEGER,
                    GL_UNSIGNED_INT,
                    zero1);

    // Bind images: curr = read, next = write, claim = read/write
    glBindImageTexture(0, sim->tex_curr,
                       0, GL_FALSE, 0,
                       GL_READ_ONLY,  GL_R8UI);
    glBindImageTexture(1, sim->tex_next,
                       0, GL_FALSE, 0,
                       GL_WRITE_ONLY, GL_R8UI);
    glBindImageTexture(2, sim->tex_claim,
                       0, GL_FALSE, 0,
                       GL_READ_WRITE, GL_R32UI);

    GLuint groupsX = (sim->gridW + 16 - 1) / 16;
    GLuint groupsY = (sim->gridH + 16 - 1) / 16;
    glDispatchCompute(groupsX, groupsY, 1);

    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

    // ping-pong: curr <-> next
    GLuint tmp = sim->tex_curr;
    sim->tex_curr = sim->tex_next;
    sim->tex_next = tmp;
}

void sand_destroy(SandSim* sim) {
    if (!sim) return;

    if (sim->tex_curr)  glDeleteTextures(1, &sim->tex_curr);
    if (sim->tex_next)  glDeleteTextures(1, &sim->tex_next);
    if (sim->tex_claim) glDeleteTextures(1, &sim->tex_claim);
    sim->tex_curr = sim->tex_next = sim->tex_claim = 0;

    if (sim->prog_sim)   glDeleteProgram(sim->prog_sim);
    if (sim->prog_paint) glDeleteProgram(sim->prog_paint);
    if (sim->prog_relax) glDeleteProgram(sim->prog_relax);
    sim->prog_sim = sim->prog_paint = sim->prog_relax = 0;
}