#include "pxl8_particles.h"

#include <stdlib.h>

#include "pxl8_gfx.h"
#include "pxl8_gfx2d.h"
#include "pxl8_palette.h"
#include "pxl8_rng.h"

struct pxl8_particles {
    pxl8_particle* particles;
    pxl8_palette* palette;
    pxl8_rng* rng;
    u32 alive_count;
    u32 count;
    u32 max_count;

    f32 x, y;
    f32 spread_x, spread_y;

    f32 drag;
    f32 gravity_x, gravity_y;
    f32 turbulence;

    f32 spawn_rate;
    f32 spawn_timer;

    u8 color_min, color_max;
    f32 life_min, life_max;
    f32 size_min, size_max;
    f32 vx_min, vx_max, vy_min, vy_max;

    pxl8_particle_render_fn render_fn;
    pxl8_particle_spawn_fn spawn_fn;
    pxl8_particle_update_fn update_fn;
    void* userdata;
};

pxl8_particles* pxl8_particles_create(u32 max_count, pxl8_rng* rng) {
    pxl8_particles* ps = calloc(1, sizeof(pxl8_particles));
    if (!ps) return NULL;

    ps->particles = calloc(max_count, sizeof(pxl8_particle));
    if (!ps->particles) {
        free(ps);
        return NULL;
    }

    ps->rng = rng;
    ps->max_count = max_count;
    ps->drag = 0.98f;
    ps->gravity_y = 100.0f;
    ps->spawn_rate = 10.0f;

    ps->color_min = ps->color_max = 15;
    ps->life_min = ps->life_max = 1.0f;
    ps->size_min = ps->size_max = 1.0f;

    return ps;
}

void pxl8_particles_destroy(pxl8_particles* ps) {
    if (!ps) return;
    free(ps->particles);
    free(ps);
}

void pxl8_particles_clear(pxl8_particles* ps) {
    if (!ps || !ps->particles) return;

    for (u32 i = 0; i < ps->max_count; i++) {
        ps->particles[i].life = 0;
        ps->particles[i].flags = 0;
    }
    ps->alive_count = 0;
    ps->spawn_timer = 0;
}

void pxl8_particles_emit(pxl8_particles* ps, u32 count) {
    if (!ps || !ps->particles) return;

    for (u32 i = 0; i < count && ps->alive_count < ps->max_count; i++) {
        for (u32 j = 0; j < ps->max_count; j++) {
            if (ps->particles[j].life <= 0) {
                pxl8_particle* p = &ps->particles[j];

                f32 life = ps->life_min + pxl8_rng_f32(ps->rng) * (ps->life_max - ps->life_min);
                p->life = life;
                p->max_life = life;

                p->x = ps->x + (pxl8_rng_f32(ps->rng) - 0.5f) * ps->spread_x;
                p->y = ps->y + (pxl8_rng_f32(ps->rng) - 0.5f) * ps->spread_y;
                p->z = 0;

                p->vx = ps->vx_min + pxl8_rng_f32(ps->rng) * (ps->vx_max - ps->vx_min);
                p->vy = ps->vy_min + pxl8_rng_f32(ps->rng) * (ps->vy_max - ps->vy_min);
                p->vz = 0;

                p->ax = ps->gravity_x;
                p->ay = ps->gravity_y;
                p->az = 0;

                u8 ramp_range = ps->color_max - ps->color_min + 1;
                u8 ramp_pos = ps->color_min + (pxl8_rng_next(ps->rng) % ramp_range);
                u8 color = ps->palette
                    ? pxl8_palette_ramp_index(ps->palette, ramp_pos)
                    : ramp_pos;
                p->color = p->start_color = p->end_color = color;

                p->size = ps->size_min + pxl8_rng_f32(ps->rng) * (ps->size_max - ps->size_min);
                p->angle = 0;
                p->spin = 0;
                p->flags = 1;

                if (ps->spawn_fn) {
                    ps->spawn_fn(ps, p);
                }

                ps->alive_count++;
                break;
            }
        }
    }
}

void pxl8_particles_render(pxl8_particles* ps, pxl8_gfx* gfx) {
    if (!ps || !ps->particles || !gfx) return;

    for (u32 i = 0; i < ps->max_count; i++) {
        pxl8_particle* p = &ps->particles[i];
        if (p->life > 0 && p->flags) {
            if (ps->render_fn) {
                ps->render_fn(gfx, p, ps->userdata);
            } else {
                i32 x = (i32)p->x;
                i32 y = (i32)p->y;
                if (x >= 0 && x < pxl8_gfx_get_width(gfx) && y >= 0 && y < pxl8_gfx_get_height(gfx)) {
                    pxl8_2d_pixel(gfx, x, y, p->color);
                }
            }
        }
    }
}

