#define _GNU_SOURCE
#include <arpa/inet.h>
#include <assert.h>
#include <err.h>
#include <fcntl.h>
#include <inttypes.h>
#include <limits.h>
#include <linux/if_packet.h>
#include <net/ethernet.h>
#include <net/if.h>
#include <net/if_arp.h>
#include <poll.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <time.h>
#include <unistd.h>

#define EM_PROTOCOL_VERSION 0
#define EM_MESSAGE_MAX_LENGTH (1500 - 2 - 2 - 2)

#define EMT_SPEAK_VERSION 0
#define EMT_STATUS_REQUEST 1
#define EMT_STATUS 2
#define EMT_MSGID_REQUEST 3
#define EMT_MSGID 4
#define EMT_MESSAGE 5
#define EMT_ACK 6

#define EMS_AVAILABLE 0
#define EMS_UNAVAILABLE 1
#define EMS_OFFLINE 2

unsigned const char veth0a_mac[6] = {0xb2, 0xc8, 0x5b, 0x78, 0xb4, 0xef}; //debg
unsigned const char veth0b_mac[6] = {0xf6, 0x18, 0xfd, 0x2a, 0x80, 0xf3}; //debg

bool running = true;

int packet_socket;
int urandom;

unsigned char own_mac[6];
unsigned char broadcast_mac[6] = {0xff, 0xff, 0xff, 0xff, 0xff, 0xff};

unsigned char own_status = EMS_AVAILABLE;
unsigned char own_nick[256];
unsigned char own_nick_length = 0;
time_t next_status_broadcast;

enum message_send_states {IDLE, QUEUED, SENDING};
enum message_send_states own_message_send_state = IDLE;
unsigned char own_message_destination_mac[6];
unsigned char own_message[EM_MESSAGE_MAX_LENGTH];
size_t own_message_length = 0;
uint16_t own_message_msgid = 0;

struct msgid_cache_entry {
	unsigned char other_mac[6];
	bool know_send;
	bool know_receive;
	uint16_t next_send;
	uint16_t next_receive;
};

struct msgid_cache_entry msgid_cache[256];
ssize_t msgid_cache_fill = 0;
unsigned char next_slot = 0;

ssize_t msgid_cache_lookup(const unsigned char mac[6]) {
	for (ssize_t i = 0; i < msgid_cache_fill; i++) {
		if (memcmp(msgid_cache[i].other_mac, mac, 6) == 0) {
			// Found it
			return i;
		}
	}

	// Did not find it
	return -1;
}

ssize_t msgid_cache_add(const unsigned char mac[6]) {
	ssize_t index = next_slot++;

	// If we are adding into a new slot (instead of overwriting one),
	// expand the fill pointer
	//
	// + 1 because msgid_cache_fill of N means that cache slots [0, N-1]
	// are in use
	if (msgid_cache_fill < index + 1) {
		msgid_cache_fill = index + 1;
	}

	memcpy(msgid_cache[index].other_mac, mac, sizeof(msgid_cache[index].other_mac));
	msgid_cache[index].know_send = false;
	msgid_cache[index].know_receive = false;

	return index;
}

unsigned char random_byte(void) {
	unsigned char randomness;
	if (read(urandom, &randomness, 1) != 1) {
		err(1, "read");
	}
	return randomness;
}

time_t monotonic_time(void) {
	struct timespec now;
	if (clock_gettime(CLOCK_MONOTONIC, &now) == -1) {
		err(1, "clock_gettime");
	}
	return now.tv_sec;
}

char hexify(int nybble) {
	assert(0 <= nybble && nybble <= 16);
	return "0123456789abcdef"[nybble];
}

void format_mac(const unsigned char binary_address[6], char formatted[18]) {
	for (size_t i = 0; i < 6; i++) {
		unsigned char byte = binary_address[i];
		formatted[3*i] = hexify(byte >> 4);
		formatted[3*i + 1] = hexify(byte & 0xf);
		formatted[3*i + 2] = ':';
	}
	formatted[17] = '\0';
}

void drop_privileges(void) {
	uid_t uid = getuid();
	gid_t gid = getgid();

	if (setresgid(gid, gid, gid) == -1) {
		err(1, "setresgid");
	}
	if (setresuid(uid, uid, uid) == -1) {
		err(1, "setresuid");
	}
}

