#include <js.h>
#include <math.h>
#include <memory.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <vec2.h>

typedef enum : uint8_t {
    TYPE_EMPTY = 0,
    TYPE_CIRCLE = 1,
    TYPE_PLANE = 2,
    TYPE_COUNT,
} body_type;

////////// Types

typedef struct {
    // Public
    body_type type;
    uint8_t reserved[3];
    vec2 pos;
    vec2 vel;
    float mass;
    union {
        struct {
            float radius;
        } circle;
        struct {
            vec2 normal;
        } plane;
    };
    // Private
    vec2 force;
} rigid_body;

typedef size_t rigid_body_index;
typedef void (*rigid_body_collision_callback_t)(rigid_body_index a, rigid_body_index b);

// `normal` points from body a toward body b; `overlap` is the penetration depth.
typedef struct {
    bool hit;
    vec2 normal;
    float overlap;
} contact;

typedef contact (*contact_fn)(const rigid_body* a, const rigid_body* b);

constexpr float COLLISION_EPSILON = 1e-6f;

////////// Prototypes

void rigid_body_resolve_collision(rigid_body_index rb);
void rigid_body_handle_collision(rigid_body_index a, rigid_body_index b);
float point_to_line_dist(vec2 point, vec2 line_point, vec2 line_norm);
rigid_body_index rigid_body_find_empty();

////////// Globals

static rigid_body* rigid_bodies = NULL;
static size_t rigid_bodies_cap = 0;
static rigid_body_collision_callback_t rigid_body_collision_callback = NULL;

static vec2 gravity = {0, 0};

// Queued during the solve and dispatched only after it finishes, so callbacks may
// free or create bodies (reallocating `rigid_bodies`) without invalidating the body
// pointers held mid-solve.
typedef struct {
    rigid_body_index a;
    rigid_body_index b;
} collision_event;

static collision_event* collision_events = NULL;
static size_t collision_events_count = 0;
static size_t collision_events_cap = 0;

////////// Functions

inline static rigid_body* rb_get(rigid_body_index idx) {
    return (rigid_bodies + (idx));
}

static void collision_event_push(rigid_body_index a, rigid_body_index b) {
    if (collision_events_count == collision_events_cap) {
        size_t new_cap = collision_events_cap ? collision_events_cap * 2 : 16;
        collision_event* grown = realloc(collision_events, new_cap * sizeof(collision_event));
        if (!grown) {
            return;
        }
        collision_events = grown;
        collision_events_cap = new_cap;
    }
    collision_events[collision_events_count].a = a;
    collision_events[collision_events_count].b = b;
    collision_events_count++;
}

static void rigid_body_dispatch_collisions() {
    if (rigid_body_collision_callback) {
        for (size_t i = 0; i < collision_events_count; i++) {
            rigid_body_collision_callback(collision_events[i].a, collision_events[i].b);
        }
    }
    collision_events_count = 0;
}

static void rigid_body_integrate(rigid_body* rb, float dt) {
    if (!isinf(rb->mass)) {
        rb->vel.x += (rb->force.x / rb->mass + gravity.x) * dt;
        rb->vel.y += (rb->force.y / rb->mass + gravity.y) * dt;

        rb->pos.x += rb->vel.x * dt;
        rb->pos.y += rb->vel.y * dt;
    }

    rb->force.x = 0;
    rb->force.y = 0;
}

JS_EXPORT rigid_body* rigid_body_get(rigid_body_index idx) {
    return rb_get(idx);
}

size_t rigid_body_new(float x, float y, float vx, float vy, float mass) {
    rigid_body_index idx = rigid_body_find_empty();
    if (idx == SIZE_MAX) {
        return SIZE_MAX;  // allocation failed
    }
    rigid_body* rb = rb_get(idx);
    rb->type = TYPE_EMPTY;

    rb->pos.x = x;
    rb->pos.y = y;

    rb->vel.x = vx;
    rb->vel.y = vy;

    rb->force.x = 0;
    rb->force.y = 0;

    rb->mass = mass;

    return idx;
}

