#include "wapp.h"
#include "common.h"
#include "raylib.h"
#include <math.h>
#include <stdint.h>
#include <stdbool.h>
#include <assert.h>
#include <string.h>
#include <pthread.h>

#define THREAD_COUNT 4

#define MOVABLE_TAG_SHIFT     0
#define RENDERABLE_TAG_SHIFT  1
#define COLLIDER_TAG_SHIFT    2
#define INSIDE_ZONE_TAG_SHIFT 3

enum EntityTag {
  ENTITY_TAG_MOVABLE     = 1 << MOVABLE_TAG_SHIFT,
  ENTITY_TAG_RENDERABLE  = 1 << RENDERABLE_TAG_SHIFT,
  ENTITY_TAG_COLLIDER    = 1 << COLLIDER_TAG_SHIFT,
  ENTITY_TAG_INSIDE_ZONE = 1 << INSIDE_ZONE_TAG_SHIFT,
};

typedef struct Entity Entity;
struct Entity {
  u32 id;
};

typedef struct Position Position;
struct Position {
  i16 x;
  i16 y;
};

typedef struct Scale Scale;
struct Scale {
  i16 width;
  i16 height;
};

typedef struct Rect Rect;
struct Rect {
  Position position;
  Scale scale;
};

typedef struct Velocity Velocity;
struct Velocity {
  i16 x;
  i16 y;
};

typedef struct Manager Manager;
struct Manager {
  Entity   *entities;
  Rect     *rects;
  Velocity *velocities;
  u8       *tags;
  u64      count;
};

typedef struct PositionThreadArgs PositionThreadArgs;
struct PositionThreadArgs {
  const Rect *zones;
  Rect       *rects;
  Velocity   *velocities;
  u8         *tags;
  u64        count;
  f32        delta;
};

typedef void *(*PThreadRoutine)(void *);
typedef void (*ScaleInitialiser)(Scale *scale, XOR256State *state);
typedef void (*VelocityInitialiser)(Velocity *velocity, XOR256State *state);
typedef void (*RaylibDrawRectFunc)(int posX, int posY, int width, int height, Color color);

void init_manager(const Allocator *allocator, Manager *manager);
void init_position(Position *position, XOR256State *state);
void init_scale_wanderer(Scale *scale, XOR256State *state);
void init_scale_zone(Scale *scale, XOR256State *state);
void init_velocity(Velocity *velocity, XOR256State *state);
void zero_velocity(Velocity *velocity, XOR256State *state);
void update_positions(u8 *tags, Rect *rects, Velocity *velocities, u64 count, f32 delta);
void *update_position_thread(PositionThreadArgs *args);
void render_entities(const u8 *tags, const Rect *rects, u64 count);
u8 collides(const Rect *rect, const Rect *collider);
i16 get_random_velocity(XOR256State *state);

int main(void) {
  i32 target_fps = 120;

  SetTraceLogLevel(LOG_NONE);
  InitWindow(WIDTH, HEIGHT, "DOD test");
  SetTargetFPS(target_fps);

  Allocator arena = wapp_mem_arena_allocator_init(MB(20));
  assert(!wapp_mem_allocator_invalid(&arena));

  XOR256State state = wapp_prng_xorshift_init_state();
  Manager manager   = {0};

  init_manager(&arena, &manager);

  ScaleInitialiser scale_initialisers[2]       = {init_scale_wanderer, init_scale_zone};
  VelocityInitialiser velocity_initialisers[2] = {init_velocity, zero_velocity};

  for (u64 i = 0; i < manager.count; ++i) {
    u8 is_zone = (i - ZONE_COUNT) >> 63;

    manager.entities[i].id  = i;
    manager.tags[i]         = (ENTITY_TAG_COLLIDER & (is_zone << COLLIDER_TAG_SHIFT)) | ENTITY_TAG_MOVABLE | ENTITY_TAG_RENDERABLE;

    init_position(&(manager.rects[i].position), &state);
    scale_initialisers[is_zone](&(manager.rects[i].scale), &state);
    velocity_initialisers[is_zone](&(manager.velocities[i]), &state);
  }

  f32 last_time = GetFrameTime();
  f32 delta, cur_time;

  while (!WindowShouldClose()) {
    f64 time = GetTime();
    if (time >= 20.0) {
      break;
    }

    cur_time = GetFrameTime();
    delta    = cur_time - last_time;

    update_positions(manager.tags, manager.rects, manager.velocities, manager.count, delta);

    BeginDrawing();

    ClearBackground(BG_COLOR);

    render_entities(manager.tags, manager.rects, manager.count);

    DrawFPS(10, 10);

    EndDrawing();
  }

  wapp_mem_arena_allocator_destroy(&arena);

  CloseWindow();

  return 0;
}