void send_frame(const unsigned char *frame, size_t frame_length) {
	if (write(packet_socket, frame, frame_length) == -1) {
		err(1, "write");
	}
}

void write_headers(unsigned char frame[14], const unsigned char destination_mac[6], unsigned char packet_type) {
	// Destination MAC
	memcpy(&frame[0], destination_mac, 6);

	// Source MAC
	memcpy(&frame[6], own_mac, 6);

	// EtherType
	frame[12] = 0xda;
	frame[13] = 0x7a;

	// Ethermess version
	frame[14] = EM_PROTOCOL_VERSION;

	// Ethermess packet type
	frame[15] = packet_type;
}

void send_speak_version(const unsigned char destination[6]) {
	unsigned char frame[14 + 2 + 1];

	write_headers(frame, destination, EMT_SPEAK_VERSION);

	// Version to speak
	frame[16] = EM_PROTOCOL_VERSION;

	send_frame(frame, sizeof(frame));
}

void send_status_request(const unsigned char destination[6]) {
	unsigned char frame[14 + 2];

	write_headers(frame, destination, EMT_STATUS_REQUEST);

	send_frame(frame, sizeof(frame));
}

void send_status(const unsigned char destination[6]) {
	unsigned char frame[14 + 2 + 1 + 1 + own_nick_length];

	write_headers(frame, destination, EMT_STATUS);

	// Status
	frame[16] = own_status;

	// Length of nick
	frame[17] = own_nick_length;

	// Nick
	memcpy(&frame[18], own_nick, own_nick_length);

	send_frame(frame, sizeof(frame));
}

void send_msgid_request(const unsigned char destination[6]) {
	unsigned char frame[14 + 2];

	write_headers(frame, destination, EMT_MSGID_REQUEST);

	send_frame(frame, sizeof(frame));
}

void send_msgid(const unsigned char destination[6]) {
	unsigned char frame[14 + 2 + 2];

	write_headers(frame, destination, EMT_MSGID);

	// Look up destination in the ID cache
	ssize_t cache_index = msgid_cache_lookup(destination);
	if (cache_index == -1) {
		// Not in the cache
		// Create a new entry
		cache_index = msgid_cache_add(destination);
	}

	if (!msgid_cache[cache_index].know_receive) {
		// We don't have receive ID stored in the cache
		// In that case, start from a random index
		msgid_cache[cache_index].next_receive = (random_byte() << 8) | random_byte();
		msgid_cache[cache_index].know_receive = true;
	}

	// Message ID of next message we're waiting to receive
	uint16_t msgid = msgid_cache[cache_index].next_receive;

	frame[16] = msgid >> 8;
	frame[17] = msgid & 0xff;

	send_frame(frame, sizeof(frame));
}

void send_message(void) {
	unsigned char frame[14 + 2 + 2 + 2 + own_message_length];

	write_headers(frame, own_message_destination_mac, EMT_MESSAGE);

	// Message ID
	frame[16] = own_message_msgid >> 8;
	frame[17] = own_message_msgid & 0xff;

	// Message length
	frame[18] = own_message_length >> 8;
	frame[19] = own_message_length & 0xff;

	// Message
	memcpy(&frame[20], own_message, own_message_length);

	send_frame(frame, sizeof(frame));
}

void send_ack(const unsigned char destination[6], uint16_t msgid) {
	unsigned char frame[14 + 2 + 2];

	write_headers(frame, destination, EMT_ACK);

	// Message ID
	frame[16] = msgid >> 8;
	frame[17] = msgid & 0xff;

	send_frame(frame, sizeof(frame));
}

void read_command(void) {
	int cmd = getchar();
	if (cmd == EOF) {
		err(1, "getchar");
	}

	unsigned const char *other_mac;
	if (memcmp(own_mac, veth0a_mac, 6) == 0) {
		other_mac = veth0b_mac;
	} else {
		other_mac = veth0a_mac;
	}

	if (cmd == 'q') {
		running = false;
	} else if (cmd == 's') {
		send_status_request(other_mac);
	} else if (cmd == 'i') {
		send_msgid_request(other_mac);
	} else if (cmd == 'm') {
		memcpy(own_message, "Hello, world!", 13);
		own_message_length = 13;
		memcpy(own_message_destination_mac, other_mac, 6);
		own_message_send_state = QUEUED;
	} else if (cmd == '\n') {
		// Ignore
	} else {
		fprintf(stderr, "?"); //debg
	}
}

