#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include "../src/wapp/wapp.h"

/* ---------------------------------------------------------------------------
 * Pool allocator (arena-backed, intrusive free list)
 * -------------------------------------------------------------------------*/

typedef struct PrPoolFreeNode PrPoolFreeNode;
struct PrPoolFreeNode {
	PrPoolFreeNode *next;
};

typedef struct {
	WpAllocator    *allocator;
	PrPoolFreeNode *free_head;
	u64             slot_size;
} PrPool;

static void prPoolInit(PrPool *pool, WpAllocator *arena, u64 slot_size) {
	pool->allocator     = arena;
	pool->free_head = NULL;
	pool->slot_size = slot_size;
}

static void *prPoolAlloc(PrPool *pool) {
	if (pool->free_head) {
		PrPoolFreeNode *node = pool->free_head;
		pool->free_head = node->next;
		return node;
	}
	return wpMemAllocatorAlloc(pool->allocator, pool->slot_size);
}

static void prPoolFree(PrPool *pool, void *slot) {
	if (!slot) { return; }
	PrPoolFreeNode *node = (PrPoolFreeNode *)slot;
	node->next  = pool->free_head;
	pool->free_head = node;
}

/* ---------------------------------------------------------------------------
 * Graph — unified nodes with values + forward/backward adjacency
 *
 * - Nodes live in a flat array; free slots are linked via next_free.
 * - Each node has a generation counter bumped on free; handles are
 *   { index, generation } so stale references are detectable.
 * - Edges are pool-allocated PrEdgeNode structs with both forward
 *   and backward linkage so deletion traces both directions.
 * -------------------------------------------------------------------------*/

#define PR_INVALID_INDEX ((u64)-1)

typedef struct PrEdgeNode PrEdgeNode;
struct PrEdgeNode {
	PrEdgeNode *pool_next;       /* used by pool free list when freed */
	PrEdgeNode *next_forward;    /* chain in source's forward list */
	PrEdgeNode *next_backward;   /* chain in target's backward list */
	u64         source_idx;
	u64         target_idx;
};

typedef struct {
	u64 index;
	u64 generation;
} PrHandle;

typedef struct {
	int         value;
	u64         generation;      /* bumped on every free */
	u64         next_free;       /* free-list index (PR_INVALID_INDEX = active) */
	PrEdgeNode *forward_head;    /* outgoing edges */
	PrEdgeNode *backward_head;   /* incoming edges */
} PrGraphNode;

typedef struct {
	PrGraphNode *nodes;          /* WpArray of PrGraphNode */
	u64          max_nodes;
	u64          max_ever;       /* highest index ever allocated */
	u64          free_head;      /* PR_INVALID_INDEX = empty */
	u64          count;          /* active node count */
	PrPool       edge_pool;
} PrGraph;

/* --- initialisation ---------------------------------------------------- */

static void prGraphInit(PrGraph *g, WpAllocator *arena, u64 max_nodes) {
	g->nodes     = wpArrayAllocCapacity(PrGraphNode, arena, max_nodes, WP_ARRAY_INIT_FILLED);
	g->max_nodes = max_nodes;
	g->max_ever  = 0;
	g->free_head = PR_INVALID_INDEX;
	g->count     = 0;
	prPoolInit(&g->edge_pool, arena, sizeof(PrEdgeNode));

	/* Build the free list — last slot's next_free stays PR_INVALID_INDEX */
	for (u64 i = 0; i < max_nodes; i++) {
		g->nodes[i].next_free = (i < max_nodes - 1) ? i + 1 : PR_INVALID_INDEX;
	}
	g->free_head = 0;
}

/* --- handle validation ------------------------------------------------ */

static b8 prHandleValid(PrGraph *g, PrHandle h) {
	if (h.index >= g->max_nodes) { return false; }
	PrGraphNode *node = &g->nodes[h.index];
	return node->generation == h.generation && node->next_free == PR_INVALID_INDEX;
}

