/*
 * 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 <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <assert.h>
#include <lauxlib.h>
#include "melisma.h"
#include "mutil.h"

#define EXPR_TYPENAME "expression"
#define SCRIPT_TYPENAME "expression script"

enum tag {ADD,SUB,MUL,DIV,IN,OUT,CONST,HALT};

struct op
{
	enum tag    tag;
	float       value;
	size_t      index;
	size_t      offset;
	ml_channel *src;
};

struct script
{
	size_t    len;
	struct op ops[];
};

struct expr
{
	size_t           num_inputs;
	size_t          *buf_lens;
	size_t          *cursors;
	float          **bufs;
	size_t           out_buf_len;
	size_t           out_cursor;
	float           *out_buf;
	float           *stack;
	size_t           stack_len;
	size_t           stack_p;
	struct script   *script;
	ml_channel      *out;
};

static size_t script_alloc_size(size_t len)
{
	return sizeof(struct script) + len*sizeof(struct op);
}

static struct script *push_script(lua_State *s, size_t len)
{
	struct script *script = lua_newuserdata(s,script_alloc_size(len));
	luaL_setmetatable(s,SCRIPT_TYPENAME);
	script->len = len;
	return script;
}

static int l_const(lua_State *s)
{
	float value = luaL_checknumber(s,1);
	struct script *script = push_script(s,1);
	script->ops[0].tag = CONST;
	script->ops[0].value = value;
	return 1;
}

static int l_input(lua_State *s)
{
	ml_channel *src = mutil_get_channel(s,1);
	size_t offset = luaL_optinteger(s,2,0);
	struct script *script = push_script(s,1);
	script->ops[0].tag = IN;
	script->ops[0].src = src;
	script->ops[0].offset = offset;
	return 1;
}

static int l_output(lua_State *s)
{
	size_t offset = luaL_checkinteger(s,1);
	struct script *script = push_script(s,1);
	script->ops[0].tag = OUT;
	script->ops[0].offset = offset;
	return 1;
}

static int convert(lua_State *s, int i)
{
	if (luaL_testudata(s,i,SCRIPT_TYPENAME))
	{
		return i;
	}

	if (lua_isnumber(s,i))
	{
		lua_pushcfunction(s,l_const);
		lua_pushvalue(s,i);
		lua_call(s,1,1);
		return lua_gettop(s);
	}

	const char *typename = lua_typename(s,lua_type(s,i));
	luaL_error(s,"expected number or expression script, got %s",typename);

	return 0;
}

static int binop(lua_State *s, int tag)
{
	int lhs_i = convert(s,1);
	int rhs_i = convert(s,2);

	struct script *lhs = luaL_checkudata(s,lhs_i,SCRIPT_TYPENAME);
	struct script *rhs = luaL_checkudata(s,rhs_i,SCRIPT_TYPENAME);

	size_t len = lhs->len + rhs->len + 1;
	struct script *res = push_script(s,len);

	memcpy(&res->ops[0],lhs->ops,sizeof(res->ops[0])*lhs->len);
	memcpy(&res->ops[lhs->len],rhs->ops,sizeof(res->ops[0])*rhs->len);

	res->ops[len-1].tag = tag;

	return 1;
}

static int l_add(lua_State *s)
{
	return binop(s,ADD);
}

static int l_sub(lua_State *s)
{
	return binop(s,SUB);
}

static int l_mul(lua_State *s)
{
	return binop(s,MUL);
}

static int l_div(lua_State *s)
{
	return binop(s,DIV);
}

static void prepare_inputs(ml_node *node, struct script *script, size_t **buf_lens)
{
	for (size_t i = 0; i < script->len; i++)
	{
		struct op *op = &script->ops[i];

		if (op->tag != IN)
		{
			continue;
		}

		/* the same input channel may have been encountered already */
		bool have_prev = false;
		for (size_t j = 0; j < i; j++)
		{
			struct op *prev = &script->ops[j];

			if (prev->tag != IN)
			{
				continue;
			}

			if (op->src == prev->src)
			{
				op->index = prev->index;
				have_prev = true;

				if (op->offset+1 > (*buf_lens)[op->index])
				{
					/* grow lookback buffer length */
					(*buf_lens)[op->index] = op->offset+1;
				}

				break;
			}
		}

		if (have_prev)
		{
			continue;
		}

		op->index = ml_num_inputs(node);
		ml_channel *dst = ml_add_input(node);
		ml_connect(op->src,dst);

		/* add buffer length */
		*buf_lens = realloc(*buf_lens,sizeof(**buf_lens)*ml_num_inputs(node));
		(*buf_lens)[op->index] = op->offset+1;
	}
}

