// vim:fileencoding=utf-8:foldmethod=marker

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

wp_intern WpLogger _log = { .name = wpStr8LitRo("dag_man") };

/* ---------------------------------------------------------------------------
 * 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;

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

wp_intern 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);
}

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

/* ---------------------------------------------------------------------------
 * Graph
 * -------------------------------------------------------------------------*/

#define INVALID_NODE_INDEX (u64)-1
#define INVALID_NODE_ID ((PrNodeId){ .index = INVALID_NODE_INDEX, .generation = INVALID_NODE_INDEX })

typedef enum {
	PR_NODE_TYPE_NONE,
	PR_NODE_TYPE_READ,
	PR_NODE_TYPE_BLUR,
	PR_NODE_TYPE_GRADE,

	COUNT_NODE_TYPES
} PrNodeType;

typedef struct {
	u64 index;
	u64 generation;
} PrNodeId;
typedef PrNodeId *PrNodeIdArray;

typedef struct {
	union {
		WpStr8   path;
		f32      blur;
		f32      gain;
	} params;
	PrNodeType type;
	u64        generation;
	u64        next_free;
} PrNode;
typedef PrNode *PrNodeArray;

/* ---------------------------------------------------------------------------
 * Graph Edge/Vertex type
 * If source or target are INVALID_NODE_ID, it should be treated as a vertex.
 * Otherwise, it should be treated as an edge. Think of it like homogeneous
 * coordinates where a vector {0, 0, 0, 0} is treated as a direction, while
 * a vector {0, 0, 0, 1} is treated as a point
 * ---------------------------------------------------------------------------*/
typedef struct PrEdgeVertex PrEdgeVertex;
struct PrEdgeVertex {
	PrEdgeVertex *next_forward;
	PrEdgeVertex *next_backward;
	PrNodeId      source;
	PrNodeId      target;
};
typedef PrEdgeVertex *PrEdgeVertexArray;

typedef struct {
	PrPool            vertex_pool;
	PrNodeArray       nodes;
	PrEdgeVertexArray vertices;
	u64               capacity;
	u64               max_count_ever;
	u64               count;
	u64               free_head;
} PrGraph;

wp_intern b8 prGraphIsStaleId(const PrGraph *graph, PrNodeId id) {
	u64 generation = graph->nodes[id.index].generation;
	return id.generation != generation;
}

wp_intern b8 prGraphIsActiveNode(const PrGraph *graph, PrNodeId id) {
	u64 next_free  = graph->nodes[id.index].next_free;
	return !prGraphIsStaleId(graph, id) && next_free == INVALID_NODE_INDEX;
}

wp_intern PrNodeId prGraphGetNode(const PrGraph *graph, u64 index) {
	return (PrNodeId){ .index = index, .generation = graph->nodes[index].generation };
}

wp_intern void prGraphInit(PrGraph *graph, WpAllocator *allocator, u64 capacity) {
	graph->nodes          = wpArrayAllocCapacity(PrNode, allocator, capacity, WP_ARRAY_INIT_FILLED);
	graph->vertices       = wpArrayAllocCapacity(PrEdgeVertex, allocator, capacity, WP_ARRAY_INIT_FILLED);
	graph->free_head      = 0;
	graph->capacity       = capacity;
	graph->max_count_ever = 0;
	graph->count          = 0;

	if (!graph->nodes || !graph->vertices) {
		graph->capacity = 0;
		return;
	}

	prPoolInit(&graph->vertex_pool, allocator, sizeof(PrEdgeVertex));
	memset(graph->nodes, 0, capacity * sizeof(PrNode));
	memset(graph->vertices, 0, capacity * sizeof(PrEdgeVertex));

	for (u64 i = 0; i < capacity; ++i) {
		graph->nodes[i].next_free = i < capacity - 1 ? i + 1 : INVALID_NODE_INDEX;
		graph->vertices[i].source = graph->vertices[i].target = INVALID_NODE_ID;
	}
}

