#include <gtest/gtest.h>
#include <syscall.h>
#include <libscap/scap.h>
#include <errno.h>
#include <fcntl.h>

/* We are supposing that if we overcome this threshold, all buffers are full.
 * Probably this threshold is too low, but it depends on the machine's workload.
 * We are running in CI so it is better to be conservative even if tests becomes
 * not so reliable...
 */
#define MAX_ITERATIONS 300

/* Number of events we want to assert */
#define EVENTS_TO_ASSERT 18

void check_event_is_not_overwritten(scap_t *h) {
	/* Start the capture */
	ASSERT_EQ(scap_start_capture(h), SCAP_SUCCESS)
	        << "unable to start the capture: " << scap_getlasterr(h) << std::endl;

	/* When the number of events is fixed for `MAX_ITERATIONS` we consider all the buffers full,
	 * this is just an approximation */
	scap_stats stats = {};
	uint64_t last_num_events = 0;
	uint16_t iterations = 0;

	while(iterations < MAX_ITERATIONS || stats.n_drops == 0) {
		ASSERT_EQ(scap_get_stats(h, &stats), SCAP_SUCCESS)
		        << "unable to get stats: " << scap_getlasterr(h) << std::endl;
		if(last_num_events == (stats.n_evts - stats.n_drops)) {
			iterations++;
		} else {
			iterations = 0;
			last_num_events = (stats.n_evts - stats.n_drops);
		}
	}

	/* Stop the capture */
	ASSERT_EQ(scap_stop_capture(h), SCAP_SUCCESS)
	        << "unable to stop the capture: " << scap_getlasterr(h) << std::endl;

	/* The idea here is to check if an event is overwritten while we still have a pointer to it.
	 * Again this is only an approximation, we don't know if new events will be written in the
	 * buffer under test...
	 *
	 * We call `scap_next` keeping the pointer to the event.
	 * An event pointer becomes invalid when we call another `scap_next`, but until that moment it
	 * should be valid!
	 */
	scap_evt *evt = NULL;
	uint16_t buffer_id;
	uint32_t flags;

	/* The first 'scap_next` could return a `SCAP_TIMEOUT` according to the chosen `buffer_mode` so
	 * we ignore it. */
	scap_next(h, &evt, &buffer_id, &flags);

	ASSERT_EQ(scap_next(h, &evt, &buffer_id, &flags), SCAP_SUCCESS)
	        << "unable to get an event with `scap_next`: " << scap_getlasterr(h) << std::endl;

	last_num_events = 0;
	iterations = 0;

	/* We save some event info to check if they are still valid after some new events */
	uint64_t prev_ts = scap_event_get_ts(evt);
	uint64_t prev_tid = scap_event_get_tid(evt);
	uint32_t prev_len = scap_event_getlen(evt);
	uint16_t prev_type = scap_event_get_type(evt);
	uint32_t prev_nparams = scap_event_get_nparams(evt);

	/* Start again the capture */
	ASSERT_EQ(scap_start_capture(h), SCAP_SUCCESS)
	        << "unable to restart the capture: " << scap_getlasterr(h) << std::endl;

	/* We use the same approximation as before */
	while(iterations < MAX_ITERATIONS) {
		ASSERT_EQ(scap_get_stats(h, &stats), SCAP_SUCCESS)
		        << "unable to get stats: " << scap_getlasterr(h) << std::endl;
		if(last_num_events == (stats.n_evts - stats.n_drops)) {
			iterations++;
		} else {
			iterations = 0;
			last_num_events = (stats.n_evts - stats.n_drops);
		}
	}

	/* We check if the previously collected event is still valid */
	ASSERT_EQ(prev_ts, scap_event_get_ts(evt)) << "different timestamp" << std::endl;
	ASSERT_EQ(prev_tid, scap_event_get_tid(evt)) << "different thread id" << std::endl;
	ASSERT_EQ(prev_len, scap_event_getlen(evt)) << "different event len" << std::endl;
	ASSERT_EQ(prev_type, scap_event_get_type(evt)) << "different event type" << std::endl;
	ASSERT_EQ(prev_nparams, scap_event_get_nparams(evt)) << "different num params" << std::endl;
}