static size_t max_output_offset(struct script *script)
{
	size_t max = 0;

	for (size_t i = 0; i < script->len; i++)
	{
		struct op *op = &script->ops[i];

		if (op->tag != OUT)
		{
			continue;
		}

		if (op->offset+1 > max)
		{
			max = op->offset+1;
		}
	}

	return max;
}

static void proc(const ml_proc_args *args);

static int l_build(lua_State *s)
{
	struct script *script = luaL_checkudata(s,1,SCRIPT_TYPENAME);

	/* copy script and add a halt instruction */
	struct script *copy = calloc(1,script_alloc_size(script->len+1));
	copy->len = script->len + 1;
	memcpy(copy->ops,script->ops,script->len*sizeof(copy->ops[0]));
	copy->ops[script->len].tag = HALT;

	struct expr *expr = calloc(1,sizeof(*expr));
	expr->script = copy;

	/* add output channel */
	ml_node *node = ml_create_node(proc,expr);
	ml_label_node(node,"expression");
	expr->out = ml_add_output(node);

	/* set input ops indices and determine lookback buffer lengths */
	prepare_inputs(node,copy,&expr->buf_lens);
	expr->num_inputs = ml_num_inputs(node);

	/* allocate lookback buffers for inputs */
	expr->bufs = calloc(1,sizeof(*expr->bufs)*expr->num_inputs);
	expr->cursors = calloc(1,sizeof(*expr->cursors)*expr->num_inputs);
	for (size_t i = 0; i < expr->num_inputs; i++)
	{
		expr->bufs[i] = calloc(1,sizeof(*expr->bufs[i])*expr->buf_lens[i]);
	}

	/* allocate lookback buffer for output */
	expr->out_buf_len = max_output_offset(copy);
	expr->out_buf = calloc(1,sizeof(*expr->out_buf)*expr->out_buf_len);

	*(void**) lua_newuserdata(s,sizeof(void*)) = expr;
	luaL_setmetatable(s,EXPR_TYPENAME);

	return 1;
}

static int l_out(lua_State *s)
{
	struct expr *e = *(void**) luaL_checkudata(s,1,EXPR_TYPENAME);
	*(void**) lua_newuserdata(s,sizeof(void*)) = e->out;
	luaL_setmetatable(s,MELISMA_OUTPUT_CHANNEL_TYPENAME);
	return 1;
}

static void push(struct expr *expr, float v)
{
	if (expr->stack_p == expr->stack_len)
	{
		expr->stack_len = expr->stack_len == 0 ? 12 : 2*expr->stack_len;
		expr->stack = realloc(expr->stack,sizeof(expr->stack[0])*expr->stack_len);
	}

	expr->stack[expr->stack_p++] = v;
}

static float pop(struct expr *expr)
{
	assert(expr->stack_p > 0);
	return expr->stack[--expr->stack_p];
}

