#include <stdio.h>
#include <err.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <sys/un.h>
#include <sys/socket.h>

#include "llm.h"

#define SOCK_PATH  "/var/run/lex.sock"
#define MODEL_PATH "qwen2.5-3b-instruct-q4_k_m.gguf"

static void usage(const char *prog)
{
	errx(1, "usage: %s [-m model_path]", prog);
}

static void loop(struct llm_ctx *llm, int lsock)
{
	int csock, c;
	FILE *cfp;
	size_t len;
	char line[4096];

	for (;;) {
		csock = accept(lsock, NULL, NULL);
		if (csock == -1) {
			if (errno != EINTR)
				warn("accept()");
			continue;
		}

		if (!(cfp = fdopen(csock, "r+"))) {
			warn("cfp fdopen()");
			close(csock);
			continue;
		}

		if (fgets(line, sizeof(line), cfp) != NULL) {
			len = strcspn(line, "\n");
			if (line[len] == '\0' && len == sizeof(line) - 1) {
				warnx("request too large");
				// drain and skip
				while ((c = fgetc(cfp)) != '\n' && c != EOF)
					;
				fclose(cfp);
				continue;
			}

			line[len] = '\0';
			if (line[0] != '\0')
				llm_run(llm, line, cfp);
		}

		fclose(cfp);
	}
}

int sock_init(void)
{
	int sock;
	struct sockaddr_un addr;

	sock = socket(AF_UNIX, SOCK_STREAM, 0);
	if (sock == -1)
		err(1, "socket()");

	memset(&addr, 0, sizeof(addr));
	addr.sun_family = AF_UNIX;
	strlcpy(addr.sun_path, SOCK_PATH, sizeof(addr.sun_path));

	unlink(SOCK_PATH); /* remove stale socket from a previous run */
	
	if (bind(sock, (struct sockaddr *)&addr, sizeof(addr)) == -1)
		err(1, "bind()");

	if (listen(sock, 1) == -1) /* backlog = 1 */
		err(1, "listen()");

	return sock;
}

int main(int argc , char *argv[])
{
	int sock, opt;
	const char *prog;
	const char *model_path;
	struct llm_ctx *llm;

	prog = argv[0];
	model_path = MODEL_PATH;

	while ((opt = getopt(argc, argv, "m:")) != -1) {
		switch (opt) {
		case 'm':
			model_path = optarg;
			break;
		default:
			usage(prog);
		}
	}

	argc -= optind;
	argv += optind;
	if (argc != 0) 
		usage(prog);

	if (unveil(model_path, "r") == -1)
		err(1, "unveil %s failed", model_path);
	if (unveil(SOCK_PATH, "rwc") == -1)
		err(1, "unveil %s failed", SOCK_PATH);
	if (unveil(NULL, NULL) == -1)
		err(1, "unveil lock failed");

	sock = sock_init();

	if (!(llm = llm_init(model_path)))
		errx(1, "Failed to load model");

	/* llm_init() calls mlock, hence pledge() after that. */
	if (pledge("stdio unix", NULL) == -1)
		err(1, "secondary pledge failed");

	loop(llm, sock); /* doesn't return */

	return 0;
}