/*
 * 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 <math.h>

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

#include "melisma.h"
#include "mutil.h"

#define WAVETABLE_TYPENAME "wavetable"
#define OSCIL_TYPENAME "oscil"

#ifndef WAVETABLE_SIZE
#define WAVETABLE_SIZE 2048
#endif

#define WAVETABLE_ALLOC_SIZE (sizeof(float)*WAVETABLE_SIZE)

struct oscil
{
	ml_node       *node;
	unsigned long  clk;
	float          f;
	float          a;
	int            mf;
	int            ma;
	float          p;
	float          v;
	float          tbl[WAVETABLE_SIZE];
};

static float *push_new_wavetable(lua_State *s)
{
	float *ptr = lua_newuserdata(s,WAVETABLE_ALLOC_SIZE);
	luaL_setmetatable(s,WAVETABLE_TYPENAME);
	memset(ptr,0,WAVETABLE_ALLOC_SIZE);
	return ptr;
}

static int l_const(lua_State *s)
{
	float v = luaL_checknumber(s,1);
	float *tbl = push_new_wavetable(s);

	for (size_t i = 0; i < WAVETABLE_SIZE; i++)
	{
		tbl[i] = v;
	}

	return 1;
}

static int l_sin(lua_State *s)
{
	float *tbl = push_new_wavetable(s);

	for (size_t i = 0; i < WAVETABLE_SIZE; i++)
	{
		float f = 2.f * M_PI * (float) i / (float) WAVETABLE_SIZE;
		tbl[i] = sinf(f);
	}

	return 1;
}

static int l_saw(lua_State *s)
{
	float *tbl = push_new_wavetable(s);

	for (size_t i = 0; i < WAVETABLE_SIZE; i++)
	{
		tbl[i] =  2 * (float) i / (float) WAVETABLE_SIZE - 1;
	}

	return 1;
}

static int l_square(lua_State *s)
{
	float *tbl = push_new_wavetable(s);

	for (size_t i = 0; i < WAVETABLE_SIZE; i++)
	{
		if (i >= WAVETABLE_SIZE/2)
		{
			tbl[i] = 1;
		}
		else
		{
			tbl[i] = -1;
		}
	}

	return 1;
}

static int l_shift(lua_State *s)
{
	float *src = luaL_checkudata(s,1,WAVETABLE_TYPENAME);
	float p = luaL_checknumber(s,2);

	while (p < 0)
	{
		p += 2*M_PI;
	}

	while (p > 2*M_PI)
	{
		p -= 2*M_PI;
	}

	size_t offset = p/2.f/M_PI*WAVETABLE_SIZE;
	float *dst = push_new_wavetable(s);

	size_t i,d;
	for (i=0, d=offset; i < WAVETABLE_SIZE; i++, d++)
	{
		dst[d % WAVETABLE_SIZE] = src[i];
	}

	return 1;
}

static int l_scale(lua_State *s)
{
	float *src = luaL_checkudata(s,1,WAVETABLE_TYPENAME);
	int n = luaL_checkinteger(s,2);

	// if (n <= 0 || (n & (n-1)) != 0)
	// {
	// 	luaL_error(s,"scaling factor must be a power of two");
	// }

	float *dst = push_new_wavetable(s);

	for (size_t i = 0; i < WAVETABLE_SIZE; i++)
	{
		size_t j = (n*i) % WAVETABLE_SIZE;
		dst[i] = src[j];
	}

	return 1;
}

static int l_add(lua_State *s)
{
	float *lhs = luaL_checkudata(s,1,WAVETABLE_TYPENAME);
	float *rhs = luaL_checkudata(s,2,WAVETABLE_TYPENAME);
	float *res = push_new_wavetable(s);

	for (size_t i = 0; i < WAVETABLE_SIZE; i++)
	{
		res[i] = lhs[i] + rhs[i];
	}

	return 1;
}

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

	unsigned long sr = ml_sample_rate();
	float p = oscil->p;
	float *dst = args->out[0];

	float *mf = oscil->mf < 0 ? NULL : args->in[oscil->mf];
	float *ma = oscil->ma < 0 ? NULL : args->in[oscil->ma];

	for (size_t i = 0; i < args->nframes; i++)
	{
		float f = mf ? mf[i] : oscil->f;
		float a = ma ? ma[i] : oscil->a;

		float t = p*WAVETABLE_SIZE;

		size_t i0 = (size_t) t;
		size_t i1 = i0 + 1;

		float ip = t - i0;

		i0 &= (WAVETABLE_SIZE-1);
		i1 &= (WAVETABLE_SIZE-1);

		float s0 = oscil->tbl[i0];
		float s1 = oscil->tbl[i1];

		dst[i] = a*(s0 + ip*(s1 - s0));

		p += f/sr;

		if (p > 1)
		{
			p -= 1;
		}
	}

	oscil->p = p;
	oscil->v = dst[args->nframes-1];
}

static int l_init(lua_State *s)
{
	float *tbl = luaL_checkudata(s,1,WAVETABLE_TYPENAME);
	float f = luaL_checknumber(s,2);
	float a = luaL_checknumber(s,3);

	struct oscil *oscil = calloc(1,sizeof(*oscil));
	oscil->clk = 0;
	oscil->f = f;
	oscil->a = a;
	oscil->mf = -1;
	oscil->ma = -1;
	oscil->p = 0;
	oscil->v = a*tbl[0];

	ml_node *node = ml_create_node(proc,oscil);
	ml_label_node(node,"oscil");
	ml_add_output(node);
	oscil->node = node;

	memcpy(oscil->tbl,tbl,WAVETABLE_ALLOC_SIZE);

	*(void**) lua_newuserdata(s,sizeof(void*)) = oscil;
	luaL_setmetatable(s,OSCIL_TYPENAME);

	return 1;
}

static int l_setf(lua_State *s)
{
	struct oscil *oscil = *(void**) luaL_checkudata(s,1,OSCIL_TYPENAME);
	float v = luaL_checknumber(s,2);
	oscil->f = v;
	return 0;
}

static int l_seta(lua_State *s)
{
	struct oscil *oscil = *(void**) luaL_checkudata(s,1,OSCIL_TYPENAME);
	float v = luaL_checknumber(s,2);
	oscil->a = v;
	return 0;
}

static int l_mapf(lua_State *s)
{
	struct oscil *oscil = *(void**) luaL_checkudata(s,1,OSCIL_TYPENAME);
	ml_channel *src = mutil_get_channel(s,2);

	oscil->mf = ml_num_inputs(oscil->node);

	ml_channel *dst = ml_add_input(oscil->node);
	ml_connect(src,dst);

	return 0;
}

static int l_mapa(lua_State *s)
{
	struct oscil *oscil = *(void**) luaL_checkudata(s,1,OSCIL_TYPENAME);
	ml_channel *src = mutil_get_channel(s,2);

	oscil->ma = ml_num_inputs(oscil->node);

	ml_channel *dst = ml_add_input(oscil->node);
	ml_connect(src,dst);

	return 0;
}

static int l_sample(lua_State *s)
{
	struct oscil *oscil = *(void**) luaL_checkudata(s,1,OSCIL_TYPENAME);
	lua_pushnumber(s,oscil->v);
	return 1;
}

static int l_out(lua_State *s)
{
	struct oscil *oscil = *(void**) luaL_checkudata(s,1,OSCIL_TYPENAME);
	*(void**) lua_newuserdata(s,sizeof(void*)) = ml_get_output(oscil->node,0);
	luaL_setmetatable(s,MELISMA_OUTPUT_CHANNEL_TYPENAME);
	return 1;
}

int luaopen_melisma_oscil(lua_State *s)
{

/*** Wavetable backed oscillators.

	local moscil = require 'melisma.oscil'
	local mexpression = require 'melisma.expression'

	local wt = moscil.wavetable.sin()

	-- Make an LFO for controlling the frequency.
	local lfo_oscil = moscil.init(wt,1,1)

	-- Map [-1, 1] to [440, 880]
	local lfo = (440 + (1 + mexpression.input(lfo_oscil,0) / 2) * 440):build()

	local oscil = moscil.init(wt,0,0.5)
	oscil:mapf(lfo)

@classmod oscil
*/
	lua_newtable(s);
	int tbl = lua_absindex(s,lua_gettop(s));

