#include "ent.h"

typedef struct {
	u32 gen, dense_id, next_free;
} EntitySlot;

typedef struct {
	Entity *entity;
	EntitySlot *slot;
	u32 *dense_to_slot;
	u64 *delete_queue;
	u64 *pid_to_id;

	u64 next_pid;

	/* caps are count not byte size */
	u32 entity_cnt, entity_cap;
	u32 slot_cnt, slot_cap;
	u32 delete_queue_cnt, delete_queue_cap;
	u32 pid_to_id_cap;
	u32 free_slot;
} EntitySystem;

static EntitySystem en = {0};

static u64
entity_id_create(u32 slot, u32 gen) {
	return ((u64)gen << 32) | (u64)slot;
}

static u32
entity_slot(u64 id) {
	return (u32)id;
}

static u32
entity_gen(u64 id) {
	return (u32)(id >> 32);
}

static void
check_entity_cap(u32 req) {
	if (req <= en.entity_cap) {
		return;
	}

	while (en.entity_cap < req) {
		en.entity_cap *= 2;
	}

	u32 *tmp_dts = realloc(en.dense_to_slot, en.entity_cap * sizeof(u32));
	Entity *tmp_ent = realloc(en.entity, en.entity_cap * sizeof(Entity));
	if (!tmp_dts || !tmp_ent) {
		perror("err: entity array realloc failed.");
		exit(1);
	}

	en.dense_to_slot = tmp_dts;
	en.entity = tmp_ent;
}

static void
check_slot_cap(u32 req) {
	if (req <= en.slot_cap) {
		return;
	}

	while (en.slot_cap < req) {
		en.slot_cap *= 2;
	}

	EntitySlot *tmp = realloc(en.slot, en.slot_cap * sizeof(EntitySlot));
	if (!tmp) {
		perror("err: entity slot array realloc failed.");
		exit(1);
	}

	en.slot = tmp;
}

static void
check_delete_queue_cap(u32 req) {
	if (req <= en.delete_queue_cap) {
		return;
	}

	while (en.delete_queue_cap < req) {
		en.delete_queue_cap *= 2;
	}

	u64 *tmp = realloc(en.delete_queue, en.delete_queue_cap * sizeof(u64));
	if (!tmp) {
		perror("err: failed to realloc delete queue.");
		exit(1);
	}

	en.delete_queue = tmp;
}

static void
check_pid_to_id_cap(u64 pid) {
	if (pid < en.pid_to_id_cap) {
		return;
	}

	u32 limit = UINT32_MAX / sizeof(u64);
	if (pid >= limit) {
		TraceLog(LOG_FATAL, "pid exceeds addressable table size");
	}

	u32 start_cap = en.pid_to_id_cap;

	if (!en.pid_to_id_cap) {
		en.pid_to_id_cap = 256;
	}

	while (en.pid_to_id_cap <= pid) {
		en.pid_to_id_cap *= 2;
	}

	u64 *items = realloc(en.pid_to_id, en.pid_to_id_cap * sizeof(u64));
	if (!items) {
		perror("entity PID table alloc failed.");
		exit(1);
	}
	en.pid_to_id = items;
	memset(items + start_cap, 0,
		(en.pid_to_id_cap - start_cap) * sizeof(u64));
}


static EntitySlot*
entity_get_slot(u64 id) {
	EntitySlot *slot;
	u32 slot_id;
	u32 gen;

	if (id == 0) {
		return NULL;
	}

	slot_id = entity_slot(id);
	gen = entity_gen(id);

	if (slot_id >= en.slot_cnt) {
		return NULL;
	}

	slot = &en.slot[slot_id];

	if (slot->dense_id == ENTITY_INVALID_IDX ||
		slot->gen != gen ||
		slot->dense_id >= en.entity_cnt) {
		return NULL;
	}

	return slot;
}

static void
entity_clear_hierarchy(u64 id) {
	Entity *ent = entity_get(id);
	if (!ent) {
		return;
	}

	entity_set_parent(id, 0);

	while (ent->tr.first_child) {
		entity_set_parent(ent->tr.first_child, 0);
	}
}

static void
entity_delete_now(u64 id) {
	if (!entity_valid(id)) {
		return;
	}

	u64 child;
	while ((child = entity_get(id)->tr.first_child) != 0) {
		entity_delete_now(child);
	}

	Entity extracted = entity_extract_now(id);
	entity_clean(&extracted);
}


