// 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 "peer.h"

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

	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)
{
	ssize_t got = read(p->socket, p->in.buffer + p->in.len, p->in.promised - p->in.len);
	if (got < 0) {
		switch (errno) {
			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)
			// 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, uint8_t *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, uint8_t *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, (uint8_t *) &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

short peer_prepare(peer *p)
{
	if (p->in.len == p->in.promised)
		next_in(p, true, sizeof(pkt_header));

	return POLLIN | (p->out.avail ? POLLOUT : 0);
}

bool peer_ready(peer *p, short revents)
{
	bool x = false;

	if (revents & POLLIN)
		x = x || peer_in_ready(p);
	if (revents & POLLOUT)
		x = x || peer_out_ready(p);

	return x;
}