#define _POSIX_C_SOURCE 200809L

#include <sys/resource.h>
#include <sys/stat.h>

#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#define PRESSURE_CHUNK (8U * 1024U * 1024U)

static void
usage(void)
{
	fprintf(stderr,
	    "usage:\n"
	    "  g7_probe empty FILE\n"
	    "  g7_probe eof FILE OFFSET\n"
	    "  g7_probe range SOURCE TARGET OFFSET LENGTH\n"
	    "  g7_probe random SOURCE TARGET SEED OPERATIONS BLOCK_SIZE\n"
	    "  g7_probe pressure GO_FILE STOP_FILE REQUEST_MIB\n");
}

static int
parse_u64(const char *value, uint64_t *result)
{
	char *end;
	unsigned long long parsed;

	errno = 0;
	parsed = strtoull(value, &end, 0);
	if (errno != 0 || *value == '\0' || *end != '\0')
		return (-1);
	*result = parsed;
	return (0);
}

static int
read_exact_at(int descriptor, unsigned char *buffer, size_t length,
    off_t offset)
{
	size_t total;
	ssize_t amount;

	total = 0;
	while (total < length) {
		amount = pread(descriptor, buffer + total, length - total,
		    offset + (off_t)total);
		if (amount < 0 && errno == EINTR)
			continue;
		if (amount <= 0)
			return (-1);
		total += (size_t)amount;
	}
	return (0);
}

static uint64_t
digest_bytes(uint64_t digest, const unsigned char *buffer, size_t length)
{
	size_t index;

	for (index = 0; index < length; index++) {
		digest ^= buffer[index];
		digest *= UINT64_C(1099511628211);
	}
	return (digest);
}

static int
empty_probe(const char *path)
{
	struct stat status;
	unsigned char byte;
	ssize_t amount;
	int descriptor;

	descriptor = open(path, O_RDONLY);
	if (descriptor < 0)
		return (1);
	if (fstat(descriptor, &status) != 0 || status.st_size != 0) {
		close(descriptor);
		return (1);
	}
	amount = read(descriptor, &byte, sizeof(byte));
	if (amount != 0 || lseek(descriptor, 0, SEEK_SET) != 0 ||
	    lseek(descriptor, 0, SEEK_END) != 0) {
		close(descriptor);
		return (1);
	}
	close(descriptor);
	printf("size=0 read_bytes=0 seek_set=0 seek_end=0 exit=0\n");
	return (0);
}

static int
eof_probe(const char *path, const char *offset_value)
{
	unsigned char byte;
	uint64_t offset;
	ssize_t amount;
	int descriptor;

	if (parse_u64(offset_value, &offset) != 0 || offset > INT64_MAX)
		return (1);
	descriptor = open(path, O_RDONLY);
	if (descriptor < 0)
		return (1);
	amount = pread(descriptor, &byte, sizeof(byte), (off_t)offset);
	close(descriptor);
	if (amount != 0)
		return (1);
	printf("offset=%" PRIu64 " read_bytes=0 errno=0 exit=0\n", offset);
	return (0);
}

static int
range_probe(const char *source_path, const char *target_path,
    const char *offset_value, const char *length_value)
{
	unsigned char *source_buffer, *target_buffer;
	uint64_t digest, length, offset;
	int source, target, result;

	if (parse_u64(offset_value, &offset) != 0 ||
	    parse_u64(length_value, &length) != 0 || length == 0 ||
	    length > SIZE_MAX || offset > INT64_MAX)
		return (1);
	source_buffer = malloc((size_t)length);
	target_buffer = malloc((size_t)length);
	if (source_buffer == NULL || target_buffer == NULL) {
		free(source_buffer);
		free(target_buffer);
		return (1);
	}
	source = open(source_path, O_RDONLY);
	target = open(target_path, O_RDONLY);
	result = 1;
	if (source >= 0 && target >= 0 &&
	    read_exact_at(source, source_buffer, (size_t)length,
	    (off_t)offset) == 0 &&
	    read_exact_at(target, target_buffer, (size_t)length,
	    (off_t)offset) == 0 &&
	    memcmp(source_buffer, target_buffer, (size_t)length) == 0) {
		digest = digest_bytes(UINT64_C(1469598103934665603),
		    source_buffer, (size_t)length);
		printf("offset=%" PRIu64 " length=%" PRIu64
		    " mismatches=0 digest=%016" PRIx64 " exit=0\n",
		    offset, length, digest);
		result = 0;
	}
	if (source >= 0)
		close(source);
	if (target >= 0)
		close(target);
	free(source_buffer);
	free(target_buffer);
	return (result);
}

