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

/* ---------------------------------------------------------------------------
 * Shared types
 * -------------------------------------------------------------------------*/

#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 type
 * -------------------------------------------------------------------------*/

typedef struct PrGraphEdge PrGraphEdge;
struct PrGraphEdge {
	PrGraphEdge *next_forward;
	PrGraphEdge *next_backward;
	PrNodeId     source;
	PrNodeId     target;
};
typedef PrGraphEdge *PrGraphEdgeArray;

/* ---------------------------------------------------------------------------
 * Graph node — compact adjacency head (two pointers, no dead fields)
 * -------------------------------------------------------------------------*/

typedef struct PrGraphNode PrGraphNode;
struct PrGraphNode {
	PrGraphEdge *next_forward;
	PrGraphEdge *next_backward;
};
typedef PrGraphNode *PrGraphNodeArray;

/* ---------------------------------------------------------------------------
 * PrGraph — owns only topology (edges + adjacency heads)
 * -------------------------------------------------------------------------*/

typedef struct {
	PrPool          edge_pool;
	PrGraphNodeArray graph_nodes;
	u64              capacity;
} PrGraph;

/* ---------------------------------------------------------------------------
 * PrNodeManager — owns compositor node data + handle lifecycle + topology
 * -------------------------------------------------------------------------*/

typedef struct {
	PrNodeArray   nodes;
	PrGraph       graph;
	u64           capacity;
	u64           max_count_ever;
	u64           count;
	u64           free_head;
} PrNodeManager;

/* ---------------------------------------------------------------------------
 * Function declarations (cross-referencing both types)
 * -------------------------------------------------------------------------*/

wp_intern void     prNodeManagerInit(PrNodeManager *mgr, WpAllocator *allocator, u64 capacity);
wp_intern b8       prNodeManagerIsStaleNode(const PrNodeManager *mgr, PrNodeId id);
wp_intern b8       prNodeManagerIsActiveNode(const PrNodeManager *mgr, PrNodeId id);
wp_intern PrNodeId prNodeManagerGetNode(const PrNodeManager *mgr, u64 index);
wp_intern PrNodeId prNodeManagerAddNode(PrNodeManager *mgr, PrNodeType type);
wp_intern void     prNodeManagerRemoveNode(PrNodeManager *mgr, PrNodeId id);
wp_intern void     prNodeManagerAddEdge(PrNodeManager *mgr, PrNodeId from, PrNodeId to);

wp_intern void          prGraphInit(PrGraph *graph, WpAllocator *allocator, u64 capacity);
wp_intern void          prGraphAddEdge(PrNodeManager *mgr, PrGraph *graph, u64 from_idx, u64 to_idx);
wp_intern void          prGraphRemoveEdges(PrGraph *graph, u64 idx);
wp_intern PrNodeIdArray prGraphTopologicalSort(const PrNodeManager *mgr, const PrGraph *graph,
                                               const WpAllocator *allocator);
wp_intern void          prGraphDump(const PrNodeManager *mgr, const PrGraph *graph);

/* ---------------------------------------------------------------------------
 * PrNodeManager implementation
 * -------------------------------------------------------------------------*/

wp_intern void prNodeManagerInit(PrNodeManager *mgr, WpAllocator *allocator, u64 capacity) {
	mgr->nodes          = wpArrayAllocCapacity(PrNode, allocator, capacity, WP_ARRAY_INIT_FILLED);
	mgr->free_head      = 0;
	mgr->capacity       = capacity;
	mgr->max_count_ever = 0;
	mgr->count          = 0;

	if (!mgr->nodes) {
		mgr->capacity = 0;
		return;
	}

	memset(mgr->nodes, 0, capacity * sizeof(PrNode));

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

	prGraphInit(&mgr->graph, allocator, capacity);
}

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

wp_intern b8 prNodeManagerIsActiveNode(const PrNodeManager *mgr, PrNodeId id) {
	u64 next_free = mgr->nodes[id.index].next_free;
	return !prNodeManagerIsStaleNode(mgr, id) && next_free == INVALID_NODE_INDEX;
}

wp_intern PrNodeId prNodeManagerGetNode(const PrNodeManager *mgr, u64 index) {
	return (PrNodeId){ .index = index, .generation = mgr->nodes[index].generation };
}

