#define _GNU_SOURCE
#include <arpa/inet.h>
#include <assert.h>
#include <err.h>
#include <errno.h>
#include <inttypes.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/types.h>
#include <unistd.h>

#define EM_PROTOCOL_VERSION 0

#define EMT_SPEAK_VERSION 0
#define EMT_STATUS_REQUEST 1
#define EMT_STATUS 2
#define EMT_MSGID_REQUEST 3
#define EMT_MSGID 4

#define EMS_AVAILABLE 0
#define EMS_UNAVAILABLE 1

bool running = true;

int packet_socket;

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] = {'n', 'o', 'r', 't', 't', 'i'};
unsigned char own_nick_length = 6;

bool own_message_queued = false;
unsigned char own_message_destination_mac[6];
unsigned char own_message[1500 - 2 - 2 - 2];
size_t own_message_length = 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
	if (msgid_cache_fill < index) {
		msgid_cache_fill++;
	}

	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;
}

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();

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

void send_frame(const unsigned char *frame, size_t frame_length) {
	errno = 0;
	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 0
		// TODO: Make it random instead
		msgid_cache[cache_index].next_receive = 0;
		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 read_command(void) {
	int cmd = getchar();
	if (cmd == EOF) {
		err(1, "getchar");
	}

	if (cmd == 'q') {
		running = false;
	} else if (cmd == 's') {
		send_status_request(broadcast_mac);
	} else if (cmd == 'i') {
		send_msgid_request(broadcast_mac);
	} 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) {
		// 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[256];
	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);

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

	if (status == EMS_UNAVAILABLE) {
		errno = 0;
		if (printf("(unavailable) ") == -1) {
			err(1, "printf");
		}
	}

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

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

	errno = 0;
	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);

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

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

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

	errno = 0;
	ssize_t res = recv(packet_socket, frame, sizeof(frame), 0);
	if (res == -1) {
		errx(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;

		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) {
		errno = 0;
		int ready = poll(pollfds, sizeof(pollfds) / sizeof(*pollfds), -1);
		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];

	// Create a packet socket
	errno = 0;
	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);
	errno = 0;
	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;
	errno = 0;
	if (bind(packet_socket, (const struct sockaddr*)&sll, sizeof(sll)) == -1) {
		err(1, "bind");
	}

	// Get our own MAC
	errno = 0;
	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);

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

	// 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)
	errno = 0;
	if (close(packet_socket) == -1) {
		err(1, "close");
	}

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

	return 0;
}