/* ---------------------------------------------------------------------------
 * Kahn's algorithm — topological sort / cycle detection
 *
 * Returns a WpArray of PrNodeId (sorted topologically). If the result count
 * is less than graph->count, the graph contains a cycle.
 * -------------------------------------------------------------------------*/

wp_intern PrNodeId *prGraphTopologicalSort(PrGraph *graph, const WpAllocator *allocator) {
	if (graph->count == 0) { return NULL; }
	if (!graph->nodes || graph->capacity == 0) { return NULL; }

	PrNodeIdArray result = wpArrayAllocCapacity(PrNodeId, allocator, graph->count, WP_ARRAY_INIT_NONE);
	if (!result) { return NULL; }

	WpAllocator local_arena = wpMemArenaAllocatorInitZero(KiB(16));

	WpU64Array in_degree = wpArrayAllocCapacity(u64, &local_arena, graph->capacity, WP_ARRAY_INIT_FILLED);
	if (!in_degree) { return result; }
	memset(in_degree, 0, wpArrayCapacity(in_degree) * sizeof(u64));

	for (u64 i = 0; i < graph->max_count_ever; i++) {
		PrNodeId id = prGraphGetNode(graph, i);
		if (!prGraphIsActiveNode(graph, id)) { continue; }

		PrEdgeVertex *vertex = &graph->vertices[i];
		PrEdgeVertex *curr   = vertex->next_forward;
		while (curr) {
			in_degree[curr->target.index]++;
			curr = curr->next_forward;
		}
	}

	WpQueue queue = wpQueueAlloc(u64, &local_arena, graph->count);

	for (u64 i = 0; i < graph->max_count_ever; i++) {
		PrNodeId id = prGraphGetNode(graph, i);
		if (!prGraphIsActiveNode(graph, id)) { continue; }
		if (in_degree[i] == 0) {
			wpQueuePush(u64, &queue, &i);
		}
	}

	while (queue.count > 0) {
		u64 *node_idx = wpQueuePop(u64, &queue);
		if (!node_idx) { break; }

		PrNodeId id = prGraphGetNode(graph, *node_idx);
		wpArrayAppendCapped(PrNodeId, result, &id);

		PrEdgeVertex *vertex = &graph->vertices[*node_idx];
		PrEdgeVertex *curr   = vertex->next_forward;
		while (curr) {
			u64 target_idx = curr->target.index;
			if (in_degree[target_idx] > 0) {
				in_degree[target_idx]--;
				if (in_degree[target_idx] == 0) {
					wpQueuePush(u64, &queue, &target_idx);
				}
			}
			curr = curr->next_forward;
		}
	}

	return result;
}

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

wp_intern void _unlinkForward(PrGraph *graph, PrNodeId from, PrEdgeVertex *edge) {
	if (!prGraphIsActiveNode(graph, from)) { return; }

	PrEdgeVertex *src  = &graph->vertices[from.index];
	PrEdgeVertex *curr = src->next_forward;
	PrEdgeVertex *prev = NULL;
	while (curr) {
		if (curr == edge) {
			if (prev) {
				prev->next_forward = curr->next_forward;
			} else {
				src->next_forward = curr->next_forward;
			}
			return;
		}

		prev = curr;
		curr = curr->next_forward;
	}
}

wp_intern void _unlinkBackward(PrGraph *graph, PrNodeId to, PrEdgeVertex *edge) {
	if (!prGraphIsActiveNode(graph, to)) { return; }

	PrEdgeVertex *dst  = &graph->vertices[to.index];
	PrEdgeVertex *curr = dst->next_backward;
	PrEdgeVertex *prev = NULL;
	while (curr) {
		if (curr == edge) {
			if (prev) {
				prev->next_backward = curr->next_backward;
			} else {
				dst->next_backward = curr->next_backward;
			}
			return;
		}

		prev = curr;
		curr = curr->next_backward;
	}
}