static uint64_t
xorshift64star(uint64_t *state)
{
	uint64_t value;

	value = *state;
	value ^= value >> 12;
	value ^= value << 25;
	value ^= value >> 27;
	*state = value;
	return (value * UINT64_C(2685821657736338717));
}

static double
elapsed_seconds(const struct timespec *start, const struct timespec *end)
{
	return ((double)(end->tv_sec - start->tv_sec) +
	    (double)(end->tv_nsec - start->tv_nsec) / 1000000000.0);
}

static int
random_probe(int argc, char **argv)
{
	struct stat source_status, target_status;
	struct timespec start, end;
	unsigned char *source_buffer, *target_buffer;
	uint64_t block_size, digest, index, input_seed, operations, offset, seed;
	uint64_t slots;
	double elapsed;
	int source, target, result;

	if (argc != 7 || parse_u64(argv[4], &seed) != 0 || seed == 0 ||
	    parse_u64(argv[5], &operations) != 0 || operations == 0 ||
	    parse_u64(argv[6], &block_size) != 0 || block_size == 0 ||
	    block_size > SIZE_MAX)
		return (1);
	input_seed = seed;
	source = open(argv[2], O_RDONLY);
	target = open(argv[3], O_RDONLY);
	if (source < 0 || target < 0 || fstat(source, &source_status) != 0 ||
	    fstat(target, &target_status) != 0 || source_status.st_size <= 0 ||
	    source_status.st_size != target_status.st_size ||
	    (uint64_t)source_status.st_size < block_size) {
		if (source >= 0)
			close(source);
		if (target >= 0)
			close(target);
		return (1);
	}
	source_buffer = malloc((size_t)block_size);
	target_buffer = malloc((size_t)block_size);
	if (source_buffer == NULL || target_buffer == NULL) {
		close(source);
		close(target);
		free(source_buffer);
		free(target_buffer);
		return (1);
	}
	slots = (uint64_t)source_status.st_size / block_size;
	digest = UINT64_C(1469598103934665603);
	result = 1;
	clock_gettime(CLOCK_MONOTONIC, &start);
	for (index = 0; index < operations; index++) {
		offset = (xorshift64star(&seed) % slots) * block_size;
		if (read_exact_at(source, source_buffer, (size_t)block_size,
		    (off_t)offset) != 0 ||
		    read_exact_at(target, target_buffer, (size_t)block_size,
		    (off_t)offset) != 0 ||
		    memcmp(source_buffer, target_buffer, (size_t)block_size) != 0)
			goto out;
		digest = digest_bytes(digest, target_buffer, (size_t)block_size);
	}
	clock_gettime(CLOCK_MONOTONIC, &end);
	elapsed = elapsed_seconds(&start, &end);
	printf("seed=0x%016" PRIx64 " operations=%" PRIu64
	    " block_size=%" PRIu64
	    " bytes=%" PRIu64 " mismatches=0 digest=%016" PRIx64
	    " elapsed_seconds=%.6f iops=%.2f mean_us=%.2f exit=0\n",
	    input_seed, operations, block_size, operations * block_size, digest, elapsed,
	    (double)operations / elapsed, elapsed * 1000000.0 / operations);
	result = 0;
out:
	close(source);
	close(target);
	free(source_buffer);
	free(target_buffer);
	return (result);
}