void init_manager(const Allocator *allocator, Manager *manager) {
  assert(allocator != NULL && manager != NULL);

  u64 total_count = (u64)WANDERER_COUNT + (u64)ZONE_COUNT;
  assert(total_count < (1lu << 32)); // Ensure we're not exceeding the maximum limit of entities

  u64 entities_size   = sizeof(Entity)   * total_count;
  u64 rects_size      = sizeof(Rect)     * total_count;
  u64 velocities_size = sizeof(Velocity) * total_count;
  u64 tags_size       = sizeof(u8)       * total_count;
  u64 allocation_size = entities_size + rects_size + velocities_size + tags_size;
  u8 *buffer          = wapp_mem_allocator_alloc(allocator, allocation_size);

  assert(buffer != NULL);

  memset(buffer, 0, allocation_size);

  manager->entities   = (Entity *)buffer;
  manager->rects      = (Rect *)(buffer + entities_size);
  manager->velocities = (Velocity *)(buffer + entities_size + rects_size);
  manager->tags       = (u8 *)(buffer + entities_size + rects_size + velocities_size);
  manager->count      = total_count;
}

void init_position(Position *position, XOR256State *state) {
  position->x = wapp_prng_xorshift_256ss(state) % WIDTH;
  position->y = wapp_prng_xorshift_256ss(state) % HEIGHT;
}

void init_scale_wanderer(Scale *scale, XOR256State *state) {
  i16 value = (i16)((wapp_prng_xorshift_256p(state) % (MAX_WANDERER_DIM + 1 - MIN_WANDERER_DIM)) + MIN_WANDERER_DIM);
  scale->width  = value;
  scale->height = value;
}

void init_scale_zone(Scale *scale, XOR256State *state) {
  scale->width  = wapp_prng_xorshift_256p(state) % ((u64)HALF_WIDTH  - MIN_ZONE_DIM) + MIN_ZONE_DIM;
  scale->height = wapp_prng_xorshift_256p(state) % ((u64)HALF_HEIGHT - MIN_ZONE_DIM) + MIN_ZONE_DIM;
}

void init_velocity(Velocity *velocity, XOR256State *state) {
  velocity->x = get_random_velocity(state);
  velocity->y = get_random_velocity(state);
}

void zero_velocity(Velocity *velocity, XOR256State *state) {
  velocity->x = 0.0f;
  velocity->y = 0.0f;
}

void update_positions(u8 *tags, Rect *rects, Velocity *velocities, u64 count, f32 delta) {
  persistent pthread_t threads[THREAD_COUNT]       = {0};
  persistent PositionThreadArgs args[THREAD_COUNT] = {0};

  u64 total_entities_count  = count - ZONE_COUNT;
  u64 thread_entities_count = (u64)(ceil((f64)total_entities_count / THREAD_COUNT));

  i32 result;
  u64 start = ZONE_COUNT;
  u64 end, args_count;

  for (u64 i = 0; i < THREAD_COUNT; ++i) {
    if (total_entities_count > thread_entities_count) {
      end = start + thread_entities_count;
      total_entities_count -= thread_entities_count;
    } else {
      end = start + total_entities_count;
      total_entities_count = 0;
    }

    args_count = end - start;

    args[i].zones      = &(rects[0]);
    args[i].tags       = &(tags[start]);
    args[i].rects      = &(rects[start]);
    args[i].velocities = &(velocities[start]);
    args[i].count      = args_count;
    args[i].delta      = delta;

    start += args_count;

    result = pthread_create(&(threads[i]), NULL, (PThreadRoutine)update_position_thread, (void *)&(args[i]));
    assert(result == 0);
  }

  for (u64 i = 0; i < THREAD_COUNT; ++i) {
    pthread_join(threads[i], NULL);
  }
}

