/*
 * 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 <errno.h>
#include <assert.h>

#include <unistd.h>
#include <sys/eventfd.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <lauxlib.h>

#include "melisma.h"
#include "melisma/ring.h"

#define SOCKADDR_TYPENAME   "sockaddr"
#define UDP_SOCKET_TYPENAME "udp socket"
#define UDP_PACKET_SIZE     65536

struct udp_socket
{
	int              fd;
	int              family;
	int              event;
	int              bound;
	int              callb;
	lua_State       *state;
	ml_ring         *recvring;
	ml_ring         *sendring;
	uint8_t          pollbuf[UDP_PACKET_SIZE];
	uint8_t          procbuf[UDP_PACKET_SIZE];
};

// TODO: CHECK BUFFER SIZES FOR ADDRESSES

static int l_sockaddr_unix(lua_State *s)
{
	size_t len;
	const char *path = luaL_checklstring(s,1,&len);

	struct sockaddr_storage *sa = lua_newuserdata(s,sizeof *sa);
	luaL_setmetatable(s,SOCKADDR_TYPENAME);
	memset(sa,0,sizeof *sa);

	((struct sockaddr_un*)sa)->sun_family = AF_UNIX;
	memcpy(((struct sockaddr_un*)sa)->sun_path,path,len);

	return 1;
}

static int l_sockaddr_inet(lua_State *s)
{
	const char *addr = luaL_checkstring(s,1);
	int port = luaL_checkinteger(s,2);

	struct in_addr in_addr;
	if (inet_pton(AF_INET,addr,&in_addr) != 1)
	{
		luaL_error(s,"invalid ipv4 address '%s'",addr);
	}

	struct sockaddr_storage *sa = lua_newuserdata(s,sizeof *sa);
	luaL_setmetatable(s,SOCKADDR_TYPENAME);
	memset(sa,0,sizeof *sa);

	((struct sockaddr_in*)sa)->sin_family = AF_INET;
	((struct sockaddr_in*)sa)->sin_port = htons(port);
	memcpy(&((struct sockaddr_in*)sa)->sin_addr,&in_addr,sizeof in_addr);

	return 1;
}

static void recv_handler(int fd, int flags, void *args)
{
	struct udp_socket *us = args;
	ssize_t len = recvfrom(fd,us->pollbuf,sizeof us->pollbuf,0,NULL,NULL);

	if (len < 0)
	{
		perror("recvfrom");
		return;
	}

	if (!ml_ring_has_space(us->recvring,len + sizeof len))
	{
		fprintf(stderr,"dropped udp packet");
	}

	ml_ring_write_blocking(us->recvring,&len,sizeof len);
	ml_ring_write_blocking(us->recvring,us->pollbuf,len);
}

static void send_handler(int fd, int flags, void *args)
{
	struct udp_socket *us = args;

	size_t sa_len,len;
	struct sockaddr_storage storage = {0};

	ml_ring_read_blocking(us->sendring,&sa_len,sizeof sa_len);
	ml_ring_read_blocking(us->sendring,&storage,sa_len);
	ml_ring_read_blocking(us->sendring,&len,sizeof len);
	ml_ring_read_blocking(us->sendring,us->pollbuf,len);

	if (sendto(us->fd,us->pollbuf,len,0,(struct sockaddr*) &storage,sa_len) < 0)
	{
		perror("sendto");
	}

	uint64_t x;
	if (read(us->event,&x,sizeof x) != sizeof x)
		;
}

static int l_socket(lua_State *s)
{
	int family = lua_tointeger(s,lua_upvalueindex(1));
	int fd = socket(family,SOCK_DGRAM,0);

	if (fd < 0)
	{
		const char *err = strerror(errno);
		luaL_error(s,"failed to create socket: %s",err);
	}

	struct udp_socket *us = calloc(1,sizeof(*us));
	us->fd = fd;
	us->family = family;
	us->sendring = ml_ring_init(UDP_PACKET_SIZE);
	us->event = eventfd(0,EFD_SEMAPHORE);

	ml_add_io_handler(us->event,POLLIN,send_handler,us);

	*(struct udp_socket**) lua_newuserdata(s,sizeof(struct udp_socket*)) = us;
	luaL_setmetatable(s,UDP_SOCKET_TYPENAME);

	return 1;
}

static const char *family_name(int family)
{
	switch (family)
	{
	case AF_UNIX: return "unix";
	case AF_INET: return "inet";
	default:      assert(0);
	}
}

static void proc(const ml_proc_args *args)
{
	struct udp_socket *us = args->userp;
	ssize_t len;

	while (!ml_ring_empty(us->recvring))
	{
		ml_ring_read_blocking(us->recvring,&len,sizeof len);
		ml_ring_read_blocking(us->recvring,us->procbuf,len);

		lua_pushlstring(us->state,(const char*)us->procbuf,len);
		int arg = luaL_ref(us->state,LUA_REGISTRYINDEX);

		ml_schedule_event(0,us->callb,arg);
	}
}

static size_t get_sockaddr_len(struct sockaddr *sa)
{
	if (sa->sa_family == AF_UNIX)
	{
		const char *path = ((struct sockaddr_un*)sa)->sun_path;

		if (path[0] == 0 && path[1] != 0)
		{
			return offsetof(struct sockaddr_un,sun_path) + strlen(path+1) + 1;
		}

		return offsetof(struct sockaddr_un,sun_path) + strlen(path);
	}
	else
	{
		return sizeof(struct sockaddr_in);
	}
}

static void check_family(lua_State *s, struct udp_socket *us, struct sockaddr *sa)
{
	if (us->family != sa->sa_family)
	{
		luaL_error(s,"sockaddr address family (%s) does not match socket address family (%s)",
		           family_name(sa->sa_family), family_name(us->family));
	}
}

static int l_bind(lua_State *s)
{
	struct udp_socket *us = *(struct udp_socket**) luaL_checkudata(s,1,UDP_SOCKET_TYPENAME);
	struct sockaddr *sa = luaL_checkudata(s,2,SOCKADDR_TYPENAME);

	luaL_checktype(s,3,LUA_TFUNCTION);

	if (us->bound)
	{
		luaL_error(s,"socket is already bound");
	}

	check_family(s,us,sa);

	size_t len = get_sockaddr_len(sa);

	if (bind(us->fd,sa,len) < 0)
	{
		const char *err = strerror(errno);
		luaL_error(s,"failed to bind socket: %s",err);
	}

	us->bound = 1;
	us->recvring = ml_ring_init(UDP_PACKET_SIZE);
	us->state = s;

	ml_add_io_handler(us->fd,POLLIN,recv_handler,us);

	lua_pushvalue(s,3);
	us->callb = luaL_ref(s,LUA_REGISTRYINDEX);

	ml_node *node = ml_create_node(proc,us);
	ml_channel *src = ml_add_output(node);
	ml_channel *dst = ml_add_input(ml_get_root_node());
	ml_connect(src,dst);

	return 0;
}

static int l_send(lua_State *s)
{
	struct udp_socket *us = *(struct udp_socket**) luaL_checkudata(s,1,UDP_SOCKET_TYPENAME);
	struct sockaddr *sa = luaL_checkudata(s,2,SOCKADDR_TYPENAME);
	size_t len;
	const char *data = luaL_checklstring(s,3,&len);

	check_family(s,us,sa);

	size_t sa_len = get_sockaddr_len(sa);
	size_t full_len = sizeof len + len + sizeof sa_len + sa_len;

	if (!ml_ring_has_space(us->sendring,full_len))
	{
		lua_pushboolean(s,0);
		return 1;
	}

	ml_ring_write_blocking(us->sendring,&sa_len,sizeof sa_len);
	ml_ring_write_blocking(us->sendring,sa,sa_len);
	ml_ring_write_blocking(us->sendring,&len,sizeof len);
	ml_ring_write_blocking(us->sendring,data,len);

	uint64_t x = 1;
	if (write(us->event,&x,sizeof x) < sizeof x)
		;

	lua_pushboolean(s,1);
	return 1;
}

int luaopen_melisma_udp(lua_State *s)
{
	lua_newtable(s);
	int tbl = lua_gettop(s);

/*** UDP support.

	local mudp = require 'melisma.udp'

	local addr = mudp.sockaddr.inet('127.0.0.1',8000)
	local sock = mudp.socket.inet()

	sock:send(addr,'hello')
	sock:bind(addr,print)

@module udp
*/