static void
entity_process_delete_queue(void) {
	for (int i = 0; i < en.delete_queue_cnt; ++i) {
		u64 id = en.delete_queue[i];

		if (entity_valid(id)) {
			entity_delete_now(id);
		}
	}
	en.delete_queue_cnt = 0;
}

void
entity_system_init(void) {
	en.entity_cap = 256;
	en.slot_cap = 256;
	en.delete_queue_cap = 32;
	en.entity = calloc(en.entity_cap, sizeof(Entity));
	en.slot = calloc(en.slot_cap, sizeof(EntitySlot));
	en.dense_to_slot = calloc(en.entity_cap, sizeof(u32));
	en.delete_queue = calloc(en.delete_queue_cap, sizeof(u64));
	en.free_slot = ENTITY_INVALID_IDX;
	en.next_pid = 1;
}

void
entity_system_exit(void) {
	for (u32 i = 0; i < en.entity_cnt; ++i) {
		entity_clean(&en.entity[i]);
	}
	free(en.entity);
	free(en.dense_to_slot);
	free(en.slot);
	free(en.delete_queue);
	free(en.pid_to_id);
	en = (EntitySystem) {0};
}

void
entity_system_clear(void) {
	for (u32 i = 0; i < en.entity_cnt; ++i) {
		entity_clean(&en.entity[i]);
	}

	memset(en.entity, 0, en.entity_cap * sizeof(Entity));
	memset(en.slot, 0, en.slot_cap * sizeof(EntitySlot));
	memset(en.dense_to_slot, 0, en.entity_cap * sizeof(u32));
	memset(en.delete_queue, 0, en.delete_queue_cap * sizeof(u64));
	memset(en.pid_to_id, 0, en.pid_to_id_cap * sizeof(u64));

	en.entity_cnt = 0;
	en.slot_cnt = 0;
	en.delete_queue_cnt = 0;
	en.free_slot = ENTITY_INVALID_IDX;
	en.next_pid = 1;


}

// static void
// entity_draw_mdl(Entity *ent) {
// 	Matrix T = MatrixTranslate(
// 				ent->tr.pos.x,
// 				ent->tr.pos.y,
// 				ent->tr.pos.z);

// 	Matrix R = MatrixRotateXYZ((Vec3) {
// 				ent->tr.rot.x * DEG2RAD,
// 				ent->tr.rot.y * DEG2RAD,
// 				ent->tr.rot.z * DEG2RAD,
// 			});

// 	Matrix S = MatrixScale(
// 				ent->tr.scale.x,
// 				ent->tr.scale.y,
// 				ent->tr.scale.z);

// 	Model mdl = *ent->mdl.mdl;
// 	mdl.transform = MatrixMultiply(S, MatrixMultiply(R, T));

// 	Color tint = ent->ed.selected ?
// 	                (Color){ 100, 100, 100, 255 } :
// 	                BLACK;

// 	DrawModel(mdl, (Vec3) {0}, 1, tint);
// }

// static void
// entity_draw_col(u64 id) {
// 	Entity *ent = entity_get(id);

// 	if (!b3Body_IsValid(ent->body_id)) {
// 		TraceLog(LOG_WARNING, "tried to draw invalid physics body");
// 		return;
// 	}

// 	if (b3Body_GetShapes(ent->body_id, &ent->shape_id, 1) == 0) {
// 		return;
// 	}

// 	const b3HullData *hull = b3Shape_GetHull(ent->shape_id);
// 	const b3Vec3 *points = b3GetHullPoints(hull);
// 	const b3HullHalfEdge *edges = b3GetHullEdges(hull);
// 	b3WorldTransform tr = b3Body_GetTransform(ent->body_id);

// 	for (int i = 0; i < hull->edgeCount; ++i) {
// 		if (i > edges[i].twin) {
// 			continue;
// 		}

// 		const b3HullHalfEdge *edge = &edges[i];
// 		const b3HullHalfEdge *twin = &edges[edge->twin];

// 		b3Pos world_a =
// 			b3TransformWorldPoint(tr, points[edge->origin]);
// 		b3Pos world_b =
// 			b3TransformWorldPoint(tr, points[twin->origin]);

// 		Color tint = GRAY;
// 		if (FLAG_CHECK(ent->flags, BIT(EFLAG_EDITOR_SELECTED))) {
// 			tint = LIME;
// 		}

// 		DrawLine3D(*(Vec3*)&world_a, *(Vec3*)&world_b, tint);
// 	}
// }

