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

#include "../src/wapp/wapp.h"
#include <inttypes.h>
#include <stdio.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 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;
	u64          source_idx;
	u64          target_idx;
};
/* ---------------------------------------------------------------------------
 * Graph vertex — compact adjacency head (two pointers + active flag)
 * -------------------------------------------------------------------------*/

typedef struct PrGraphVertex PrGraphVertex;
struct PrGraphVertex {
	PrGraphEdge *next_forward;
	PrGraphEdge *next_backward;
	b8           active;
};
typedef PrGraphVertex *PrGraphVertexArray;

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

typedef struct {
	PrPool             edge_pool;
	PrGraphVertexArray vertices;
	u64                capacity;
	u64                max_vertex_ever;
	u64                vertex_count;
} 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     prNodeManagerDumpGraph(const PrNodeManager *mgr);

wp_intern void       prGraphInit(PrGraph *graph, WpAllocator *allocator, u64 capacity);
wp_intern void       prGraphAddVertex(PrGraph *graph, u64 idx);
wp_intern void       prGraphRemoveVertex(PrGraph *graph, u64 idx);
wp_intern b8         prGraphAddEdge(PrGraph *graph, u64 from_idx, u64 to_idx);
wp_intern b8         prGraphEdgeExists(const PrGraph *graph, u64 from_idx, u64 to_idx);
wp_intern u64        prGraphVertexCount(const PrGraph *graph);
wp_intern WpU64Array prGraphTopologicalSort(const PrGraph *graph, const WpAllocator *allocator);

wp_intern void _unlinkForward(PrGraph *graph, u64 from_idx, PrGraphEdge *edge);
wp_intern void _unlinkBackward(PrGraph *graph, u64 to_idx, PrGraphEdge *edge);
/* ---------------------------------------------------------------------------
 * 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; }

	prGraphAddVertex(&mgr->graph, idx);

	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 and mark vertex inactive */
	prGraphRemoveVertex(&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; }
	if (prGraphEdgeExists(&mgr->graph, from.index, to.index)) { return; }
	prGraphAddEdge(&mgr->graph, from.index, to.index);
}

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

wp_intern void prGraphInit(PrGraph *graph, WpAllocator *allocator, u64 capacity) {
	graph->vertices      = wpArrayAllocCapacity(PrGraphVertex, allocator, capacity, WP_ARRAY_INIT_FILLED);
	graph->capacity      = capacity;
	graph->max_vertex_ever = 0;
	graph->vertex_count  = 0;

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

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

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

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

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

/* --- vertex lifecycle ------------------------------------------------- */

wp_intern void prGraphAddVertex(PrGraph *graph, u64 idx) {
	graph->vertices[idx].active = true;
	graph->vertex_count++;
	if (idx >= graph->max_vertex_ever) {
		graph->max_vertex_ever = idx + 1;
	}
}

wp_intern void prGraphRemoveVertex(PrGraph *graph, u64 idx) {
	PrGraphVertex *vtx = &graph->vertices[idx];
	if (!vtx->active) { return; }

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

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

	vtx->next_forward  = NULL;
	vtx->next_backward = NULL;
	vtx->active        = false;
	graph->vertex_count--;
}

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

wp_intern b8 prGraphEdgeExists(const PrGraph *graph, u64 from_idx, u64 to_idx) {
	PrGraphVertex *vtx  = &graph->vertices[from_idx];
	PrGraphEdge   *curr = vtx->next_forward;
	while (curr) {
		if (curr->target_idx == to_idx) { return true; }
		curr = curr->next_forward;
	}
	return false;
}

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

	edge->source_idx = from_idx;
	edge->target_idx = to_idx;

	/* Link into adjacency chains */
	PrGraphVertex *src    = &graph->vertices[from_idx];
	PrGraphVertex *dst    = &graph->vertices[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));
	WpU64Array  sorted   = prGraphTopologicalSort(graph, &scratch);
	u64         sorted_n = sorted ? wpArrayCount(sorted) : 0;
	if (sorted_n < graph->vertex_count) {
		_unlinkForward(graph, from_idx, edge);
		_unlinkBackward(graph, to_idx, edge);
		prPoolFree(&graph->edge_pool, edge);
		return false;
	}
	return true;
}