#if defined(__NR_close) && defined(__NR_openat) && defined(__NR_listen) &&           \
        defined(__NR_accept4) && defined(__NR_getegid) && defined(__NR_getgid) &&    \
        defined(__NR_geteuid) && defined(__NR_getuid) && defined(__NR_bind) &&       \
        defined(__NR_connect) && defined(__NR_sendto) && defined(__NR_getsockopt) && \
        defined(__NR_recvmsg) && defined(__NR_recvfrom) && defined(__NR_socket) &&   \
        defined(__NR_socketpair)

void check_event_order(scap_t *h) {
	uint32_t events_to_assert[EVENTS_TO_ASSERT] = {PPME_SYSCALL_CLOSE_X,
	                                               PPME_SYSCALL_OPENAT_2_E,
	                                               PPME_SYSCALL_OPENAT_2_X,
	                                               PPME_SOCKET_LISTEN_X,
	                                               PPME_SOCKET_ACCEPT4_6_X,
	                                               PPME_SYSCALL_GETEGID_X,
	                                               PPME_SYSCALL_GETGID_X,
	                                               PPME_SYSCALL_GETEUID_X,
	                                               PPME_SYSCALL_GETUID_X,
	                                               PPME_SOCKET_BIND_X,
	                                               PPME_SOCKET_CONNECT_E,
	                                               PPME_SOCKET_CONNECT_X,
	                                               PPME_SOCKET_SENDTO_X,
	                                               PPME_SOCKET_GETSOCKOPT_X,
	                                               PPME_SOCKET_RECVMSG_X,
	                                               PPME_SOCKET_RECVFROM_X,
	                                               PPME_SOCKET_SOCKET_X,
	                                               PPME_SOCKET_SOCKETPAIR_X};

	/* Start the capture */
	ASSERT_EQ(scap_start_capture(h), SCAP_SUCCESS)
	        << "unable to start the capture: " << scap_getlasterr(h) << std::endl;

	/* 1. Generate a `close` exit event. */
	syscall(__NR_close, -1);

	/* 2. Generate an `openat` event pair */
	syscall(__NR_openat, 0, "/**mock_path**/", 0, 0);

	/* 3. Generate a `listen` exit event. */
	syscall(__NR_listen, -1, -1);

	/* 4. Generate an `accept4` exit event. */
	syscall(__NR_accept4, -1, NULL, NULL, 0);

	/* 5. Generate a `getegid` exit event. */
	syscall(__NR_getegid);

	/* 6. Generate a `getgid` exit event. */
	syscall(__NR_getgid);

	/* 7. Generate a `geteuid` exit event. */
	syscall(__NR_geteuid);

	/* 8. Generate a `getuid` exit event. */
	syscall(__NR_getuid);

	/* 9. Generate a `bind` exit event. */
	syscall(__NR_bind, -1, NULL, 0);

	/* 10. Generate a `connect` event pair */
	syscall(__NR_connect, -1, NULL, 0);

	/* 11. Generate a `sendto` exit event. */
	syscall(__NR_sendto, -1, NULL, 0, 0, NULL, 0);

	/* 12. Generate a `getsockopt` exit event. */
	syscall(__NR_getsockopt, -1, 0, 0, NULL, NULL);

	/* 13. Generate a `recvmsg` exit event. */
	syscall(__NR_recvmsg, -1, NULL, 0);

	/* 14. Generate a `recvmsg` exit event. */
	syscall(__NR_recvfrom, -1, NULL, 0, 0, NULL, 0);

	/* 15. Generate a `socket` exit event. */
	syscall(__NR_socket, 0, 0, 0);

	/* 16. Generate a `socketpair` exit event. */
	syscall(__NR_socketpair, 0, 0, 0, 0);

	/* Stop the capture */
	ASSERT_EQ(scap_stop_capture(h), SCAP_SUCCESS)
	        << "unable to stop the capture: " << scap_getlasterr(h) << std::endl;

	scap_evt *evt = NULL;
	uint16_t buffer_id = 0;
	uint32_t flags = 0;
	int ret = 0;
	uint64_t actual_pid = getpid();
	/* if we hit 5 consecutive timeouts it means that all buffers are empty (approximation) */
	uint16_t timeouts = 0;

	for(int i = 0; i < EVENTS_TO_ASSERT; i++) {
		while(true) {
			ret = scap_next(h, &evt, &buffer_id, &flags);
			if(ret == SCAP_SUCCESS) {
				timeouts = 0;
				if(scap_event_get_tid(evt) == actual_pid &&
				   scap_event_get_type(evt) == events_to_assert[i]) {
					/* We found our event */
					break;
				}
			} else if(ret == SCAP_TIMEOUT) {
				timeouts++;
				if(timeouts == 5) {
					FAIL() << "we didn't find event '" << events_to_assert[i] << "' at position '"
					       << i << "'" << std::endl;
				}
			}
		}
	}
}

