#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <error.h>
#include <poll.h>
#include <arpa/inet.h>
#include <linux/icmp.h>
#include <linux/icmpv6.h>
#include <sys/uio.h>
#include <sys/socket.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <stdbool.h>
#include <pthread.h>
#include <signal.h>
int done = 0;
struct timespec timeout;
struct timespec seq[65536];
int sock4 = -1;
int sock6 = -1;
void * reader (void * unused __attribute__((unused))) {
	struct pollfd fds[] = {
		{
			.fd = sock4,
			.events = POLLIN
		},
		{
			.fd = sock6,
			.events = POLLIN
		}
	};
	int waited = 0;
	while (true) {
		int r = ppoll(fds, 2, &timeout, NULL);
		if (r == -1) {
			perror("ppoll");
			return NULL;
		}
		if (r == 0) {
			if (waited)
				return NULL;
			if (done) {
				fprintf(stderr, "Done sending. Waiting for any new packets ...\n");
				waited = 1;
			}
			continue;
		}
		struct timespec timestamp;
		if (clock_gettime(CLOCK_MONOTONIC, &timestamp) == -1) {
			perror("clock_gettime r");
			return NULL;
		}
		char buf[32];
		uint32_t secr = htonl(timestamp.tv_sec);
		uint32_t nsecr = htonl(timestamp.tv_nsec);
		memcpy(buf+8, &secr, 4);
		memcpy(buf+12, &nsecr, 4);
		if (fds[0].revents & POLLIN) {
			struct sockaddr_in from;
			char recvbuf[512];
			struct icmphdr hdr;
			socklen_t socklen = sizeof from;
			int recvret = recvfrom(sock4, recvbuf, sizeof recvbuf, 0, (struct sockaddr *) &from, &socklen);
			if (recvret == -1) {
				perror("recvfrom 4");
				return NULL;
			}
			memcpy(&hdr, recvbuf + (recvbuf[0] & 0x0f)*4, sizeof hdr);
			if (recvret < (int) sizeof hdr) {
				fprintf(stderr, "4: too small packet received!\n");
			}
			if (hdr.un.echo.id != (getpid() & 0xffff)) {
				fprintf(stderr, "4: received not own packet ID=%d. ID should be %d. skipping.\n", hdr.un.echo.id, getpid() & 0xffff);
				continue;
			}
			uint32_t secs = htonl(seq[hdr.un.echo.sequence].tv_sec);
			uint32_t nsecs = htonl(seq[hdr.un.echo.sequence].tv_nsec);
			uint16_t ffff = 0xffff;
			memcpy(buf, &secs, 4);
			memcpy(buf+4, &nsecs, 4);
			memcpy(buf+16+8+2, &ffff, 2);
			memcpy(buf+16+8+2+2, &from.sin_addr.s_addr, 4);
			if (fwrite(buf, 32, 1, stdout) == 0) {
				fprintf(stderr, "fwrite failed 4!\n");
				return NULL;
			}
			waited = 0;
		}
		if (fds[1].revents & POLLIN) {
			struct sockaddr_in6 from;
			struct icmp6hdr hdr;
			socklen_t socklen = sizeof from;
			if (recvfrom(sock6, &hdr, sizeof hdr, 0, (struct sockaddr *) &from, &socklen) == -1) {
				perror("recvfrom 6");
				return NULL;
			}
			if (hdr.icmp6_dataun.u_echo.identifier != (getpid() & 0xffff)) {
				fprintf(stderr, "received not own packet ID=%d. ID should be %d. skipping.\n", hdr.icmp6_dataun.u_echo.identifier, getpid() & 0xffff);
				continue;
			}
			uint32_t secs = htonl(seq[hdr.icmp6_dataun.u_echo.sequence].tv_sec);
			uint32_t nsecs = htonl(seq[hdr.icmp6_dataun.u_echo.sequence].tv_nsec);
			memcpy(buf, &secs, 4);
			memcpy(buf+4, &nsecs, 4);
			memcpy(buf+16, from.sin6_addr.s6_addr, 16);
			if (fwrite(buf, 32, 1, stdout) == 0) {
				fprintf(stderr, "fwrite failed 6!\n");
				return NULL;
			}
			waited = 0;
		}
		fflush(stdout);
	}
}
int main (int argc, char ** argv) {
	if (argc < 2) {
		fprintf(stderr, "SEND_DELAY=1000 PKTCNT=1000 WAIT_AT_END=1000 %s 2a01:261:e44:1300::1 ...\nSEND_DELAY is in microseconds\nWAIT_AT_END is in milliseconds\nPKTCNT=-1 for infinite profiling\n", argv[0]);
		return 1;
	}
	int send_delay = 1000;
	if (getenv("SEND_DELAY"))
		send_delay = atoi(getenv("SEND_DELAY"));
	int pktcnt = 1000;
	if (getenv("PKTCNT"))
		pktcnt = atoi(getenv("PKTCNT"));
	int wait_at_end = 1000;
	if (getenv("WAIT_AT_END"))
		wait_at_end = atoi(getenv("WAIT_AT_END"));
	timeout.tv_sec = wait_at_end/1000;
	timeout.tv_nsec = (wait_at_end%1000)*10000000;
	sock6 = socket(AF_INET6, SOCK_RAW | SOCK_CLOEXEC, IPPROTO_ICMPV6);
	if (sock6 == -1) {
		perror("socket 6");
		return 1;
	}
	sock4 = socket(AF_INET, SOCK_RAW | SOCK_CLOEXEC, IPPROTO_ICMP);
	if (sock4 == -1) {
		perror("socket 4");
		return 1;
	}
	struct icmp6hdr hdr6 = {
		.icmp6_type = 128,
		.icmp6_code = 0,
		.icmp6_dataun.u_echo.identifier = (getpid() & 0xffff)
	};
	struct icmphdr hdr4 = {
		.type = 8,
		.code = 0,
		.un.echo.id = (getpid() & 0xffff)
	};
	uint16_t seqnr = 0;
	struct iovec iov6[2] = {
		{
			.iov_base = &hdr6,
			.iov_len = sizeof hdr6-sizeof seqnr
		},
		{
			.iov_base = &seqnr,
			.iov_len = sizeof seqnr
		}
	};
	struct iovec iov4[2] = {
		{
			.iov_base = &hdr4,
			.iov_len = sizeof hdr4-sizeof seqnr
		},
		{
			.iov_base = &seqnr,
			.iov_len = sizeof seqnr
		}
	};
	struct sockaddr_in targets4[argc-1];
	struct sockaddr_in6 targets6[argc-1];
	int t4c = 0;
	int t6c = 0;
	memset(targets4, 0, sizeof targets4);
	memset(targets6, 0, sizeof targets6);
	for (int i = 0; i < argc-1; i++) {
		if (inet_pton(AF_INET6, argv[1+i], targets6[t6c].sin6_addr.s6_addr) == 1) {
			targets6[t6c++].sin6_family = AF_INET6;
			continue;
		}
		if (inet_pton(AF_INET, argv[1+i], &(targets4[t4c].sin_addr.s_addr)) == 1) {
			targets4[t4c++].sin_family = AF_INET;
			continue;
		}
		fprintf(stderr, "address %s is neither IPv4 nor IPv6!\n", argv[1+i]);
		return 1;
	}
	struct mmsghdr mm4[t4c];
	memset(mm4, 0, sizeof mm4);
	struct mmsghdr mm6[t6c];
	memset(mm6, 0, sizeof mm6);
	for (int i = 0; i < t4c; i++) {
		mm4[i].msg_hdr.msg_name = &(targets4[i]);
		mm4[i].msg_hdr.msg_namelen = sizeof(struct sockaddr_in);
		mm4[i].msg_hdr.msg_iov = iov4;
		mm4[i].msg_hdr.msg_iovlen = 2;
	}
	for (int i = 0; i < t6c; i++) {
		mm6[i].msg_hdr.msg_name = &(targets6[i]);
		mm6[i].msg_hdr.msg_namelen = sizeof(struct sockaddr_in6);
		mm6[i].msg_hdr.msg_iov = iov6;
		mm6[i].msg_hdr.msg_iovlen = 2;
	}
	pthread_t readerthread;
	pthread_create(&readerthread, NULL, reader, NULL);
	while (pktcnt != 0) {
		if (pktcnt > 0)
			pktcnt--;
		if (sendmmsg(sock6, mm6, t6c, 0) == -1) {
			perror("sendmmsg 6");
			return 1;
		}
		hdr4.checksum = htons(ntohs(seqnr)+8*256+ntohs(hdr4.un.echo.id)) ^ 0xffff;
		if (sendmmsg(sock4, mm4, t4c, 0) == -1) {
			perror("sendmmsg 4");
			return 1;
		}
		if (clock_gettime(CLOCK_MONOTONIC, &seq[seqnr]) == -1) {
			perror("clock_gettime");
			return 1;
		}
		seqnr++;
		usleep(send_delay);
	}
	done = 1;
	pthread_join(readerthread, NULL);
	close(sock4);
	close(sock6);
	return 0;
}