#define _POSIX_C_SOURCE 200809L
#include "sleep_utp.h"

#include <assert.h>
#include <errno.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <poll.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

static char received[4096];
static size_t received_len;

static void on_data(sleep_utp_t *utp, const uint8_t *data, size_t length)
{
	(void)utp;
	assert(length <= sizeof(received) - received_len);
	memcpy(received + received_len, data, length);
	received_len += length;
}

static ssize_t receive_packet(int fd, uint8_t *packet, size_t capacity,
			      struct sockaddr_storage *source, socklen_t *source_len)
{
	struct pollfd pfd = {fd, POLLIN, 0};
	assert(poll(&pfd, 1, 1000) == 1);
	return recvfrom(fd, packet, capacity, 0, (struct sockaddr *)source, source_len);
}

int main(void)
{
	int a_fd = socket(AF_INET, SOCK_DGRAM, 0), b_fd = socket(AF_INET, SOCK_DGRAM, 0);
	struct sockaddr_in a_addr, b_addr;
	socklen_t a_len = sizeof(a_addr), b_len = sizeof(b_addr);
	struct sockaddr_storage source;
	socklen_t source_len = sizeof(source);
	sleep_utp_t a, b;
	sleep_utp_packet_t packet;
	uint8_t wire[1400];
	ssize_t n;
	uint16_t sent_sequence;
	assert(a_fd >= 0 && b_fd >= 0);
	memset(&a_addr, 0, sizeof(a_addr)); memset(&b_addr, 0, sizeof(b_addr));
	a_addr.sin_family = AF_INET; a_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
	b_addr.sin_family = AF_INET; b_addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
	assert(bind(a_fd, (struct sockaddr *)&a_addr, sizeof(a_addr)) == 0);
	assert(bind(b_fd, (struct sockaddr *)&b_addr, sizeof(b_addr)) == 0);
	assert(getsockname(a_fd, (struct sockaddr *)&a_addr, &a_len) == 0);
	assert(getsockname(b_fd, (struct sockaddr *)&b_addr, &b_len) == 0);
	assert(sleep_utp_init(&a, a_fd) == 0 && sleep_utp_init(&b, b_fd) == 0);
	a.on_data = on_data; b.on_data = on_data;
	assert(sleep_utp_connect(&a, (struct sockaddr *)&b_addr, b_len) == 0);
	n = receive_packet(b_fd, wire, sizeof(wire), &source, &source_len);
	assert(n > 0 && sleep_utp_packet_decode(&packet, wire, (size_t)n) == 0);
	assert(packet.type == SLEEP_UTP_SYN);
	assert(sleep_utp_accept(&b, &packet, (struct sockaddr *)&source, source_len) == 0);
	n = receive_packet(a_fd, wire, sizeof(wire), &source, &source_len);
	assert(n > 0 && sleep_utp_input(&a, wire, (size_t)n) == 0);
	assert(a.state == SLEEP_UTP_CONNECTED);
	{
		uint8_t oversized[SLEEP_UTP_MTU + 1];
		size_t wire_len;
		memset(&packet, 0, sizeof(packet));
		packet.type = SLEEP_UTP_DATA; packet.version = 1;
		packet.payload = oversized; packet.payload_len = sizeof(oversized);
		assert(sleep_utp_packet_encode(&packet, wire, sizeof(wire), &wire_len) == -1);
		assert(errno == EMSGSIZE);
		memset(wire, 0, sizeof(wire));
		wire[0] = (uint8_t)(SLEEP_UTP_DATA << 4) | 1;
		assert(sleep_utp_packet_decode(&packet, wire, 20 + sizeof(oversized)) == -1);
		assert(errno == EMSGSIZE);
		wire[0] = (uint8_t)(SLEEP_UTP_STATE << 4) | 1;
		assert(sleep_utp_packet_decode(&packet, wire, 21) == -1);
		assert(errno == EPROTO);
	}
	{
		sleep_utp_slot_t *slot = &a.unacked[a.sequence % SLEEP_UTP_WINDOW];
		struct pollfd pfd = {b_fd, POLLIN, 0};
		slot->used = 1;
		assert(sleep_utp_send(&a, (const uint8_t *)"blocked", 7) == -1);
		assert(errno == EAGAIN);
		assert(poll(&pfd, 1, 20) == 0);
		slot->used = 0;
	}
	packet.type = SLEEP_UTP_STATE; packet.version = 1;
	packet.connection_id = (uint16_t)(a.receive_id + 1);
	{
		size_t wire_len;
		assert(sleep_utp_packet_encode(&packet, wire, sizeof(wire), &wire_len) == 0);
		assert(sleep_utp_input(&a, wire, wire_len) == -1 && errno == EPROTO);
	}
	assert(sleep_utp_send(&a, (const uint8_t *)"first", 5) == 0);
	n = receive_packet(b_fd, wire, sizeof(wire), &source, &source_len);
	assert(n > 0 && sleep_utp_input(&b, wire, (size_t)n) == 0);
	assert(received_len == 5 && !memcmp(received, "first", 5));
	n = receive_packet(a_fd, wire, sizeof(wire), &source, &source_len);
	assert(n > 0 && sleep_utp_input(&a, wire, (size_t)n) == 0);
	assert(a.bytes_in_flight == 0);
	memset(&packet, 0, sizeof(packet));
	packet.type = SLEEP_UTP_DATA; packet.version = 1; packet.connection_id = a.send_id;
	packet.sequence = 3; packet.window_size = 1024 * 1024;
	packet.payload = (const uint8_t *)"later"; packet.payload_len = 5;
	{
		size_t wire_len;
		assert(sleep_utp_packet_encode(&packet, wire, sizeof(wire), &wire_len) == 0);
		assert(sendto(a_fd, wire, wire_len, 0, (struct sockaddr *)&b_addr, b_len) == (ssize_t)wire_len);
	}
	n = receive_packet(b_fd, wire, sizeof(wire), &source, &source_len);
	assert(n > 0 && sleep_utp_input(&b, wire, (size_t)n) == 0);
	assert(received_len == 5); /* sequence 3 waits for sequence 2 */
	packet.sequence = 2; packet.payload = (const uint8_t *)"next"; packet.payload_len = 4;
	{
		size_t wire_len;
		assert(sleep_utp_packet_encode(&packet, wire, sizeof(wire), &wire_len) == 0);
		assert(sendto(a_fd, wire, wire_len, 0, (struct sockaddr *)&b_addr, b_len) == (ssize_t)wire_len);
	}
	n = receive_packet(b_fd, wire, sizeof(wire), &source, &source_len);
	assert(n > 0 && sleep_utp_input(&b, wire, (size_t)n) == 0);
	assert(received_len == 14 && !memcmp(received, "firstnextlater", 14));
	for (int ack = 0; ack < 2; ack++) {
		n = receive_packet(a_fd, wire, sizeof(wire), &source, &source_len);
		assert(n > 0 && sleep_utp_input(&a, wire, (size_t)n) == 0);
	}
	{
		uint8_t large[2401];
		size_t before = received_len;
		/* The preceding out-of-order packets advanced the receiver through seq 3. */
		a.sequence = 4;
		memset(large, 0x5a, sizeof(large));
		assert(sleep_utp_send(&a, large, sizeof(large)) == 0);
		for (int part = 0; part < 3; part++) {
			n = receive_packet(b_fd, wire, sizeof(wire), &source, &source_len);
			assert(n > 20 && sleep_utp_input(&b, wire, (size_t)n) == 0);
			n = receive_packet(a_fd, wire, sizeof(wire), &source, &source_len);
			assert(n >= 20 && sleep_utp_input(&a, wire, (size_t)n) == 0);
		}
		assert(received_len == before + sizeof(large));
		for (size_t i = before; i < received_len; i++) assert(received[i] == 0x5a);
		assert(a.bytes_in_flight == 0);
	}
	assert(sleep_utp_send(&a, (const uint8_t *)"retry", 5) == 0);
	n = receive_packet(b_fd, wire, sizeof(wire), &source, &source_len);
	assert(n > 0 && sleep_utp_packet_decode(&packet, wire, (size_t)n) == 0);
	sent_sequence = packet.sequence;
	assert(sleep_utp_tick(&a, 1) == 0);
	assert(sleep_utp_tick(&a, 2000) == 0);
	n = receive_packet(b_fd, wire, sizeof(wire), &source, &source_len);
	assert(n > 0 && sleep_utp_packet_decode(&packet, wire, (size_t)n) == 0);
	assert(packet.type == SLEEP_UTP_DATA && packet.sequence == sent_sequence);
	close(a_fd); close(b_fd);
	return 0;
}