bool check_padding(const unsigned char *data, size_t index, size_t data_length) {
	// Valid padding is all zero bytes
	assert(index <= data_length);
	for (size_t i = index; i < data_length; i++) {
		if (data[i] != 0) {
			// Check failed
			return false;
		}
	}

	// Check succeeded
	return true;
}

void handle_status(const unsigned char source_mac[6], const unsigned char *data, size_t data_length) {
	if (data_length < 2) {
		// Too short
		fprintf(stderr, "Data too short: %zu\n", data_length); // debg
		return;
	}

	unsigned char status = data[0];

	if (status != EMS_AVAILABLE && status != EMS_UNAVAILABLE && status != EMS_OFFLINE) {
		// Unknown status, throw away
		fprintf(stderr, "Unknown status %u\n", status); // debg
		return;
	}

	unsigned char nick_length = data[1];

	if (nick_length > data_length - 2) {
		// Malformed length field
		fprintf(stderr, "Nick length %u, remaining packet length %zu\n", nick_length, data_length); // debg
		return;
	}

	unsigned char nick[nick_length];
	memcpy(nick, &data[2], nick_length);

	if (!check_padding(data, 2 + nick_length, data_length)) {
		// Malformed padding
		return;
	}

	// TODO: check that nick is valid utf-8 with no control chars

	char mac[18];
	format_mac(source_mac, mac);

	if (printf("%s status: ", mac) == -1) {
		err(1, "printf");
	}

	if (status == EMS_UNAVAILABLE) {
		if (printf("(unavailable) ") == -1) {
			err(1, "printf");
		}
	} else if (status == EMS_OFFLINE) {
		if (printf("(offline) ") == -1) {
			err(1, "printf");
		}
	}

	for (size_t i = 0; i < (size_t)nick_length; i++) {
		if (putchar(nick[i]) == EOF) {
			err(1, "putchar");
		}
	}

	if (putchar('\n') == EOF) {
		err(1, "putchar");
	}

	if (fflush(stdout) == EOF) {
		err(1, "fflush");
	}
}

void handle_msgid(const unsigned char source_mac[6], const unsigned char *data, size_t data_length) {
	if (data_length < 2) {
		// Too short
		fprintf(stderr, "Data too short: %zu\n", data_length); // debg
		return;
	}

	uint16_t msgid = (data[0] << 8) | data[1];

	if (!check_padding(data, 2, data_length)) {
		// Malformed padding
		return;
	}

	ssize_t cache_index = msgid_cache_lookup(source_mac);
	if (cache_index == -1) {
		// Not in the cache, so add it there
		cache_index = msgid_cache_add(source_mac);
	}

	msgid_cache[cache_index].next_send = msgid;
	msgid_cache[cache_index].know_send = true;

	char mac[18];
	format_mac(source_mac, mac);

	if (printf("%s awaits message ID %" PRIu16 "\n", mac, msgid_cache[cache_index].next_send) == 0) {
		err(1, "printf");
	}

	if (fflush(stdout) == EOF) {
		err(1, "fflush");
	}
}

void handle_message(const unsigned char source_mac[6], const unsigned char *data, size_t data_length) {
	if (data_length < 4) {
		// Too short
		fprintf(stderr, "Data too short: %zu\n", data_length); // debg
		return;
	}

	uint16_t msgid = (data[0] << 8) | data[1];

	uint16_t message_length = (data[2] << 8) | data[3];

	if (message_length > data_length - 4 || message_length > EM_MESSAGE_MAX_LENGTH) {
		// Malformed length field
		fprintf(stderr, "Message length %u, remaining packet length %zu, max message length %u\n", message_length, data_length, EM_MESSAGE_MAX_LENGTH); // debg
		return;
	}

	unsigned char message[message_length];
	memcpy(message, &data[4], message_length);

	if (!check_padding(data, 4 + message_length, data_length)) {
		// Malformed padding
		return;
	}

	// TODO: Check that the message is valid utf-8 with newline as the only control char

	// TODO: Check whether we've received this message already and update msgid cache

	char mac[18];
	format_mac(source_mac, mac);

	if (printf("%s (%" PRIu16 "): ", mac, msgid) == -1) {
		err(1, "printf");
	}

	for (size_t i = 0; i < message_length; i++) {
		if (putchar(message[i]) == EOF) {
			err(1, "putchar");
		}
	}

	if (putchar('\n') == EOF) {
		err(1, "putchar");
	}

	if (fflush(stdout) == EOF) {
		err(1, "fflush");
	}

	send_ack(source_mac, msgid);
}