void pxl8_particles_update(pxl8_particles* ps, f32 dt) {
    if (!ps || !ps->particles) return;

    if (ps->spawn_rate > 0.0f) {
        ps->spawn_timer += dt;
        f32 spawn_interval = 1.0f / ps->spawn_rate;
        u32 max_spawns_per_frame = ps->max_count / 10;
        if (max_spawns_per_frame < 1) max_spawns_per_frame = 1;
        u32 spawn_count = 0;
        while (ps->spawn_timer >= spawn_interval && spawn_count < max_spawns_per_frame) {
            pxl8_particles_emit(ps, 1);
            ps->spawn_timer -= spawn_interval;
            spawn_count++;
        }
    }

    for (u32 i = 0; i < ps->max_count; i++) {
        pxl8_particle* p = &ps->particles[i];
        if (p->life > 0) {
            if (ps->update_fn) {
                ps->update_fn(p, dt, ps->userdata);
            } else {
                p->vx += p->ax * dt;
                p->vy += p->ay * dt;
                p->vz += p->az * dt;

                p->vx *= ps->drag;
                p->vy *= ps->drag;
                p->vz *= ps->drag;

                p->x += p->vx * dt;
                p->y += p->vy * dt;
                p->z += p->vz * dt;

                p->angle += p->spin * dt;
            }

            p->life -= dt / p->max_life;
            if (p->life <= 0) {
                p->flags = 0;
                ps->alive_count--;
            }
        }
    }
}

u32 pxl8_particles_count(const pxl8_particles* ps) {
    return ps ? ps->alive_count : 0;
}

u32 pxl8_particles_max_count(const pxl8_particles* ps) {
    return ps ? ps->max_count : 0;
}

pxl8_particle* pxl8_particles_get(pxl8_particles* ps, u32 index) {
    if (!ps || index >= ps->max_count) return NULL;
    return &ps->particles[index];
}

pxl8_rng* pxl8_particles_rng(pxl8_particles* ps) {
    return ps ? ps->rng : NULL;
}

f32 pxl8_particles_get_drag(const pxl8_particles* ps) {
    return ps ? ps->drag : 0.0f;
}

f32 pxl8_particles_get_gravity_x(const pxl8_particles* ps) {
    return ps ? ps->gravity_x : 0.0f;
}

f32 pxl8_particles_get_gravity_y(const pxl8_particles* ps) {
    return ps ? ps->gravity_y : 0.0f;
}

f32 pxl8_particles_get_spawn_rate(const pxl8_particles* ps) {
    return ps ? ps->spawn_rate : 0.0f;
}

f32 pxl8_particles_get_spread_x(const pxl8_particles* ps) {
    return ps ? ps->spread_x : 0.0f;
}

f32 pxl8_particles_get_spread_y(const pxl8_particles* ps) {
    return ps ? ps->spread_y : 0.0f;
}

f32 pxl8_particles_get_turbulence(const pxl8_particles* ps) {
    return ps ? ps->turbulence : 0.0f;
}

void* pxl8_particles_get_userdata(const pxl8_particles* ps) {
    return ps ? ps->userdata : NULL;
}

f32 pxl8_particles_get_x(const pxl8_particles* ps) {
    return ps ? ps->x : 0.0f;
}

f32 pxl8_particles_get_y(const pxl8_particles* ps) {
    return ps ? ps->y : 0.0f;
}

void pxl8_particles_set_colors(pxl8_particles* ps, u8 color_min, u8 color_max) {
    if (!ps) return;
    ps->color_min = color_min;
    ps->color_max = color_max;
}

void pxl8_particles_set_drag(pxl8_particles* ps, f32 drag) {
    if (ps) ps->drag = drag;
}

void pxl8_particles_set_gravity(pxl8_particles* ps, f32 gx, f32 gy) {
    if (!ps) return;
    ps->gravity_x = gx;
    ps->gravity_y = gy;
}

void pxl8_particles_set_life(pxl8_particles* ps, f32 life_min, f32 life_max) {
    if (!ps) return;
    ps->life_min = life_min;
    ps->life_max = life_max;
}

void pxl8_particles_set_palette(pxl8_particles* ps, pxl8_palette* palette) {
    if (ps) ps->palette = palette;
}

void pxl8_particles_set_position(pxl8_particles* ps, f32 x, f32 y) {
    if (!ps) return;
    ps->x = x;
    ps->y = y;
}

void pxl8_particles_set_render_fn(pxl8_particles* ps, pxl8_particle_render_fn fn) {
    if (ps) ps->render_fn = fn;
}

void pxl8_particles_set_size(pxl8_particles* ps, f32 size_min, f32 size_max) {
    if (!ps) return;
    ps->size_min = size_min;
    ps->size_max = size_max;
}

void pxl8_particles_set_spawn_fn(pxl8_particles* ps, pxl8_particle_spawn_fn fn) {
    if (ps) ps->spawn_fn = fn;
}

void pxl8_particles_set_spawn_rate(pxl8_particles* ps, f32 rate) {
    if (ps) ps->spawn_rate = rate;
}

void pxl8_particles_set_spread(pxl8_particles* ps, f32 spread_x, f32 spread_y) {
    if (!ps) return;
    ps->spread_x = spread_x;
    ps->spread_y = spread_y;
}

void pxl8_particles_set_turbulence(pxl8_particles* ps, f32 turbulence) {
    if (ps) ps->turbulence = turbulence;
}

void pxl8_particles_set_update_fn(pxl8_particles* ps, pxl8_particle_update_fn fn) {
    if (ps) ps->update_fn = fn;
}

void pxl8_particles_set_userdata(pxl8_particles* ps, void* userdata) {
    if (ps) ps->userdata = userdata;
}

void pxl8_particles_set_velocity(pxl8_particles* ps, f32 vx_min, f32 vx_max, f32 vy_min, f32 vy_max) {
    if (!ps) return;
    ps->vx_min = vx_min;
    ps->vx_max = vx_max;
    ps->vy_min = vy_min;
    ps->vy_max = vy_max;
}