wp_intern PrNodeId prNodeManagerAddNode(PrNodeManager *mgr, PrNodeType type) {
	u64 idx = mgr->free_head;
	if (idx == INVALID_NODE_INDEX) { return INVALID_NODE_ID; }

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

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

	mgr->count++;

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

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

wp_intern void prNodeManagerRemoveNode(PrNodeManager *mgr, PrNodeId id) {
	if (!prNodeManagerIsActiveNode(mgr, id)) { return; }

	/* Tear down all edges incident to this node */
	prGraphRemoveEdges(&mgr->graph, id.index);

	/* Return node slot to free list with bumped generation */
	PrNode *node = &mgr->nodes[id.index];
	node->generation++;
	node->next_free = mgr->free_head;
	mgr->free_head  = id.index;
	mgr->count--;
}

wp_intern void prNodeManagerAddEdge(PrNodeManager *mgr, PrNodeId from, PrNodeId to) {
	if (!prNodeManagerIsActiveNode(mgr, from) || !prNodeManagerIsActiveNode(mgr, to)) { return; }
	if (from.index == to.index) { return; }
	prGraphAddEdge(mgr, &mgr->graph, from.index, to.index);
}

/* ---------------------------------------------------------------------------
 * PrGraph implementation
 * -------------------------------------------------------------------------*/

wp_intern void prGraphInit(PrGraph *graph, WpAllocator *allocator, u64 capacity) {
	graph->graph_nodes = wpArrayAllocCapacity(PrGraphNode, allocator, capacity, WP_ARRAY_INIT_FILLED);
	graph->capacity    = capacity;

	if (!graph->graph_nodes) {
		graph->capacity = 0;
		return;
	}

	prPoolInit(&graph->edge_pool, allocator, sizeof(PrGraphEdge));
	memset(graph->graph_nodes, 0, capacity * sizeof(PrGraphNode));
}

/* --- internal: unlink edge helpers (raw indices, caller guarantees validity) -- */

wp_intern void _unlinkForward(PrGraph *graph, u64 from_idx, PrGraphEdge *edge) {
	PrGraphNode *gnode = &graph->graph_nodes[from_idx];
	PrGraphEdge *curr  = gnode->next_forward;
	PrGraphEdge *prev  = NULL;
	while (curr) {
		if (curr == edge) {
			if (prev) {
				prev->next_forward = curr->next_forward;
			} else {
				gnode->next_forward = curr->next_forward;
			}
			return;
		}
		prev = curr;
		curr = curr->next_forward;
	}
}

wp_intern void _unlinkBackward(PrGraph *graph, u64 to_idx, PrGraphEdge *edge) {
	PrGraphNode *gnode = &graph->graph_nodes[to_idx];
	PrGraphEdge *curr  = gnode->next_backward;
	PrGraphEdge *prev  = NULL;
	while (curr) {
		if (curr == edge) {
			if (prev) {
				prev->next_backward = curr->next_backward;
			} else {
				gnode->next_backward = curr->next_backward;
			}
			return;
		}
		prev = curr;
		curr = curr->next_backward;
	}
}

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

wp_intern void prGraphAddEdge(PrNodeManager *mgr, PrGraph *graph, u64 from_idx, u64 to_idx) {
	PrGraphEdge *edge = (PrGraphEdge *)prPoolAlloc(&graph->edge_pool);
	if (!edge) { return; }

	PrNodeId from_id = prNodeManagerGetNode(mgr, from_idx);
	PrNodeId to_id   = prNodeManagerGetNode(mgr, to_idx);

	edge->source = from_id;
	edge->target = to_id;

	/* Link into adjacency chains */
	PrGraphNode *src   = &graph->graph_nodes[from_idx];
	PrGraphNode *dst   = &graph->graph_nodes[to_idx];
	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(mgr, graph, &scratch);
	u64          sorted_cnt = sorted ? wpArrayCount(sorted) : 0;
	if (sorted_cnt < mgr->count) {
		_unlinkForward(graph, from_idx, edge);
		_unlinkBackward(graph, to_idx, edge);
		prPoolFree(&graph->edge_pool, edge);
	}
}

wp_intern void prGraphRemoveEdges(PrGraph *graph, u64 idx) {
	PrGraphNode *gnode = &graph->graph_nodes[idx];

	/* Free outgoing edges: unlink from each target's backward list */
	PrGraphEdge *curr = gnode->next_forward;
	while (curr) {
		PrGraphEdge *next = curr->next_forward;
		_unlinkBackward(graph, curr->target.index, curr);
		prPoolFree(&graph->edge_pool, curr);
		curr = next;
	}

	/* Free incoming edges: unlink from each source's forward list */
	curr = gnode->next_backward;
	while (curr) {
		PrGraphEdge *next = curr->next_backward;
		_unlinkForward(graph, curr->source.index, curr);
		prPoolFree(&graph->edge_pool, curr);
		curr = next;
	}

	gnode->next_forward = gnode->next_backward = NULL;
}

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

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

	PrNodeIdArray result = wpArrayAllocCapacity(PrNodeId, allocator, mgr->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 < mgr->max_count_ever; i++) {
		PrNodeId id = prNodeManagerGetNode(mgr, i);
		if (!prNodeManagerIsActiveNode(mgr, id)) { continue; }

		PrGraphNode *gnode = &graph->graph_nodes[i];
		PrGraphEdge *curr  = gnode->next_forward;
		while (curr) {
			in_degree[curr->target.index]++;
			curr = curr->next_forward;
		}
	}

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

	for (u64 i = 0; i < mgr->max_count_ever; i++) {
		PrNodeId id = prNodeManagerGetNode(mgr, i);
		if (!prNodeManagerIsActiveNode(mgr, 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 = prNodeManagerGetNode(mgr, *node_idx);
		wpArrayAppendCapped(PrNodeId, result, &id);

		PrGraphNode *gnode = &graph->graph_nodes[*node_idx];
		PrGraphEdge *curr  = gnode->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;
}

/* ---------------------------------------------------------------------------
 * Dump
 * -------------------------------------------------------------------------*/

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

		PrGraphNode *gnode = &graph->graph_nodes[i];

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

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

		printf("\n");
	}

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

		PrGraphNode *gnode = &graph->graph_nodes[i];

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

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

		printf("\n");
	}
}