void handle_ack(const unsigned char source_mac[6], const unsigned char *data, size_t data_length) {
	if (data_length < 2) {
		// Too short
		fprintf(stderr, "Data too short: %zu\n", data_length); // debg
		return;
	}

	uint16_t msgid = (data[0] << 8) | data[1];

	if (!check_padding(data, 2, data_length)) {
		// Malformed padding
		return;
	}

	char mac[18];
	format_mac(source_mac, mac);

	if (printf("%s ack %" PRIu16 "\n", mac, msgid) == -1) {
		err(1, "printf");
	}

	if (own_message_send_state == SENDING && msgid == own_message_msgid) {
		own_message_send_state = IDLE;
	}
}

void process_frame(void) {
	unsigned char frame[1518]; // Largest a 802.3 frame can be without FCS

	ssize_t res = recv(packet_socket, frame, sizeof(frame), 0);
	if (res == -1) {
		err(1, "recv");
	} else if (res < 16) {
		// Frame too short to contain enough information
		return;
	}
	size_t packet_length = (size_t)res;

	// Check that the packet is Ethermess (EtherType DA7A)
	if (frame[12] != 0xda || frame[13] != 0x7a) {
		return;
	}

	if (memcmp(frame, own_mac, 6) == 0) {
		// Targetted at us
		fprintf(stderr, "."); // debg
	} else if (memcmp(frame, broadcast_mac, 6) == 0) {
		// Broadcast
		fprintf(stderr, "^"); // debg
	} else {
		// Does not concern us
		return;
	}

	// Extract source MAC
	unsigned char source_mac[6];
	memcpy(source_mac, &frame[6], sizeof(source_mac));

	// Extract version
	unsigned char version = frame[14];

	// If they speak a version we don't understand, tell them to speak ours
	if (version > EM_PROTOCOL_VERSION) {
		fprintf(stderr, "Protocol version mismatch: %u\n", version); // debg
		send_speak_version(source_mac);
		return;
	}

	// Extract Ethermess packet type
	unsigned char packet_type = frame[15];

	// Process the packet based on the packet type
	switch (packet_type) {
		case EMT_STATUS_REQUEST:
			if (check_padding(&frame[16], 0, packet_length - 16)) {
				send_status(source_mac);
			}
			break;

		case EMT_STATUS:
			handle_status(source_mac, &frame[16], packet_length - 16);
			break;

		case EMT_MSGID_REQUEST:
			if (check_padding(&frame[16], 0, packet_length - 16)) {
				send_msgid(source_mac);
			}
			break;

		case EMT_MSGID:
			handle_msgid(source_mac, &frame[16], packet_length - 16);
			break;

		case EMT_MESSAGE:
			handle_message(source_mac, &frame[16], packet_length - 16);
			break;

		case EMT_ACK:
			handle_ack(source_mac, &frame[16], packet_length - 16);
			break;

		default:
			fprintf(stderr, "Ignoring packet of type %i\n", packet_type);
	}
}