wp_intern void prGraphAddEdge(PrGraph *graph, PrNodeId from, PrNodeId to) {
	if (!prGraphIsActiveNode(graph, from) || !prGraphIsActiveNode(graph, to)) { return; }
	if (from.index == to.index) { return; }

	PrEdgeVertex *edge = (PrEdgeVertex *)prPoolAlloc(&graph->vertex_pool);
	if (!edge) { return; }

	edge->source = from;
	edge->target = to;

	/* Link into adjacency chains */
	PrEdgeVertex *src   = &graph->vertices[from.index];
	PrEdgeVertex *dst   = &graph->vertices[to.index];
	edge->next_forward  = src->next_forward;
	edge->next_backward = dst->next_backward;
	src->next_forward   = edge;
	dst->next_backward  = edge;

	/* Check whether the new edge created a cycle */
	WpAllocator scratch  = wpMemArenaAllocatorInitZero(KiB(16));
	PrNodeIdArray sorted = prGraphTopologicalSort(graph, &scratch);
	u64 sorted_count     = sorted ? wpArrayCount(sorted) : 0;
	if (sorted_count < graph->count) {
		_unlinkForward(graph, from, edge);
		_unlinkBackward(graph, to, edge);
		prPoolFree(&graph->vertex_pool, edge);
	}

	wpMemArenaAllocatorDestroy(&scratch);
}

wp_intern PrNodeId prGraphAddNode(PrGraph *graph, PrNodeType type) {
	u64 idx = graph->free_head;
	if (idx == INVALID_NODE_INDEX) { return INVALID_NODE_ID; }

	PrNode *node = &graph->nodes[idx];

	graph->free_head = node->next_free;
	node->next_free  = INVALID_NODE_INDEX;
	node->type       = type;
	memset(&node->params, 0, sizeof(node->params));

	graph->count++;

	if (idx + 1 > graph->max_count_ever) { graph->max_count_ever = idx + 1; }

	return (PrNodeId){ .index = idx, .generation = node->generation };
}

wp_intern void prGraphRemoveNode(PrGraph *graph, PrNodeId id) {
	if (!prGraphIsActiveNode(graph, id)) { return; }

	PrNode *node = &graph->nodes[id.index];
	u64 next_gen = node->generation + 1;
	memset(node, 0, sizeof(PrNode));

	node->generation = next_gen;
	node->next_free  = graph->free_head;
	graph->free_head = id.index;

	PrEdgeVertex *vertex = &graph->vertices[id.index];
	if (!(vertex->next_forward) && !(vertex->next_backward)) {
		goto REMOVE_NODE_UNSET_POINTERS;
	}

	PrEdgeVertex *curr = vertex->next_forward;
	while (curr) {
		PrEdgeVertex *next = curr->next_forward;
		_unlinkBackward(graph, curr->target, curr);
		prPoolFree(&graph->vertex_pool, curr);
		curr = next;
	}

	curr = vertex->next_backward;
	while (curr) {
		PrEdgeVertex *trailing = curr->next_backward;
		_unlinkForward(graph, curr->source, curr);
		prPoolFree(&graph->vertex_pool, curr);
		curr = trailing;
	}

REMOVE_NODE_UNSET_POINTERS:
	vertex->next_forward = vertex->next_backward = NULL;

	graph->count--;
}

wp_intern void prGraphDump(const PrGraph *graph) {
	printf("==============INPUTS==============\n");
	for (u64 i = 0; i < graph->max_count_ever; ++i) {
		PrNodeId id = prGraphGetNode(graph, i);
		if (!prGraphIsActiveNode(graph, id)) { continue; }

		PrEdgeVertex *vertex = &graph->vertices[i];

		printf("%" PRIu64 ":", id.index + 1);

		if (!(vertex->next_backward)) {
			printf(" (none)");
		} else {
			PrEdgeVertex *curr = vertex->next_backward;
			while (curr) {
				printf(" %" PRIu64, curr->source.index + 1);
				curr = curr->next_backward;
			}
		}

		printf("\n");
	}

	printf("==============OUTPUTS==============\n");
	for (u64 i = 0; i < graph->max_count_ever; ++i) {
		PrNodeId id = prGraphGetNode(graph, i);
		if (!prGraphIsActiveNode(graph, id)) { continue; }

		PrEdgeVertex *vertex = &graph->vertices[i];

		printf("%" PRIu64 ":", id.index + 1);

		if (!(vertex->next_forward)) {
			printf(" (none)");
		} else {
			PrEdgeVertex *curr = vertex->next_forward;
			while (curr) {
				printf(" %" PRIu64, curr->target.index + 1);
				curr = curr->next_forward;
			}
		}

		printf("\n");
	}
}

