#include <stdlib.h>
#include <stdbool.h>
#include <stddef.h>
#include <string.h>
#include <assert.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/ip_icmp.h>
#include <arpa/inet.h>
#include "icmp_client.h"
#include "icmp_server.h"
#include "util.h"


static void make_sockaddr_u32(struct sockaddr_in *dst, uint32_t addr) {
	memset(dst, 0, sizeof *dst);
	dst->sin_family = AF_INET;
	dst->sin_port = htons(0);
	dst->sin_addr.s_addr = addr;
}

__attribute__((unused))
static void make_sockaddr(struct sockaddr_in *dst, const char *ip_address) {
	make_sockaddr_u32(dst, inet_addr(ip_address));
}

static uint16_t compute_checksum(const void *buf_, size_t buflen) {
	const uint8_t *buf_bytes = (const uint8_t*)buf_;
	const uint16_t *buf = (const uint16_t*)buf_;

	uint32_t res = 0;
	for (size_t i = 0; i < buflen / 2; i++) {
		res += buf[i];
		res = (res + (res >> 16)) & 0xffff;
	}

	if (buflen % 2 == 1) {
		res += (uint16_t)buf_bytes[buflen - 1];
		res = (res + (res >> 16)) & 0xffff;
	}

	return res;
}

// TODO: Maybe SOCK_{DGRAM,RAW} | SOCK_NONBLOCK to make non-blocking sockets;
// this allows non-blocking reads in the icmpd thread.
int icmp_client_open_socket(void) {
	return socket(PF_INET, SOCK_DGRAM, IPPROTO_ICMP);
}

int icmp_server_open_socket(void) {
	return socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
}

struct icmp_incoming icmp_client_communicate(
		int sock, uint32_t addr,
		int seqnum,
		const void *data, size_t length) {

	if (icmp_client_send(sock, addr, seqnum, data, length) < 0) {
		return (struct icmp_incoming){.data = NULL, .length = 0, .id = 0, .seqnum = 0, .source_addr = 0};
	}

	return icmp_client_receive(sock);
}

struct icmp_incoming icmp_server_receive(int sock) {
	static uint8_t buffer[ICMP_MAX_PAYLOAD_LENGTH];

	struct icmp_incoming errret =
			(struct icmp_incoming){.data = NULL, .length = 0, .id = 0, .seqnum = 0, .source_addr = 0};

	char buf[MAX_IP_PACKET_SIZE];

	struct iovec iov1;
	memset(&iov1, 0, sizeof iov1);
	iov1.iov_base = buf;
	iov1.iov_len = sizeof buf;

	struct msghdr recv_msghdr;
	memset(&recv_msghdr, 0, sizeof recv_msghdr);
	recv_msghdr.msg_name = NULL;
	recv_msghdr.msg_iov = &iov1;
	recv_msghdr.msg_iovlen = 1;

	// TODO: check whether the above stuff needs re-initialisation
	while (true) {
		recv_msghdr.msg_flags = 0;

		ssize_t nr = recvmsg(sock, &recv_msghdr, 0);
		if (nr < 0) return errret;

		// buf now contains received data starting at the IP header

		struct iphdr *hdr = (struct iphdr*)buf;
		int hdr_len = hdr->ihl * 4;
		uint32_t saddr = hdr->saddr;

		struct icmp_echo *msg = (struct icmp_echo*)(buf + hdr_len);
		int msg_len = nr - hdr_len;
		int payloadlen = msg_len - offsetof(struct icmp_echo, payload);

		if (msg->type == ICMP_ECHO) {
			memcpy(buffer, msg->payload, payloadlen);

			return (struct icmp_incoming){
				.data = buffer,
				.length = payloadlen,
				.id = ntohs(msg->id),
				.seqnum = ntohs(msg->seqnum),
				.source_addr = saddr
			};
		}
	}
}

struct icmp_incoming icmp_client_receive(int sock) {
	static uint8_t buffer[ICMP_MAX_PAYLOAD_LENGTH];

	struct icmp_incoming errret =
			(struct icmp_incoming){.data = NULL, .length = 0, .id = 0, .seqnum = 0, .source_addr = 0};

	struct msghdr replyhdr;
	memset(&replyhdr, 0, sizeof replyhdr);

	struct icmp_echo msg;

	struct iovec iov1;
	memset(&iov1, 0, sizeof iov1);
	iov1.iov_base = &msg;
	iov1.iov_len = sizeof msg;

	struct sockaddr_in addr;

	replyhdr.msg_name = &addr;
	replyhdr.msg_namelen = sizeof addr;
	replyhdr.msg_iov = &iov1;
	replyhdr.msg_iovlen = 1;

	ssize_t nr = recvmsg(sock, &replyhdr, 0);
	if (nr < 0) {
		return errret;
	}

	size_t payloadlen = nr - ICMP_PAYLOAD_OFFSET;
	assert(payloadlen <= MAX_DATAGRAM_SIZE);

	memcpy(buffer, msg.payload, payloadlen);

	return (struct icmp_incoming){
		.data = buffer,
		.length = payloadlen,
		.id = ntohs(msg.id),
		.seqnum = ntohs(msg.seqnum),
		.source_addr = addr.sin_addr.s_addr
	};
}

static int icmp_server_send(
			int sock, uint32_t addr_u32,
			int id, int seqnum,
			const void *data_, size_t length,
			int type) {

	const uint8_t *data = (const uint8_t*)data_;

	assert(length <= ICMP_MAX_PAYLOAD_LENGTH);

	struct sockaddr_in addr;
	make_sockaddr_u32(&addr, addr_u32);

	struct icmp_echo msg;
	memset(&msg, 0, sizeof msg);
	msg.type = type;
	msg.code = 0;
	msg.id = htons(id);
	msg.seqnum = htons(seqnum);
	msg.checksum = 0;

	memcpy(msg.payload, data, length);

	size_t total_length = ICMP_PAYLOAD_OFFSET + length;

	msg.checksum = ~compute_checksum(&msg, total_length);

	int ret = sendto(sock, &msg, ICMP_PAYLOAD_OFFSET + length, 0, (struct sockaddr*)&addr, sizeof addr);
	if (ret < 0) return -1;

	return 0;
}

int icmp_server_send_reply(
			int sock, uint32_t addr_u32,
			int id, int seqnum,
			const void *data, size_t length) {

	return icmp_server_send(sock, addr_u32, id, seqnum, data, length, ICMP_ECHOREPLY);
}

int icmp_server_send_echo(
			int sock, uint32_t addr_u32,
			int id, int seqnum,
			const void *data, size_t length) {

	return icmp_server_send(sock, addr_u32, id, seqnum, data, length, ICMP_ECHO);
}

int icmp_client_send(
		int sock, uint32_t addr_u32,
		int seqnum,
		const void *data_, size_t length) {

	const uint8_t *data = (const uint8_t*)data_;

	assert(length <= ICMP_MAX_PAYLOAD_LENGTH);

	struct sockaddr_in addr;
	make_sockaddr_u32(&addr, addr_u32);

	struct icmp_echo msg;
	memset(&msg, 0, sizeof msg);
	msg.type = ICMP_ECHO;
	msg.code = 0;
	msg.id = 0;  // overwritten by kernel
	msg.seqnum = htons(seqnum);

	if (data == NULL) assert(length == 0);
	else memcpy(msg.payload, data, length);

	if (sendto(sock, &msg, ICMP_PAYLOAD_OFFSET + length, 0, (struct sockaddr*)&addr, sizeof addr) < 0) {
		return -1;
	}

	return 0;
}