/*
 * melisma - extensible and scriptable audio engine
 * Copyright (C) 2026  Ludvig Gunne Lindström
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <stdarg.h>
#include <string.h>
#include <stdbool.h>

#include <lua.h>
#include <lauxlib.h>

#include "melisma.h"
#include "graph.h"
#include "mem.h"

#ifndef BUFFER_SIZE
#define BUFFER_SIZE 4096
#endif

struct ml_channel
{
	ml_node        *node;
	enum {IN,OUT}   dir;
	ml_channel     *src;
	ml_channel     *next;
	size_t          buf;
	int             refs;
};

struct ml_node
{
	ml_proc_fn   proc;
	void        *userp;
	ml_channel  *chans[2];
	ml_node     *next;
	bool         visited;
	bool         cycle;
	char        *label;
};

struct chunk
{
	char          *label;
	ml_proc_fn     proc;
	ml_proc_args   args;
	struct chunk  *next;
	float         *data[];
};

struct free_list
{
	size_t            buf;
	struct free_list *next;
};

static ml_node          *root;
static ml_node          *active;
static ml_node          *active_tail;
static struct free_list *free_list;
static size_t            counter;
static float            *buffers;
static struct chunk     *graph;

ml_node *ml_create_node(ml_proc_fn proc, void *userp)
{
	ml_node *node = scalloc(1,sizeof(*node));
	node->proc = proc;
	node->userp = userp;
	return node;
}

void ml_label_node(ml_node *node, const char *fmt, ...)
{
	va_list ap_alloc,ap_print;

	va_start(ap_alloc,fmt);
	va_copy(ap_print,ap_alloc);

	size_t len = 1 + vsnprintf(NULL,0,fmt,ap_alloc);
	va_end(ap_alloc);

	char *buf = smalloc(len);
	vsnprintf(buf,len,fmt,ap_print);

	node->label = buf;
}

void *ml_get_userp(ml_node *node)
{
	return node->userp;
}

ml_channel *add_channel(ml_node *node, int dir)
{
	ml_channel *chan = scalloc(1,sizeof(*chan));
	chan->dir = dir;
	chan->node = node;

	if (node->chans[dir] == NULL)
	{
		node->chans[dir] = chan;
	}
	else
	{
		ml_channel *insert = node->chans[dir];
		while (insert->next)
		{
			insert = insert->next;
		}
		insert->next = chan;
	}

	return chan;
}

ml_channel *ml_add_input(ml_node *node)
{
	return add_channel(node,IN);
}

ml_channel *ml_add_output(ml_node *node)
{
	return add_channel(node,OUT);
}

static ml_channel *get_channel(ml_node *node, size_t i, int dir)
{
	ml_channel *chan = node->chans[dir];

	while (i > 0)
	{
		if (chan == NULL)
		{
			break;
		}

		i--;
		chan = chan->next;
	}

	return chan;
}

ml_channel *ml_get_input(ml_node *node, size_t i)
{
	return get_channel(node,i,IN);
}

ml_channel *ml_get_output(ml_node *node, size_t i)
{
	return get_channel(node,i,OUT);
}

static size_t num_channels(ml_node *node, int dir)
{
	ml_channel *chan = node->chans[dir];
	size_t n = 0;

	while (chan != NULL)
	{
		n++;
		chan = chan->next;
	}

	return n;
}

size_t ml_num_inputs(ml_node *node)
{
	return num_channels(node,IN);
}

size_t ml_num_outputs(ml_node *node)
{
	return num_channels(node,OUT);
}

void ml_connect(ml_channel *src, ml_channel *dst)
{
	assert(src->dir == OUT);
	assert(dst->dir == IN);
	src->refs++;
	dst->src = src;
}

ml_node *ml_get_root_node(void)
{
	return root;
}

void init_graph(lua_State *s)
{
	root = ml_create_node(NULL,NULL);
	ml_label_node(root,"root");

	luaL_newmetatable(s,MELISMA_NODE_TYPENAME);
	luaL_newmetatable(s,MELISMA_INPUT_CHANNEL_TYPENAME);
	luaL_newmetatable(s,MELISMA_OUTPUT_CHANNEL_TYPENAME);
}

static const char *topo_sort(ml_node *node)
{
	if (node->visited)
	{
		return NULL;
	}

	if (node->cycle)
	{
		return "cycle in audio graph detected";
	}

	node->visited = true;
	node->cycle = true;

	ml_channel *in = node->chans[IN];

	while (in != NULL)
	{
		if (in->src == NULL)
		{
			return "disconnected input channel in audio graph detected";
		}

		const char *err = topo_sort(in->src->node);

		if (err != NULL)
		{
			return err;
		}

		in = in->next;
	}

	if (active_tail)
	{
		active_tail->next = node;
	}
	else
	{
		active = node;
	}

	active_tail = node;
	node->cycle = false;

	return NULL;
}

static size_t alloc_buf(void)
{
	if (free_list)
	{
		struct free_list *head = free_list;
		size_t buf = head->buf;
		free_list = head->next;
		free(head);
		return buf;
	}

	return counter++;
}

static void free_buf(size_t buf)
{
	struct free_list *head = scalloc(1,sizeof(*head));
	head->buf = buf;
	head->next = free_list;
	free_list = head;
}

static void assign_bufs()
{
	ml_node *node = active;

	while (node)
	{
		ml_channel *in, *out;
		in = node->chans[IN];
		out = node->chans[OUT];

		while (out)
		{
			out->buf = alloc_buf();
			out = out->next;
		}

		while (in)
		{
			if (--in->src->refs == 0)
			{
				free_buf(in->src->buf);
			}
			in = in->next;
		}

		node = node->next;
	}
}

static void make_chunks(void)
{
	struct chunk *chunk, *tail = NULL;
	struct ml_node *node = active;

	while (node)
	{
		size_t icnt = ml_num_inputs(node);
		size_t ocnt = ml_num_outputs(node);
		size_t nchans = icnt + ocnt;
		size_t size = sizeof(*chunk) + sizeof(chunk->data[0]) * nchans;

		chunk = scalloc(1,size);
		chunk->label = node->label ? strdup(node->label) : NULL;
		chunk->proc = node->proc;
		chunk->args.userp = node->userp;
		chunk->args.icnt = icnt;
		chunk->args.ocnt = ocnt;
		chunk->args.in = chunk->data;
		chunk->args.out = chunk->data + icnt;

		ml_channel *in, *out;
		size_t i = 0;

		in = node->chans[IN];

		while (in != NULL)
		{
			chunk->data[i++] = &buffers[BUFFER_SIZE*in->src->buf];
			in = in->next;
		}

		out = node->chans[OUT];

		while (out != NULL)
		{
			chunk->data[i++] = &buffers[BUFFER_SIZE*out->buf];
			out = out->next;
		}

		assert(i == nchans);

		if (tail == NULL)
		{
			graph = chunk;
		}
		else
		{
			tail->next = chunk;
		}

		tail = chunk;
		node = node->next;
	}
}

const char *build_graph(void)
{
	const char *err = topo_sort(root);

	if (err != NULL)
	{
		return err;
	}

	assign_bufs();

	buffers = malloc(sizeof(*buffers)*BUFFER_SIZE*counter);

	make_chunks();

	return NULL;
}

static void proc_once(size_t n)
{
	struct chunk *c = graph;
	while (c)
	{
		if (c->proc)
		{
			c->args.nframes = n;
			c->proc(&c->args);
		}

		c = c->next;
	}
}

void proc_graph(size_t n)
{
	while (n > 0)
	{
		size_t m = n > BUFFER_SIZE ? BUFFER_SIZE : n;
		proc_once(m);
		n -= m;
	}
}

static void add_to_layout(lua_State *s, int layout, ml_node *node)
{
	int top = lua_absindex(s,lua_gettop(s));

	lua_pushlightuserdata(s,node);
	lua_gettable(s,layout);

	if (!lua_isnil(s,-1))
	{
		return;
	}

	lua_newtable(s);
	int tbl = lua_absindex(s,lua_gettop(s));

	if (node->label)
	{
		lua_pushstring(s,node->label);
		lua_setfield(s,tbl,"label");
	}

	lua_newtable(s);
	int list = lua_absindex(s,lua_gettop(s));
	ml_channel *in = node->chans[IN];
	int i = 1;

	while (in)
	{
		add_to_layout(s,layout,in->src->node);
		lua_pushlightuserdata(s,in->src->node);
		lua_rawseti(s,list,i++);
		in = in->next;
	}

	lua_pushvalue(s,list);
	lua_setfield(s,tbl,"inputs");

	lua_pushlightuserdata(s,node);
	lua_pushvalue(s,tbl);
	lua_settable(s,layout);

	lua_settop(s,top);
}

int l_graph_layout(lua_State *s)
{
	lua_newtable(s);
	int layout = lua_absindex(s,lua_gettop(s));
	add_to_layout(s,layout,root);
	lua_pushvalue(s,layout);
	return 1;
}