wp_intern u64 prGraphVertexCount(const PrGraph *graph) {
	return graph->vertex_count;
}

/* ---------------------------------------------------------------------------
 * 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 WpU64Array prGraphTopologicalSort(const PrGraph *graph, const WpAllocator *allocator) {
	if (graph->vertex_count == 0) { return NULL; }
	if (!graph->vertices || graph->capacity == 0) { return NULL; }

	WpU64Array result = wpArrayAllocCapacity(u64, allocator, graph->vertex_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_vertex_ever; i++) {
		if (!graph->vertices[i].active) { continue; }

		PrGraphEdge *curr = graph->vertices[i].next_forward;
		while (curr) {
			in_degree[curr->target_idx]++;
			curr = curr->next_forward;
		}
	}

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

	for (u64 i = 0; i < graph->max_vertex_ever; i++) {
		if (!graph->vertices[i].active) { continue; }
		if (in_degree[i] == 0) {
			wpQueuePush(u64, &queue, &i);
		}
	}

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

		wpArrayAppendCapped(u64, result, node_idx);

		PrGraphEdge *curr = graph->vertices[*node_idx].next_forward;
		while (curr) {
			u64 target_idx = curr->target_idx;
			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 prNodeManagerDumpGraph(const PrNodeManager *mgr) {
	const PrGraph *graph = &mgr->graph;
	printf("==============INPUTS==============\n");
	for (u64 i = 0; i < mgr->max_count_ever; ++i) {
		PrNodeId id = prNodeManagerGetNode(mgr, i);
		if (!prNodeManagerIsActiveNode(mgr, id)) { continue; }

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

		PrGraphEdge *curr = graph->vertices[i].next_backward;
		if (!curr) {
			printf(" (none)");
		} else {
			do {
				printf(" %" PRIu64, curr->source_idx + 1);
				curr = curr->next_backward;
			} while (curr);
		}
		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; }

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

		PrGraphEdge *curr = graph->vertices[i].next_forward;
		if (!curr) {
			printf(" (none)");
		} else {
			do {
				printf(" %" PRIu64, curr->target_idx + 1);
				curr = curr->next_forward;
			} while (curr);
		}
		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);

	prNodeManagerDumpGraph(&mgr);
	
	prNodeManagerRemoveNode(&mgr, n3);

	printf("\n");
	prNodeManagerDumpGraph(&mgr);

	PrNodeId n6 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);

	prNodeManagerAddEdge(&mgr, n4, n6);

	printf("\n");
	prNodeManagerDumpGraph(&mgr);

	prNodeManagerRemoveNode(&mgr, n5);

	printf("\n");
	prNodeManagerDumpGraph(&mgr);

	PrNodeId n7 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);

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

	printf("\n");
	prNodeManagerDumpGraph(&mgr);

	PrNodeId n8 = prNodeManagerAddNode(&mgr, PR_NODE_TYPE_READ);

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

	printf("\n");
	prNodeManagerDumpGraph(&mgr);

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

	prNodeManagerDumpGraph(&mgr);

	printf("\n=== Topological sort ===\n");
	WpU64Array sorted = prGraphTopologicalSort(&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++) {
			PrNodeId id = prNodeManagerGetNode(&mgr, sorted[i]);
			printf("  [%llu] idx=%llu gen=%llu\n",
			       (unsigned long long)i,
			       (unsigned long long)sorted[i],
			       (unsigned long long)id.generation);
		}
	}

	wpMemArenaAllocatorDestroy(&arena);

	return 0;
}