// static void
// entity_set_b3_transform(u64 id) {
// 	Entity *ent = entity_get(id);

// 	if (!b3Body_IsValid(ent->body_id)) {
// 		return;
// 	}

// 	b3WorldTransform tr = b3Body_GetTransform(ent->body_id);

// 	// if in edit mode, entity transform is authoritative
// 	if (FLAG_CHECK(app.flags, BIT(APP_FLAG_DEBUG))) {
// 		Quaternion eq = QuaternionFromEuler(
// 					ent->rot.x * DEG2RAD,
// 					ent->rot.y * DEG2RAD,
// 					ent->rot.z * DEG2RAD);

// 		Quaternion bq = {
// 			.x = tr.q.v.x,
// 			.y = tr.q.v.y,
// 			.z = tr.q.v.z,
// 			.w = tr.q.s
// 		};

// 		if (Vector3Equals(ent->pos, *(Vec3*)&tr.p) &&
// 			QuaternionEquals(eq, bq)) {
// 			return;
// 		}

// 		b3Quat b3rot = { .v = { eq.x, eq.y, eq.z }, .s = eq.w };
// 		b3Body_SetTransform(ent->body_id, *(b3Pos*)&ent->pos, b3rot);
// 	} else {
// 		ent->pos.x = tr.p.x;
// 		ent->pos.y = tr.p.y;
// 		ent->pos.z = tr.p.y;

// 		Vec3 euler_angles =
// 			QuaternionToEuler((Quaternion) {
// 				.x = tr.q.v.x,
// 				.y = tr.q.v.y,
// 				.z = tr.q.v.z,
// 				.w = tr.q.s
// 			});

// 		ent->rot.x = euler_angles.x * RAD2DEG;
// 		ent->rot.y = euler_angles.y * RAD2DEG;
// 		ent->rot.z = euler_angles.z * RAD2DEG;
// 	}
// }

void
entity_system_update(void) {
	entity_process_delete_queue();
}

void
entity_system_draw_3d(void) {
	// for (u32 i = 0; i < entity_count(); ++i) {
	// 	u64 id = entity_id_at(i);
	// 	Entity *ent = entity_get(id);

	// 	if (ent->mdl.has && ent->mdl.draw) {
	// 		entity_draw_mdl(ent);
	// 	}
	// }
}

void
entity_system_set_next_pid(u64 pid) {
	en.next_pid = pid;
}

void
entities_clear_selected(void) {
	for (int i = 0; i < en.entity_cnt; ++i) {
		Entity *ent = entity_get(entity_id_at(i));
		ent->ed.selected = false;
	}
}

u64
entity_create(bool transient) {
	check_entity_cap(en.entity_cnt + 1);

	u32 slot_id = 0;
	EntitySlot *slot;

	if (en.free_slot != ENTITY_INVALID_IDX) {
		slot_id = en.free_slot;
		slot = &en.slot[slot_id];
		en.free_slot = slot->next_free;
	} else {
		check_slot_cap(en.slot_cnt + 1);
		slot_id = en.slot_cnt;
		en.slot_cnt += 1;
		slot = &en.slot[slot_id];
		slot->gen = 1;
	}

	u32 dense_id = en.entity_cnt;
	en.entity_cnt += 1;

	slot->dense_id = dense_id;
	slot->next_free = ENTITY_INVALID_IDX;
	en.dense_to_slot[dense_id] = slot_id;

	en.entity[dense_id] = (Entity){0};
	en.entity[dense_id].tr.rot = QuaternionIdentity();
	en.entity[dense_id].tr.scale = Vector3One();

	u64 id = entity_id_create(slot_id, slot->gen);
  	if (!transient) {
  		u64 pid = en.next_pid++;
  		check_pid_to_id_cap(pid);
  		en.entity[dense_id].name = strdup("entity");
  		en.entity[dense_id].pid = pid;
  		en.pid_to_id[pid] = id;
  	}
	return id;
}

u64
entity_restore(u64 pid) {
	if (!pid || entity_id_from_pid(pid)) {
		return 0;
	}

	check_pid_to_id_cap(pid);

	u64 id = entity_create(true);
	Entity *ent = entity_get(id);
	ent->pid = pid;
	en.pid_to_id[pid] = id;
	en.next_pid = MAX(en.next_pid, pid + 1);

	return id;
}

