// SPDX-FileCopyrightText: 2024 Lizzy Fleckenstein <lizzy@vlhl.dev>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <poll.h>
#include <errno.h>
#include <netdb.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include "net.h"

void invalid_pkt(peer *from, str pkt)
{
	// TODO: maybe inform peer about failure? not sure
	// FIXME: hexdumping a gazillon bytes to stderr might be an issue
	fprintf(stderr, "invalid pkt from %*s: ", PSTR(*from->name));
	for (size_t i = 0; i < pkt.len; i++)
		fprintf(stderr, "%02x%c", (uint8_t) pkt.data[i], i+1 == pkt.len ? '\n' : ' ');
}

int socket_create(const char *host, const char *port, bool server)
{
	struct addrinfo *info = NULL, hints = {0};
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;

	int err;
	if ((err = getaddrinfo(host, port, &hints, &info))) {
		fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(err));
		return -1;
	}

#define TRY(op, val) if ((val) < 0) { perror(op); close(fd); freeaddrinfo(info); return -1; }

	int fd = -1;
	TRY("socket", fd = socket(info->ai_family, info->ai_socktype, info->ai_protocol))

	int flag = 1;
	TRY("setsockopt", setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (void *) &flag, sizeof flag))

	if (server) {
		TRY("bind", bind(fd, info->ai_addr, info->ai_addrlen))
		TRY("listen", listen(fd, 3))
	} else {
		TRY("connect", connect(fd, info->ai_addr, info->ai_addrlen))
		TRY("fcntl", SET_NONBLOCK(fd))
	}

#undef TRY

	freeaddrinfo(info);
	return fd;
}

int socket_accept(int accept_fd)
{
	int socket = accept(accept_fd, NULL, NULL); // TODO: save address
	if (socket < 0) {
		perror("accept");
		return -1;
	}

	if (SET_NONBLOCK(socket) < 0) {
		perror("fcntl");
		close(socket);
		return -1;
	}

	return socket;
}

void peer_init(peer *p, int socket, str *name)
{
	p->socket = socket;
	p->disco = false;
	p->name = name;

	p->in.header = true;
	p->in.len = 0;
	p->in.promised = 0;
	p->in.buffer = malloc(PEER_INBUFFER_SIZE);

	p->out.cursor = 0;
	p->out.avail = 0;
	p->out.buffer = malloc(PEER_OUTBUFFER_SIZE);
}

void peer_free(peer *p)
{
	close(p->socket);
	free(p->in.buffer);
	free(p->out.buffer);
}

// in

static void next_in(peer *p, bool header, size_t len)
{
	p->in.header = header;
	p->in.len = 0;
	p->in.buffer = malloc(p->in.promised = len);
}

static bool peer_in_ready(peer *p)
{
	errno = 0;
	ssize_t got = read(p->socket, p->in.buffer + p->in.len, p->in.promised - p->in.len);
	if (got <= 0) {
		switch (errno) {
			case 0: case ECONNRESET:
				p->disco = true;
				return true;
			case EINTR:
				return peer_in_ready(p); // retry
			default:
				perror("read");
				return false;
		}
	}

	p->in.len += got;
	if (p->in.len == p->in.promised && p->in.len != 0) {
		if (!p->in.header)
			return true;

		size_t len = *(pkt_header *) p->in.buffer;
		if (len > PEER_INBUFFER_SIZE || len == 0)
			// TODO: figure out what to do if packet too large (disconnect?)
			next_in(p, true, sizeof(pkt_header));
		else
			next_in(p, false, len);
	}

	return false;
}

// out

static void send_raw(peer *p, void *data, size_t len)
{
	memcpy(p->out.buffer + p->out.cursor + p->out.avail, data, len);
	p->out.avail += len;
}

static bool out_space(peer *p, size_t len)
{
	if (len + p->out.avail > PEER_OUTBUFFER_SIZE)
		return false;

	if (p->out.cursor + p->out.avail + len > PEER_OUTBUFFER_SIZE) {
		memmove(p->out.buffer, p->out.buffer + p->out.cursor, p->out.avail);
		p->out.cursor = 0;
	}

	return true;
}

bool peer_send(peer *p, void *data, size_t len)
{
	if (len > PEER_INBUFFER_SIZE)
		return false;

	pkt_header hdr = (pkt_header) len;

	if (!out_space(p, sizeof hdr + len))
		return false;

	send_raw(p, &hdr, sizeof hdr);
	send_raw(p, data, len);

	return true;
}

static bool peer_out_ready(peer *p)
{
	ssize_t written = write(p->socket, p->out.buffer + p->out.cursor, p->out.avail);
	if (written < 0) {
		switch (errno) {
			case ECONNRESET:
				p->disco = true;
				return true;
			case EINTR:
				return peer_out_ready(p);
			default:
				perror("write");
				return false;
		}
	}

	p->out.avail -= written;
	if (p->out.avail == 0)
		p->out.cursor = 0;

	return false;
}

// poll

struct pollfd peer_prepare(peer *p)
{
	struct pollfd pfd = {0};

	pfd.fd = p->socket;
	if (p->in.len == p->in.promised)
		next_in(p, true, sizeof(pkt_header));

	pfd.events = POLLIN | (p->out.avail ? POLLOUT : 0);

	return pfd;
}

str peer_recv(peer *p, struct pollfd pfd)
{
	if (pfd.revents & (POLLHUP | POLLERR))
		p->disco = true;

	if (!p->disco && pfd.revents & POLLOUT)
		peer_out_ready(p);

	if (!p->disco && pfd.revents & POLLIN)
		if (peer_in_ready(p))
			return (str) { p->in.len, (char *) p->in.buffer };

	return NILS;
}