#else

void check_event_order(scap_t *h) {
	GTEST_SKIP() << "Some syscalls required by the test are not defined" << std::endl;
}
#endif

/* Right now this is used only by the modern bpf
 * This is extracted from `libbpf_num_possible_cpus()`.
 * We avoid to include libbpf just for this helper.
 */
static int parse_cpu_mask_str(const char *s, bool **mask, int *mask_sz) {
	int err = 0, n, len, start, end = -1;
	bool *tmp;

	*mask = NULL;
	*mask_sz = 0;

	/* Each sub string separated by ',' has format \d+-\d+ or \d+ */
	while(*s) {
		if(*s == ',' || *s == '\n') {
			s++;
			continue;
		}
		n = sscanf(s, "%d%n-%d%n", &start, &len, &end, &len);
		if(n <= 0 || n > 2) {
			fprintf(stderr, "Failed to get CPU range %s: %d\n", s, n);
			err = -EINVAL;
			goto cleanup;
		} else if(n == 1) {
			end = start;
		}
		if(start < 0 || start > end) {
			fprintf(stderr, "Invalid CPU range [%d,%d] in %s\n", start, end, s);
			err = -EINVAL;
			goto cleanup;
		}
		tmp = (bool *)realloc(*mask, end + 1);
		if(!tmp) {
			err = -ENOMEM;
			goto cleanup;
		}
		*mask = tmp;
		memset(tmp + *mask_sz, 0, start - *mask_sz);
		memset(tmp + start, 1, end - start + 1);
		*mask_sz = end + 1;
		s += len;
	}
	if(!*mask_sz) {
		fprintf(stderr, "Empty CPU range\n");
		return -EINVAL;
	}
	return 0;
cleanup:
	free(*mask);
	*mask = NULL;
	return err;
}

static int parse_cpu_mask_file(const char *fcpu, bool **mask, int *mask_sz) {
	int fd, err = 0, len;
	char buf[128];

	fd = open(fcpu, O_RDONLY | O_CLOEXEC);
	if(fd < 0) {
		err = -errno;
		fprintf(stderr, "Failed to open cpu mask file %s: %d\n", fcpu, err);
		return err;
	}
	len = read(fd, buf, sizeof(buf));
	close(fd);
	if(len <= 0) {
		err = len ? -errno : -EINVAL;
		fprintf(stderr, "Failed to read cpu mask from %s: %d\n", fcpu, err);

		return err;
	}
	if(len >= sizeof(buf)) {
		fprintf(stderr, "CPU mask is too big in file %s\n", fcpu);

		return -E2BIG;
	}
	buf[len] = '\0';

	return parse_cpu_mask_str(buf, mask, mask_sz);
}

int num_possible_cpus(void) {
	const char *fcpu = "/sys/devices/system/cpu/possible";
	int cpus = 0;
	int n = 0;

	/* array of bools for each CPU */
	bool *mask = nullptr;
	int err = parse_cpu_mask_file(fcpu, &mask, &n);
	if(err)
		return -1;

	for(int i = 0; i < n; i++) {
		if(mask[i])
			cpus++;
	}

	if(mask) {
		free(mask);
	}
	return cpus;
}