/* deferred delete */
void
entity_delete(u64 id) {
	if (!entity_valid(id)) {
		return;
	}

	check_delete_queue_cap(en.delete_queue_cnt + 1);
	en.delete_queue[en.delete_queue_cnt] = id;
	en.delete_queue_cnt += 1;
}

static void
entity_physics_clear(Entity *ent) {
	PhysicsCmp *phys = &ent->phys;

	if (phys->runtime.joints) {
		for (int i = 0; i < phys->joint_cnt; ++i) {
			b3JointId joint_id = phys->runtime.joints[i];
			if (b3Joint_IsValid(joint_id)) {
				b3DestroyJoint(joint_id, false);
			}
		}
	}

	if (phys->runtime.bodies) {
		for (int i = 0; i < phys->body_cnt; ++i) {
			PhysicsBodyRuntime *bdy = &phys->runtime.bodies[i];
			if (b3Body_IsValid(bdy->id)) {
				b3DestroyBody(bdy->id);
			}
			free(bdy->shapes);
		}
	}

	free(phys->runtime.joints);
	free(phys->runtime.bodies);
	phys->runtime.joints = NULL;
	phys->runtime.bodies = NULL;
}

void
entity_clean(Entity *ent) {
	free(ent->name);

	/* model component data */
	free(ent->mdl.key);

	/* physics component data */
	entity_physics_clear(ent); // runtime data
	for (int i = 0; i < ent->phys.body_cnt; ++i) {
		PhysicsBody *body = &ent->phys.bodies[i];

		for (int j = 0; j < body->shape_cnt; ++j) {
			PhysicsShape *shape = &body->shapes[j];

			switch (shape->tag) {
			case PHYS_SHAPE_CONVEX_HULL:
			free(shape->ConvexHull.points);
			break;

			case PHYS_SHAPE_STATIC_TRIMESH:
			free(shape->TriMesh.points);
			free(shape->TriMesh.indices);
			break;

			default:
			break;
			}
		}

		free(body->shapes);
	}

	free(ent->phys.bodies);
	free(ent->phys.joints);
}

Entity
entity_extract_now(u64 id) {
	EntitySlot *slot = entity_get_slot(id);
	assert(slot);

	u32 slot_id = entity_slot(id);
	u32 dense_id = slot->dense_id;
	u32 last_dense_id = en.entity_cnt - 1;

	assert(!en.entity[dense_id].tr.first_child);

	entity_physics_clear(&en.entity[dense_id]);
	entity_set_parent(id, 0);
	Entity extracted = en.entity[dense_id];

	if (extracted.pid) {
		en.pid_to_id[extracted.pid] = 0;
	}

	if (dense_id != last_dense_id) {
		u32 moved_slot_id = en.dense_to_slot[last_dense_id];
		en.entity[dense_id] = en.entity[last_dense_id];
		en.dense_to_slot[dense_id] = moved_slot_id;
		en.slot[moved_slot_id].dense_id = dense_id;
	}

	en.entity[last_dense_id] = (Entity){0};
	en.dense_to_slot[last_dense_id] = ENTITY_INVALID_IDX;
	en.entity_cnt--;

	slot->dense_id = ENTITY_INVALID_IDX;
	slot->gen++;

	if (slot->gen == 0) {
		slot->gen = 1;
	}

	slot->next_free = en.free_slot;
	en.free_slot = slot_id;

	return extracted;
}

static void
entity_unlink_parent(u64 id) {
	Entity *ent = entity_get(id);
	if (!ent) {
		return;
	}

	TransformCmp *tr = &ent->tr;

	if (tr->prev) {
		Entity *prev = entity_get(tr->prev);
		if (prev) {
			prev->tr.next = tr->next;
		}
	} else if (tr->parent) {
		Entity *parent = entity_get(tr->parent);
		if (parent && parent->tr.first_child == id) {
			parent->tr.first_child = tr->next;
		}
	}

	if (tr->next) {
		Entity *next = entity_get(tr->next);
		if (next) {
			next->tr.prev = tr->prev;
		}
	}

	tr->parent = 0;
	tr->prev = 0;
	tr->next = 0;
}

static bool
entity_is_ancestor(u64 ancestor_id, u64 id) {
	while (id) {
		if (id == ancestor_id) {
			return true;
		}

		Entity *ent = entity_get(id);
		if (!ent) {
			return false;
		}

		id = ent->tr.parent;
	}

	return false;
}