/* --- node management -------------------------------------------------- */

static PrHandle prNodeAdd(PrGraph *g, int value) {
	if (g->free_head == PR_INVALID_INDEX) {
		return (PrHandle){ PR_INVALID_INDEX, 0 };
	}

	u64         idx  = g->free_head;
	PrGraphNode *n   = &g->nodes[idx];
	g->free_head     = n->next_free;

	n->value         = value;
	n->next_free     = PR_INVALID_INDEX;
	n->forward_head  = NULL;
	n->backward_head = NULL;
	g->count++;

	if (idx >= g->max_ever) { g->max_ever = idx + 1; }

	return (PrHandle){ idx, n->generation };
}

static int *prNodeGetValue(PrGraph *g, PrHandle h) {
	if (!prHandleValid(g, h)) { return NULL; }
	return &g->nodes[h.index].value;
}

/* --- internal: unlink edge helpers ------------------------------------ */

static PrEdgeNode *_unlinkForward(PrGraph *g, u64 from_idx, u64 target_idx) {
	PrGraphNode *src = &g->nodes[from_idx];
	PrEdgeNode *prev = NULL;
	PrEdgeNode *curr = src->forward_head;
	while (curr) {
		if (curr->target_idx == target_idx) {
			if (prev) { prev->next_forward  = curr->next_forward; }
			else      { src->forward_head   = curr->next_forward; }
			return curr;
		}
		prev = curr;
		curr = curr->next_forward;
	}
	return NULL;
}

static void _unlinkBackward(PrGraph *g, u64 to_idx, u64 source_idx) {
	PrGraphNode *dst = &g->nodes[to_idx];
	PrEdgeNode *prev = NULL;
	PrEdgeNode *curr = dst->backward_head;
	// This is intended to handle cases where the target node might not be the one
	// immediately feeding the node represented by to_idx
	while (curr) {
		if (curr->source_idx == source_idx) {
			if (prev) { prev->next_backward = curr->next_backward; }
			else      { dst->backward_head  = curr->next_backward; }
			return;
		}
		prev = curr;
		curr = curr->next_backward;
	}
}

/* --- edge management -------------------------------------------------- */

static b8 prEdgeAdd(PrGraph *g, PrHandle from, PrHandle to) {
	if (!prHandleValid(g, from) || !prHandleValid(g, to)) { return false; }

	PrEdgeNode *edge = (PrEdgeNode *)prPoolAlloc(&g->edge_pool);
	if (!edge) { return false; }

	edge->source_idx = from.index;
	edge->target_idx = to.index;

	PrGraphNode *src = &g->nodes[from.index];
	edge->next_forward  = src->forward_head;
	src->forward_head    = edge;

	PrGraphNode *dst = &g->nodes[to.index];
	edge->next_backward = dst->backward_head;
	dst->backward_head   = edge;

	return true;
}

/* --- node removal (tears down incident edges) ------------------------- */

static void prNodeRemove(PrGraph *g, PrHandle h) {
	if (!prHandleValid(g, h)) { return; }

	u64         idx  = h.index;
	PrGraphNode *n   = &g->nodes[idx];

	/* Free outgoing edges:  unlink from each target's backward list */
	PrEdgeNode *edge = n->forward_head;
	while (edge) {
		PrEdgeNode *next = edge->next_forward;
		_unlinkBackward(g, edge->target_idx, idx);
		prPoolFree(&g->edge_pool, edge);
		edge = next;
	}

	/* Free incoming edges:  unlink from each source's forward list */
	edge = n->backward_head;
	while (edge) {
		PrEdgeNode *next = edge->next_backward;
		/* edge is also in the source's forward list — unlink by target_idx */
		_unlinkForward(g, edge->source_idx, idx);
		prPoolFree(&g->edge_pool, edge);
		edge = next;
	}

	/* Return node slot to free list with bumped generation */
	n->generation++;
	n->next_free   = g->free_head;
	g->free_head   = idx;
	g->count--;
}