i32 main(void) {
	WpAllocator arena = wpMemArenaAllocatorInitZero(MiB(16));
	if (wpMemAllocatorInvalid(&arena)) {
		wpLogFatal(&_log, wpStr8Lit("arena init failed"));
		return 1;
	}

	PrGraph graph = {0};
	prGraphInit(&graph, &arena, 128);

	PrNodeId n1 = prGraphAddNode(&graph, PR_NODE_TYPE_READ);
	PrNodeId n2 = prGraphAddNode(&graph, PR_NODE_TYPE_READ);
	PrNodeId n3 = prGraphAddNode(&graph, PR_NODE_TYPE_READ);
	PrNodeId n4 = prGraphAddNode(&graph, PR_NODE_TYPE_READ);
	PrNodeId n5 = prGraphAddNode(&graph, PR_NODE_TYPE_READ);

	prGraphAddEdge(&graph, n1, n2);
	prGraphAddEdge(&graph, n1, n3);
	prGraphAddEdge(&graph, n1, n5);
	prGraphAddEdge(&graph, n2, n4);
	prGraphAddEdge(&graph, n3, n4);
	prGraphAddEdge(&graph, n3, n5);

	prGraphDump(&graph);
	
	prGraphRemoveNode(&graph, n3);

	printf("\n");
	prGraphDump(&graph);

	PrNodeId n6 = prGraphAddNode(&graph, PR_NODE_TYPE_READ);

	prGraphAddEdge(&graph, n4, n6);

	printf("\n");
	prGraphDump(&graph);

	prGraphRemoveNode(&graph, n5);

	printf("\n");
	prGraphDump(&graph);

	PrNodeId n7 = prGraphAddNode(&graph, PR_NODE_TYPE_READ);

	prGraphAddEdge(&graph, n1, n7);
	prGraphAddEdge(&graph, n2, n7);

	printf("\n");
	prGraphDump(&graph);

	PrNodeId n8 = prGraphAddNode(&graph, PR_NODE_TYPE_READ);

	prGraphAddEdge(&graph, n6, n8);
	prGraphAddEdge(&graph, n7, n8);

	printf("\n");
	prGraphDump(&graph);

	printf("\n=== Try adding cycle 8→1 (should be rejected) ===\n");
	prGraphAddEdge(&graph, n8, n1);

	printf("\n=== Try adding 8→1 again (still rejected) ===\n");
	prGraphAddEdge(&graph, n8, n1);

	printf("\n=== Try adding cycle 7→2 (should be rejected) ===\n\n");
	prGraphAddEdge(&graph, n7, n2);

	prGraphDump(&graph);

	printf("\n=== Topological sort ===\n");
	PrNodeId *sorted = prGraphTopologicalSort(&graph, &arena);
	if (sorted) {
		u64 n = wpArrayCount(sorted);
		printf("  count: %llu / %llu active\n",
		       (unsigned long long)n,
		       (unsigned long long)graph.count);
		for (u64 i = 0; i < n; i++) {
			printf("  [%llu] idx=%llu gen=%llu\n",
			       (unsigned long long)i,
			       (unsigned long long)sorted[i].index,
			       (unsigned long long)sorted[i].generation);
		}
	}

	wpMemArenaAllocatorDestroy(&arena);

	return 0;
}