/*** Socket address description.
@type sockaddr
*/
	lua_newtable(s);
	int sockaddr_tbl = lua_gettop(s);

/*** Create a unix domain address description.
@static
@function unix
@tparam string path
@treturn sockaddr
*/
	lua_pushcfunction(s,l_sockaddr_unix);
	lua_setfield(s,sockaddr_tbl,"unix");


/*** Create a internet domain address description.
@static
@function inet
@tparam string host
@tparam integer port
@treturn sockaddr
*/
	lua_pushcfunction(s,l_sockaddr_inet);
	lua_setfield(s,sockaddr_tbl,"inet");

	lua_pushvalue(s,sockaddr_tbl);
	lua_setfield(s,tbl,"sockaddr");

	luaL_newmetatable(s,SOCKADDR_TYPENAME);
	int sockaddr_mt = lua_gettop(s);

	lua_pushvalue(s,sockaddr_tbl);
	lua_setfield(s,sockaddr_mt,"__index");

/*** Handle to a socket.
@type socket
*/

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

/*** Create a unix domain socket.
@static
@function unix
@treturn socket
*/
	lua_pushinteger(s,AF_UNIX);
	lua_pushcclosure(s,l_socket,1);
	lua_setfield(s,socket_tbl,"unix");


/*** Create an internet domain socket.
@static
@function inet
@treturn socket
*/
	lua_pushinteger(s,AF_INET);
	lua_pushcclosure(s,l_socket,1);
	lua_setfield(s,socket_tbl,"inet");

/*** Bind a socket to an address.
@function bind
@tparam sockaddr addr Address description.
@tparam function fn Packet callback function.
*/
	lua_pushcfunction(s,l_bind);
	lua_setfield(s,socket_tbl,"bind");

/*** Send a packet on a socket.
@function send
@tparam sockaddr addr Address description.
@tparam string pkt Packet to send.
*/
	lua_pushcfunction(s,l_send);
	lua_setfield(s,socket_tbl,"send");

	lua_pushvalue(s,socket_tbl);
	lua_setfield(s,tbl,"socket");

	luaL_newmetatable(s,UDP_SOCKET_TYPENAME);
	int socket_mt = lua_gettop(s);

	lua_pushvalue(s,socket_tbl);
	lua_setfield(s,socket_mt,"__index");

	lua_settop(s,tbl);
	return 1;
}