/* --- traversal helpers ------------------------------------------------ */

static void prDump(PrGraph *g) {
	printf("Active nodes: %llu\n\n", (unsigned long long)g->count);

	for (u64 i = 0; i < g->max_ever; i++) {
		PrGraphNode *n = &g->nodes[i];
		if (n->next_free != PR_INVALID_INDEX) { continue; }

		printf("  [%llu] g=%llu  value=%d\n",
		       (unsigned long long)i,
		       (unsigned long long)n->generation,
		       n->value);

		printf("    forward  ──▶");
		PrEdgeNode *e = n->forward_head;
		if (!e) { printf(" (none)"); }
		while (e) {
			PrGraphNode *tv = &g->nodes[e->target_idx];
			printf(" %llu[%d]", (unsigned long long)e->target_idx, tv->value);
			e = e->next_forward;
			if (e) { printf(","); }
		}
		printf("\n");

		printf("    backward ◀──");
		e = n->backward_head;
		if (!e) { printf(" (none)"); }
		while (e) {
			PrGraphNode *sv = &g->nodes[e->source_idx];
			printf(" %llu[%d]", (unsigned long long)e->source_idx, sv->value);
			e = e->next_backward;
			if (e) { printf(","); }
		}
		printf("\n");
	}
}

/* ---------------------------------------------------------------------------
 * Demo
 * -------------------------------------------------------------------------*/

int main(void) {
	srand((unsigned)time(NULL));

	WpAllocator arena = wpMemArenaAllocatorInitZero(KiB(64));
	if (wpMemAllocatorInvalid(&arena)) {
		fprintf(stderr, "arena init failed\n");
		return 1;
	}

	PrGraph g;
	prGraphInit(&g, &arena, 16);

	printf("=== Add 6 nodes ===\n");
	PrHandle nodes[10];
	for (u64 i = 0; i < 6; i++) {
		nodes[i] = prNodeAdd(&g, rand() % 100);
		printf("  node[%llu] = handle{%llu g%llu} val=%d\n",
		       (unsigned long long)i,
		       (unsigned long long)nodes[i].index,
		       (unsigned long long)nodes[i].generation,
		       *prNodeGetValue(&g, nodes[i]));
	}

	printf("\n=== Add edges: 0→1, 0→2, 1→3, 2→3, 3→4, 4→5 ===\n");
	u64 edge_list[][2] = { {0,1}, {0,2}, {1,3}, {2,3}, {3,4}, {4,5} };
	for (u64 i = 0; i < 6; i++) {
		u64 f = edge_list[i][0], t = edge_list[i][1];
		prEdgeAdd(&g, nodes[f], nodes[t]);
	}
	prDump(&g);

	printf("=== Remove node 2 (index %llu) ===\n",
	       (unsigned long long)nodes[2].index);
	prNodeRemove(&g, nodes[2]);
	prDump(&g);

	printf("=== Stale check ===\n");
	printf("  nodes[2] handle{%llu g%llu} valid? %s\n",
	       (unsigned long long)nodes[2].index,
	       (unsigned long long)nodes[2].generation,
	       prHandleValid(&g, nodes[2]) ? "YES" : "NO");

	printf("=== Add edge 5→0 (reuses freed edge slots) ===\n");
	prEdgeAdd(&g, nodes[5], nodes[0]);
	prDump(&g);

	printf("=== Add a new node (reuses freed node slot) ===\n");
	nodes[6] = prNodeAdd(&g, 42);
	printf("  new node = handle{%llu g%llu} val=%d\n",
	       (unsigned long long)nodes[6].index,
	       (unsigned long long)nodes[6].generation,
	       *prNodeGetValue(&g, nodes[6]));
	prDump(&g);

	wpMemArenaAllocatorDestroy(&arena);
	return 0;
}