static int
wait_for_path(const char *path, unsigned int attempts)
{
	struct timespec delay;
	unsigned int attempt;

	delay.tv_sec = 0;
	delay.tv_nsec = 100000000;
	for (attempt = 0; attempt < attempts; attempt++) {
		if (access(path, F_OK) == 0)
			return (0);
		nanosleep(&delay, NULL);
	}
	return (-1);
}

static int
pressure_probe(const char *go_path, const char *stop_path,
    const char *request_value)
{
	struct rusage usage;
	unsigned char *recovery;
	unsigned char **chunks;
	uint64_t allocated, index, page_size, request_mib, requested;
	int failure_errno, recovery_ok;

	if (parse_u64(request_value, &request_mib) != 0 || request_mib == 0 ||
	    request_mib > 4096)
		return (1);
	requested = request_mib * 1024 * 1024;
	chunks = calloc((size_t)(requested / PRESSURE_CHUNK + 1), sizeof(*chunks));
	if (chunks == NULL)
		return (1);
	printf("state=waiting pid=%ld requested_bytes=%" PRIu64 "\n",
	    (long)getpid(), requested);
	fflush(stdout);
	if (wait_for_path(go_path, 600) != 0) {
		free(chunks);
		return (1);
	}
	page_size = (uint64_t)sysconf(_SC_PAGESIZE);
	allocated = 0;
	failure_errno = 0;
	for (index = 0; allocated < requested; index++) {
		size_t offset;

		errno = 0;
		chunks[index] = malloc(PRESSURE_CHUNK);
		if (chunks[index] == NULL) {
			failure_errno = errno;
			break;
		}
		for (offset = 0; offset < PRESSURE_CHUNK; offset += page_size)
			chunks[index][offset] = (unsigned char)(index + offset);
		allocated += PRESSURE_CHUNK;
	}
	getrusage(RUSAGE_SELF, &usage);
	printf("state=holding pid=%ld allocated_bytes=%" PRIu64
	    " allocation_failure=%s failure_errno=%d maxrss=%ld\n",
	    (long)getpid(), allocated,
	    failure_errno == ENOMEM ? "ENOMEM" : "none", failure_errno,
	    usage.ru_maxrss);
	fflush(stdout);
	if (wait_for_path(stop_path, 3000) != 0) {
		for (index = 0; index < allocated / PRESSURE_CHUNK; index++)
			free(chunks[index]);
		free(chunks);
		return (1);
	}
	for (index = 0; index < allocated / PRESSURE_CHUNK; index++)
		free(chunks[index]);
	free(chunks);
	recovery = malloc(1024 * 1024);
	recovery_ok = recovery != NULL;
	if (recovery != NULL) {
		recovery[0] = 1;
		free(recovery);
	}
	printf("state=released pid=%ld released_bytes=%" PRIu64
	    " recovery=%s exit=%d\n",
	    (long)getpid(), allocated, recovery_ok ? "ok" : "failed",
	    failure_errno == ENOMEM && recovery_ok ? 0 : 1);
	return (failure_errno == ENOMEM && recovery_ok ? 0 : 1);
}

int
main(int argc, char **argv)
{
	if (argc == 3 && strcmp(argv[1], "empty") == 0)
		return (empty_probe(argv[2]));
	if (argc == 4 && strcmp(argv[1], "eof") == 0)
		return (eof_probe(argv[2], argv[3]));
	if (argc == 6 && strcmp(argv[1], "range") == 0)
		return (range_probe(argv[2], argv[3], argv[4], argv[5]));
	if (argc == 7 && strcmp(argv[1], "random") == 0)
		return (random_probe(argc, argv));
	if (argc == 5 && strcmp(argv[1], "pressure") == 0)
		return (pressure_probe(argv[2], argv[3], argv[4]));
	usage();
	return (64);
}