JS_EXPORT rigid_body_index rigid_body_new_circle(float x, float y, float vx, float vy, float mass, float radius) {
    rigid_body_index idx = rigid_body_new(x, y, vx, vy, mass);
    if (idx == SIZE_MAX) {
        return idx;
    }
    rigid_body* rb = rb_get(idx);

    rb->type = TYPE_CIRCLE;
    rb->circle.radius = radius;

    return idx;
}

JS_EXPORT rigid_body_index rigid_body_new_plane(float x, float y, float nx, float ny) {
    rigid_body_index idx = rigid_body_new(x, y, 0, 0, INFINITY);
    if (idx == SIZE_MAX) {
        return idx;
    }
    rigid_body* rb = rb_get(idx);

    rb->type = TYPE_PLANE;
    rb->plane.normal = vec2_normalize((vec2){nx, ny});

    return idx;
}

JS_EXPORT void rigid_body_free(rigid_body_index idx) {
    memset(rb_get(idx), 0, sizeof(rigid_body));
}

JS_EXPORT void rigid_body_update(rigid_body_index idx, float dt) {
    rigid_body_integrate(rb_get(idx), dt);
    rigid_body_resolve_collision(idx);
    rigid_body_dispatch_collisions();
}

JS_EXPORT void rigid_body_update_all(float dt) {
    for (rigid_body_index idx = 0; idx < rigid_bodies_cap; idx++) {
        rigid_body* current = rb_get(idx);
        if (current->type == TYPE_EMPTY) {
            continue;
        }
        rigid_body_integrate(current, dt);
    }

    for (rigid_body_index idx = 0; idx < rigid_bodies_cap; idx++) {
        rigid_body* current = rb_get(idx);
        if (current->type == TYPE_EMPTY) {
            continue;
        }
        rigid_body_resolve_collision(idx);
    }

    rigid_body_dispatch_collisions();
}

JS_EXPORT void rigid_body_add_force(rigid_body_index idx, float fx, float fy) {
    rigid_body* rb = rb_get(idx);
    rb->force.x += fx;
    rb->force.y += fy;
}

JS_EXPORT void rigid_body_add_global_force(float fx, float fy) {
    for (rigid_body_index idx = 0; idx < rigid_bodies_cap; idx++) {
        rigid_body* current = rb_get(idx);
        if (current->type == TYPE_EMPTY) {
            continue;
        }
        rigid_body_add_force(idx, fx, fy);
    }
}

JS_EXPORT void rigid_body_set_collision_callback(rigid_body_collision_callback_t callback) {
    rigid_body_collision_callback = callback;
}

JS_EXPORT void rigid_body_set_gravity(float gx, float gy) {
    gravity.x = gx;
    gravity.y = gy;
}

void rigid_body_resolve_collision(rigid_body_index idx) {
    rigid_body* rb = rb_get(idx);
    int is_static = isinf(rb->mass);
    for (rigid_body_index i = idx + 1; i < rigid_bodies_cap; i++) {
        rigid_body* current = rb_get(i);
        if (!is_static || !isinf(current->mass)) {
            rigid_body_handle_collision(idx, i);
        }
    }
}

////////// Contact generators

static contact contact_circle_circle(const rigid_body* a, const rigid_body* b) {
    vec2 d = vec2_sub(b->pos, a->pos);
    float distance = vec2_mag(d);

    float overlap = a->circle.radius + b->circle.radius - distance;
    if (overlap < 0) {
        return (contact){.hit = false};
    }

    // Coincident centres have no separation axis; pick an arbitrary one.
    vec2 n = (distance > COLLISION_EPSILON) ? vec2_div(d, distance) : (vec2){1.0f, 0.0f};
    return (contact){.hit = true, .normal = n, .overlap = overlap};
}

static contact contact_plane_circle(const rigid_body* a, const rigid_body* b) {
    // a: plane (static), b: circle.
    float distance = point_to_line_dist(b->pos, a->pos, a->plane.normal) - b->circle.radius;

    float overlap = -distance;
    if (overlap < 0) {
        return (contact){.hit = false};
    }

    return (contact){.hit = true, .normal = a->plane.normal, .overlap = overlap};
}