static void proc(const ml_proc_args *args)
{
	struct expr *e = args->userp;

	static const void *tbl[] = {
		[ADD] = &&op_add,
		[SUB] = &&op_sub,
		[MUL] = &&op_mul,
		[DIV] = &&op_div,
		[IN] = &&op_in,
		[OUT] = &&op_out,
		[CONST] = &&op_const,
		[HALT] = &&op_halt,
	};

	float *out = args->out[0];

	for (size_t i = 0; i < args->nframes; i++)
	{
		/* update lookback buffers */
		for (size_t j = 0; j < e->num_inputs; j++)
		{
			e->bufs[j][e->cursors[j]] = args->in[j][i];
		}

		/* execute script */
		size_t pc = 0;
		struct op *op;
		float rhs,lhs,res;
		size_t l;

		#define NEXT()\
			op = &e->script->ops[pc++];\
			goto *tbl[op->tag];

		NEXT();

		op_add:
			rhs = pop(e);
			lhs = pop(e);
			push(e,lhs+rhs);
			NEXT();

		op_sub:
			rhs = pop(e);
			lhs = pop(e);
			push(e,lhs-rhs);
			NEXT();

		op_mul:
			rhs = pop(e);
			lhs = pop(e);
			push(e,lhs*rhs);
			NEXT();

		op_div:
			rhs = pop(e);
			lhs = pop(e);
			push(e,lhs/rhs);
			NEXT();

		op_in:
			l = e->cursors[op->index] + e->buf_lens[op->index] - op->offset;
			while (l >= e->buf_lens[op->index])
			{
				l -= e->buf_lens[op->index];
			}
			push(e,e->bufs[op->index][l]);
			NEXT();

		op_out:
			l = e->out_cursor + e->out_buf_len - op->offset;
			while (l >= e->out_buf_len)
			{
				l -= e->out_buf_len;
			}
			push(e,e->out_buf[l]);
			NEXT();

		op_const:
			push(e,op->value);
			NEXT();

		op_halt:
			assert(e->stack_p == 1);
			out[i] = pop(e);

		/* update output lookback buffer */
		if (e->out_buf_len > 0)
		{
			e->out_buf[e->out_cursor++] = out[i];
			if (e->out_cursor == e->out_buf_len)
			{
				e->out_cursor = 0;
			}
		}

		/* update input looback buffer cursors */
		for (size_t j = 0; j < e->num_inputs; j++)
		{
			e->cursors[j]++;
			if (e->cursors[j] == e->buf_lens[j])
			{
				e->cursors[j] = 0;
			}
		}
	}
}

int luaopen_melisma_expression(lua_State *s)
{
/*** Build filters with arithmetic expressions.

	local mexpression = require 'melisma.expression'
	local input = mepxression.input

	local lhs = ...
	local rhs = ...

	local script = input(lhs,0) + 0.5 * input(rhs,10)
	local expression = script:build()

@module expression
*/
	lua_newtable(s);
	int tbl = lua_gettop(s);

	lua_pushcfunction(s,l_const);
	lua_setfield(s,tbl,"const");

/*** Introduce an input signal in the expression.
@static
@function input
@tparam channel sig Input signal.
@tparam integer off Offset in no. of samples.
@treturn script
*/
	lua_pushcfunction(s,l_input);
	lua_setfield(s,tbl,"input");


/*** Introduce the output signal in the expression.
@static
@function output
@tparam integer off Offset in no. of samples.
@treturn script
*/
	lua_pushcfunction(s,l_output);
	lua_setfield(s,tbl,"output");

/*** Represents an expression before compilation.
Implements `__add`, `__sub`, `__mul` and `__div` for
combining scripts in to larger ones.
@type script
*/
/*** Compile an expression script.
@function script:build
@treturn expression
*/
	lua_pushcfunction(s,l_build);
	lua_setfield(s,tbl,"build");

	luaL_newmetatable(s,SCRIPT_TYPENAME);
	int script_mt = lua_gettop(s);

	lua_pushvalue(s,tbl);
	lua_setfield(s,script_mt,"__index");

	lua_pushcfunction(s,l_add);
	lua_setfield(s,script_mt,"__add");

	lua_pushcfunction(s,l_sub);
	lua_setfield(s,script_mt,"__sub");

	lua_pushcfunction(s,l_mul);
	lua_setfield(s,script_mt,"__mul");

	lua_pushcfunction(s,l_div);
	lua_setfield(s,script_mt,"__div");

	lua_newtable(s);
	int expr_tbl = lua_gettop(s);

	luaL_newmetatable(s,EXPR_TYPENAME);
	int expr_mt = lua_gettop(s);

	lua_pushvalue(s,expr_tbl);
	lua_setfield(s,expr_mt,"__index");

	lua_pushcfunction(s,l_out);
	lua_setfield(s,expr_tbl,MUTIL_OUTPUT_CHANNEL_FUNCTION);

	lua_settop(s,tbl);
	return 1;
}