/* ---------------------------------------------------------------------------
 * Main
 * -------------------------------------------------------------------------*/

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

	PrNodeManager mgr = {0};
	prNodeManagerInit(&mgr, &arena, 128);

	PrNodeId n1 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);
	PrNodeId n2 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);
	PrNodeId n3 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);
	PrNodeId n4 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);
	PrNodeId n5 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);

	prNodeManagerAddEdge(&mgr, n1, n2);
	prNodeManagerAddEdge(&mgr, n1, n3);
	prNodeManagerAddEdge(&mgr, n1, n5);
	prNodeManagerAddEdge(&mgr, n2, n4);
	prNodeManagerAddEdge(&mgr, n3, n4);
	prNodeManagerAddEdge(&mgr, n3, n5);

	prGraphDump(&mgr, &mgr.graph);
	
	prNodeManagerRemoveNode(&mgr, n3);

	printf("\n");
	prGraphDump(&mgr, &mgr.graph);

	PrNodeId n6 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);

	prNodeManagerAddEdge(&mgr, n4, n6);

	printf("\n");
	prGraphDump(&mgr, &mgr.graph);

	prNodeManagerRemoveNode(&mgr, n5);

	printf("\n");
	prGraphDump(&mgr, &mgr.graph);

	PrNodeId n7 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);

	prNodeManagerAddEdge(&mgr, n1, n7);
	prNodeManagerAddEdge(&mgr, n2, n7);

	printf("\n");
	prGraphDump(&mgr, &mgr.graph);

	PrNodeId n8 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);

	prNodeManagerAddEdge(&mgr, n6, n8);
	prNodeManagerAddEdge(&mgr, n7, n8);

	printf("\n");
	prGraphDump(&mgr, &mgr.graph);

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

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

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

	prGraphDump(&mgr, &mgr.graph);

	printf("\n=== Topological sort ===\n");
	PrNodeId *sorted = prGraphTopologicalSort(&mgr, &mgr.graph, &arena);
	if (sorted) {
		u64 n = wpArrayCount(sorted);
		printf("  count: %llu / %llu active\n",
		       (unsigned long long)n,
		       (unsigned long long)mgr.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;
}