void eventloop(void) {
	// Listen on both stdin for commands and network interface for packets
	struct pollfd pollfds[2];

	// stdin
	pollfds[0].fd = 0;
	pollfds[0].events = POLLIN;

	// Network interface
	pollfds[1].fd = packet_socket;
	pollfds[1].events = POLLIN;

	while (running) {
		// (Attempt) to process a message send
		if (own_message_send_state == QUEUED) {
			// We need to have the correct msgid to be able to send
			ssize_t cache_index = msgid_cache_lookup(own_message_destination_mac);

			if (cache_index == -1 || !msgid_cache[cache_index].know_send) {
				// We don't know what the msgid should be
				// -> ask the other side
				// TODO: Implement a timer
				send_msgid_request(own_message_destination_mac);
			} else {
				// It is in the cache
				own_message_msgid = msgid_cache[cache_index].next_send++;
				own_message_send_state = SENDING;
			}
		}

		if (own_message_send_state == SENDING) {
			send_message();
		}

		// Figure out how many ms to wait
		int wait_ms;
		time_t now = monotonic_time();
		if (next_status_broadcast <= now) {
			// The time has come to send the status broadcast
			send_status(broadcast_mac);
			// Do next one in about 5 minutes
			next_status_broadcast = now + 5 * 60 + random_byte() / 64;
		} else {
			if (INT_MAX / 1000 >= next_status_broadcast - now) {
				// Wail until next status broadcast is due
				wait_ms = (next_status_broadcast - now) * 1000;
			} else {
				// Would overflow, wait INT_MAX ms
				wait_ms = INT_MAX;
			}
		}

		int ready = poll(pollfds, sizeof(pollfds) / sizeof(*pollfds), wait_ms);
		if (ready == -1) {
			err(1, "poll");
		}

		// stdin
		if (ready > 0 && pollfds[0].revents != 0) {
			ready--;

			if (pollfds[0].revents & POLLIN) {
				// Read a command
				read_command();
			} else {
				errx(1, "Got poll event %hd on stdin\n", pollfds[0].revents);
			}
		}

		// packet_socket
		if (ready > 0 && pollfds[1].revents != 0) {
			ready--;

			if (pollfds[1].revents & POLLIN) {
				// Process a frame
				process_frame();
			} else {
				errx(1, "Got poll event %hd on packet socket\n", pollfds[1].revents);
			}
		}

		if (ready > 0) {
			errx(1, "poll(1) says we have ready fds, but neither was ready");
		}
	}
}

int main(int argc, char **argv) {
	if (argc != 2) {
		fprintf(stderr, "Usage: %s interface\n", argv[0]);
		exit(1);
	}

	const char *interface_name = argv[1];

	// Open /dev/urandom for getting randomness
	urandom = open("/dev/urandom", O_RDONLY);

	// Create a packet socket
	packet_socket = socket(PF_PACKET, SOCK_RAW, htons(ETH_P_ALL));
	if (packet_socket == -1) {
		err(1, "socket");
	}

	// Only creating the socket requires root privs
	drop_privileges();

	// Find the index of the network interface
	struct ifreq ifr;
	strncpy(ifr.ifr_name, interface_name, IFNAMSIZ);
	if (ioctl(packet_socket, SIOCGIFINDEX, &ifr) == -1) {
		err(1, "ioctl");
	}

	// Bind to the network interface
	struct sockaddr_ll sll;
	sll.sll_family = AF_PACKET;
	sll.sll_protocol = htons(ETH_P_ALL);
	sll.sll_ifindex = ifr.ifr_ifindex;
	if (bind(packet_socket, (const struct sockaddr*)&sll, sizeof(sll)) == -1) {
		err(1, "bind");
	}

	// Get our own MAC
	strncpy(ifr.ifr_name, interface_name, IFNAMSIZ);
	if (ioctl(packet_socket, SIOCGIFHWADDR, &ifr) == -1) {
		err(1, "ioctl");
	}
	if (ifr.ifr_hwaddr.sa_family != ARPHRD_ETHER) {
		errx(1, "Not an Ethernet interface");
	}
	memcpy(own_mac, ifr.ifr_hwaddr.sa_data, sizeof(own_mac));

	// Print it out
	char own_mac_str[18];
	format_mac(own_mac, own_mac_str);
	fprintf(stderr, "%s\n", own_mac_str);

	// Set our nick
	if (memcmp(own_mac, veth0a_mac, 6) == 0) {
		memcpy(own_nick, "foo", 3);
		own_nick_length = 3;
	} else {
		memcpy(own_nick, "bar", 3);
		own_nick_length = 3;
	}

	// Initialize the message id cache
	memset(msgid_cache, 0, sizeof(msgid_cache));

	// Broadcast our status to the network to let them know we're here
	send_status(broadcast_mac);

	// Schedule next broadcast of our status about 5 min in the future
	next_status_broadcast = monotonic_time() + 5 * 60 + random_byte() / 64;

	// Request status from everyone, so that we can get an idea of who is on the network
	send_status_request(broadcast_mac);

	// Start the event loop
	eventloop();

	// Close the socket (tho I'm not 100% sure it's needed)
	if (close(packet_socket) == -1) {
		err(1, "close");
	}

	// Flush stdout
	if (fflush(stdout) == EOF) {
		err(1, "fflush");
	}

	// Close urandom
	if (close(urandom) == -1) {
		err(1, "close");
	}

	return 0;
}