// sand_sim.c
#include "sand_sim.h"
#include "gl_utils.h"   // for create_compute_program_from_file

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

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

    // 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->prog_sim = 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 program\n");
        return false;
    }

    return true;
}

void sand_step_gpu(SandSim* sim) {
    if (!sim || !sim->prog_sim) return;

    glUseProgram(sim->prog_sim);

    // Set grid size uniform
    GLint uGridLoc = glGetUniformLocation(sim->prog_sim, "u_gridSize");
    if (uGridLoc >= 0) {
        glUniform2i(uGridLoc, sim->gridW, sim->gridH);
    }

    // Bind images (must match bindings in sand_step.comp)
    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);

    GLuint groupsX = (sim->gridW + 15) / 16;
    GLuint groupsY = (sim->gridH + 15) / 16;

    glDispatchCompute(groupsX, groupsY, 1);
    glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

    // Ping-pong
    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);
        sim->tex_curr = 0;
    }
    if (sim->tex_next) {
        glDeleteTextures(1, &sim->tex_next);
        sim->tex_next = 0;
    }
    if (sim->prog_sim) {
        glDeleteProgram(sim->prog_sim);
        sim->prog_sim = 0;
    }
}