// Indexed [a->type][b->type] in canonical order (a->type >= b->type); a NULL entry
// means the pair does not collide.
static contact_fn contact_table[TYPE_COUNT][TYPE_COUNT] = {
    [TYPE_CIRCLE][TYPE_CIRCLE] = contact_circle_circle,
    [TYPE_PLANE][TYPE_CIRCLE] = contact_plane_circle,
};

// `n` points from a toward b; infinite-mass bodies are treated as immovable.
static void resolve_contact(rigid_body* a, rigid_body* b, vec2 n, float overlap) {
    float factor = (isinf(a->mass) || isinf(b->mass)) ? 1 : 0.5;
    vec2 n_overlap = vec2_mul(n, overlap * factor);

    if (!isinf(a->mass)) {
        a->pos = vec2_sub(a->pos, n_overlap);
    }
    if (!isinf(b->mass)) {
        b->pos = vec2_add(b->pos, n_overlap);
    }

    float van = vec2_dot(a->vel, n);
    float vbn = vec2_dot(b->vel, n);

    // Skip the impulse if the bodies are already separating along the normal.
    if (van - vbn <= 0) {
        return;
    }

    float van_new = van;
    float vbn_new = vbn;

    if (!isinf(a->mass) && !isinf(b->mass)) {
        // Elastic collision: each body weighted by the *other* body's mass.
        van_new = van - 2.0f * (van - vbn) * b->mass / (a->mass + b->mass);
        vbn_new = vbn + 2.0f * (van - vbn) * a->mass / (a->mass + b->mass);
    } else if (isinf(a->mass)) {
        vbn_new = -vbn;
    } else if (isinf(b->mass)) {
        van_new = -van;
    } else {
        return;
    }

    vec2 vat = vec2_sub(a->vel, vec2_mul(n, van));
    vec2 vbt = vec2_sub(b->vel, vec2_mul(n, vbn));

    if (!isinf(a->mass)) {
        a->vel = vec2_add(vat, vec2_mul(n, van_new));
    }
    if (!isinf(b->mass)) {
        b->vel = vec2_add(vbt, vec2_mul(n, vbn_new));
    }
}

void rigid_body_handle_collision(rigid_body_index idx1, rigid_body_index idx2) {
    // The greater type enum becomes `a` so each pair needs only one table entry; ties
    // keep the caller's order, preserving the argument order seen by the callback.
    rigid_body_index ia = idx1;
    rigid_body_index ib = idx2;
    if (rb_get(idx2)->type > rb_get(idx1)->type) {
        ia = idx2;
        ib = idx1;
    }

    rigid_body* a = rb_get(ia);
    rigid_body* b = rb_get(ib);

    contact_fn generate = contact_table[a->type][b->type];
    if (!generate) {
        return;
    }

    contact c = generate(a, b);
    if (!c.hit) {
        return;
    }

    collision_event_push(ia, ib);
    resolve_contact(a, b, c.normal, c.overlap);
}

float point_to_line_dist(vec2 p, vec2 line_point, vec2 normal) {
    vec2 diff = vec2_sub(p, line_point);
    return vec2_dot(diff, normal);
}

rigid_body_index rigid_body_find_empty() {
    if (rigid_bodies) {
        for (rigid_body_index idx = 0; idx < rigid_bodies_cap; idx++) {
            rigid_body* current = rb_get(idx);
            if (current->type == TYPE_EMPTY) {
                memset(current, 0, sizeof(rigid_body));
                return idx;
            }
        }
        size_t old_cap = rigid_bodies_cap;
        size_t new_cap = old_cap * 2;
        rigid_body* grown = realloc(rigid_bodies, new_cap * sizeof(rigid_body));
        if (!grown) {
            return SIZE_MAX;  // keep the old buffer intact; caller handles failure
        }
        rigid_bodies = grown;
        memset(rigid_bodies + old_cap, 0, (new_cap - old_cap) * sizeof(rigid_body));
        rigid_bodies_cap = new_cap;
        return old_cap;  // first slot of the freshly grown region
    } else {
        rigid_body* allocated = malloc(2 * sizeof(rigid_body));
        if (!allocated) {
            return SIZE_MAX;
        }
        rigid_bodies = allocated;
        rigid_bodies_cap = 2;
        memset(rigid_bodies, 0, rigid_bodies_cap * sizeof(rigid_body));
        return 0;
    }
}