/*** Initialize an oscillator.
@static
@function init
@tparam wavetable wt
@tparam number freq
@tparam number amp
@treturn oscil
*/
	lua_pushcfunction(s,l_init);
	lua_setfield(s,tbl,"init");

/*** Set oscillator frequency.
@function setf
@tparam number freq
*/
	lua_pushcfunction(s,l_setf);
	lua_setfield(s,tbl,"setf");

/*** Set oscillator amplitude.
@function seta
@tparam number amp
*/
	lua_pushcfunction(s,l_seta);
	lua_setfield(s,tbl,"seta");

/*** Map a signal to the frequency.
@function mapf
@tparam channel sig
*/
	lua_pushcfunction(s,l_mapf);
	lua_setfield(s,tbl,"mapf");

/*** Map a signal to the amplitude.
@function mapa
@tparam channel sig
*/
	lua_pushcfunction(s,l_mapa);
	lua_setfield(s,tbl,"mapa");

/*** Get the current oscillator value.
@function sample
@treturn number
*/
	lua_pushcfunction(s,l_sample);
	lua_setfield(s,tbl,"sample");

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

	luaL_newmetatable(s,OSCIL_TYPENAME);
	int oscil_mt = lua_gettop(s);

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

/***
@type wavetable
*/
	lua_newtable(s);
	int wt_tbl = lua_absindex(s,lua_gettop(s));

	lua_pushvalue(s,wt_tbl);
	lua_setfield(s,tbl,"wavetable");

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

/*** Sine wave.
@static
@function wavetable.sin
@treturn wavetable
*/
	lua_pushcfunction(s,l_sin);
	lua_setfield(s,wt_tbl,"sin");

/*** Saw wave.
@static
@function wavetable.saw
@treturn wavetable
*/
	lua_pushcfunction(s,l_saw);
	lua_setfield(s,wt_tbl,"saw");

/*** Square wave.
@static
@function wavetable.square
@treturn wavetable
*/
	lua_pushcfunction(s,l_square);
	lua_setfield(s,wt_tbl,"square");

	lua_pushcfunction(s,l_shift);
	lua_setfield(s,wt_tbl,"shift");

	lua_pushcfunction(s,l_scale);
	lua_setfield(s,wt_tbl,"scale");

	luaL_newmetatable(s,WAVETABLE_TYPENAME);
	int wt_mt = lua_gettop(s);

	lua_pushvalue(s,wt_tbl);
	lua_setfield(s,wt_mt,"__index");

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

	lua_settop(s,tbl);
	return 1;
}