void *update_position_thread(PositionThreadArgs *args) {
  persistent f32 multipliers[2] = {1.0f, 0.5f};
  persistent u64 inside_zone_mask = 0x7;

  u8 index = 0;
  f32 pos_x, pos_y;
  f32 max_x, max_y;

  for (u64 i = 0; i < args->count; ++i) {
    args->tags[i] &= inside_zone_mask;
    for (u64 j = 0; j < ZONE_COUNT; ++j) {
      args->tags[i] |= collides(&args->rects[i], &args->zones[j]) << INSIDE_ZONE_TAG_SHIFT;
    }

    index = (args->tags[i] & ENTITY_TAG_INSIDE_ZONE) >> INSIDE_ZONE_TAG_SHIFT;

    max_x = WIDTH - args->rects[i].scale.width;
    max_y = HEIGHT - args->rects[i].scale.height;

    pos_x = args->rects[i].position.x + args->velocities[i].x * multipliers[index] * args->delta;
    pos_y = args->rects[i].position.y + args->velocities[i].y * multipliers[index] * args->delta;

    if (pos_x < 0 || pos_x >= max_x) {
      pos_x = min(max(pos_x, 0), max_x);
      args->velocities[i].x *= -1;
    }

    if (pos_y < 0 || pos_y >= max_y) {
      pos_y = min(max(pos_y, 0), max_y);
      args->velocities[i].y *= -1;
    }

    args->rects[i].position.x = roundf(pos_x);
    args->rects[i].position.y = roundf(pos_y);
  }

  return args;
}

void render_entities(const u8 *tags, const Rect *rects, u64 count) {
  persistent RaylibDrawRectFunc renderers[2] = {DrawRectangle, DrawRectangleLines};
  persistent Color colors[2]                 = {FG_COLOR, ZONE_COLOR};

  for (u64 i = 0; i < count; ++i) {
    u8 func_index  = (tags[i] & ENTITY_TAG_COLLIDER) >> COLLIDER_TAG_SHIFT;
    u8 color_index = func_index | ((tags[i] & ENTITY_TAG_INSIDE_ZONE) >> INSIDE_ZONE_TAG_SHIFT);

    renderers[func_index](
      rects[i].position.x,
      rects[i].position.y,
      rects[i].scale.width,
      rects[i].scale.height,
      colors[color_index]
    );
  }
}

u8 collides(const Rect *rect, const Rect *collider) {
  i16 rect_min_x = rect->position.x + rect->scale.width;
  i16 rect_min_y = rect->position.y + rect->scale.height;
  i16 rect_max_x = rect->position.x;
  i16 rect_max_y = rect->position.y;

  i16 zone_x0 = collider->position.x;
  i16 zone_y0 = collider->position.y;
  i16 zone_x1 = collider->position.x + collider->scale.width;
  i16 zone_y1 = collider->position.y + collider->scale.height;

  return (u8)(rect_min_x > zone_x0 && rect_min_y > zone_y0 &&
          rect_max_x < zone_x1 && rect_max_y < zone_y1);
}

i16 get_random_velocity(XOR256State *state) {
  return (wapp_prng_xorshift_256(state) % (MAX_ABS_VELOCITY + 1 - MIN_ABS_VELOCITY)) + MIN_ABS_VELOCITY;
}