bool
entity_set_parent(u64 id, u64 parent_id) {
	Entity *ent = entity_get(id);
	if (!ent) {
		TraceLog(LOG_WARNING, "entity_set_parent: entity id null");
		return false;
	}

	if (parent_id == ent->tr.parent) {
		return true;
	}

	Entity *parent = NULL;
	if (parent_id) {
		parent = entity_get(parent_id);
		if (!parent ||
			parent_id == id ||
			entity_is_ancestor(id, parent_id)) {
			TraceLog(LOG_ERROR, "entity hierarchy failed.");
			return false;
		}
	}

	entity_unlink_parent(id);

	if (!parent) {
		return true;
	}

	u64 first_child = parent->tr.first_child;
	ent->tr.parent = parent_id;
	ent->tr.next = first_child;

	if (first_child) {
		Entity *first = entity_get(first_child);
		if (first) {
			first->tr.prev = id;
		}
	}

	parent->tr.first_child = id;
	return true;
}

bool
entity_valid(u64 id) {
	return entity_get_slot(id) != NULL;
}

Entity*
entity_get(u64 id) {
	EntitySlot *slot = entity_get_slot(id);
	if (!slot) {
		return NULL;
	}

	return &en.entity[slot->dense_id];
}

u64
entity_id_at(u32 dense_id) {
	if (dense_id >= en.entity_cnt) {
		return 0;
	}

	u32 slot_id = en.dense_to_slot[dense_id];
	EntitySlot *slot = &en.slot[slot_id];

	return entity_id_create(slot_id, slot->gen);
}

u64
entity_id_from_pid(u64 pid) {
	return pid < en.pid_to_id_cap ? en.pid_to_id[pid] : 0;
}

Entity*
entity_data(void) {
	return en.entity;
}

u32
entity_count(void) {
	return en.entity_cnt;
}

u64
entity_next_pid(void) {
	return en.next_pid;
}

/* entity transform section */
Matrix
entity_world_matrix(u64 id) {
	Matrix world = MatrixIdentity();

	while (id) {
		Entity *ent = entity_get(id);
		assert(ent);

		TransformCmp *tr = &ent->tr;

		Matrix scale = MatrixScale(
			tr->scale.x, tr->scale.y, tr->scale.z);
		Matrix rotation = QuaternionToMatrix(tr->rot);
		Matrix translation = MatrixTranslate(
			tr->pos.x, tr->pos.y, tr->pos.z);

		/* Raymath order: apply scale, rotation, translation. */
		Matrix local = MatrixMultiply(
			MatrixMultiply(scale, rotation), translation);

		world = MatrixMultiply(world, local);
		id = tr->parent;
	}

	return world;
}

Vec3
entity_point_to_world(u64 space, Vec3 point) {
	return Vector3Transform(point, entity_world_matrix(space));
}

bool
entity_point_to_local(u64 space, Vec3 point, Vec3 *out) {
	if (space) {
		Entity *ent = entity_get(space);
		assert(ent);

		TransformCmp *tr = &ent->tr;
		if (!tr->scale.x || !tr->scale.y || !tr->scale.z) {
			return false;
		}

		if (!entity_point_to_local(tr->parent, point, &point)) {
			return false;
		}
		point = Vector3Subtract(point, tr->pos);
		point = Vector3RotateByQuaternion(point,
		                        	QuaternionInvert(tr->rot));
		point = Vector3Divide(point, tr->scale);
	}

	*out = point;
	return true;
}

Vec3
entity_world_pos(u64 id) {
	return entity_point_to_world(id, Vector3Zero());
}

Quaternion
entity_world_rot(u64 id) {
	Quaternion rot = QuaternionIdentity();

	while (id) {
		Entity *ent = entity_get(id);
		assert(ent);

		rot = QuaternionMultiply(ent->tr.rot, rot);
		id = ent->tr.parent;
	}

	return QuaternionNormalize(rot);
}

void
entity_set_world_pos(Entity *ent, Vec3 pos) {
	if (!entity_point_to_local(ent->tr.parent, pos, &ent->tr.pos)) {
		TraceLog(LOG_WARNING, "Cannot set world position through a "
			"zero-scale ancestor.");
	}
}

void
entity_set_world_rot(Entity *ent, Quaternion rot) {
	f32 qlen = QuaternionLength(rot);
	if (!ent || !qlen || !isfinite(qlen)) {
		return;
	}

	ent->tr.rot =
		QuaternionNormalize(QuaternionMultiply(
	                QuaternionInvert(entity_world_rot(ent->tr.parent)),
	        		QuaternionNormalize(rot)));
}