#define _GNU_SOURCE

#include <dirent.h>
#include <errno.h>
#include <getopt.h>
#include <inttypes.h>
#include <limits.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>

#include <doca_buf.h>
#include <doca_buf_inventory.h>
#include <doca_compress.h>
#include <doca_ctx.h>
#include <doca_dev.h>
#include <doca_error.h>
#include <doca_mmap.h>
#include <doca_types.h>

#ifndef DOCA_WORKQ_RETRIEVE_FLAGS_NONE
#define DOCA_WORKQ_RETRIEVE_FLAGS_NONE 0
#endif

#define DEFAULT_OUTPUT "decompressed.out"
#define NUM_BUFS 2
#define WORKQ_DEPTH 8

struct one_result {
	uint64_t compressed_bytes;
	uint64_t decompressed_bytes;
	double elapsed_us;
	double bw_mibps;
	uint32_t crc;
	uint32_t adler;
};

static const char *errstr(doca_error_t err)
{
	return doca_get_error_string(err);
}

static void usage(const char *prog)
{
	fprintf(stderr,
		"Usage:\n"
		"  Single file:\n"
		"    %s -p <dpu-pci-bdf> -f <compressed-input> -s <dst-size> [-o output-file] [--no-write]\n"
		"\n"
		"  Directory mode:\n"
		"    %s -p <dpu-pci-bdf> -d <chunk-directory> -S <dst-size-per-chunk>\n"
		"\n"
		"Examples:\n"
		"  sudo ./%s -p 03:00.0 -f 1M.bin.compressed -s 1048576 --no-write\n"
		"  sudo ./%s -p 03:00.0 -d bins/1G -S 1M\n",
		prog, prog, prog, prog);
}

static uint64_t now_nsec(void)
{
	struct timespec ts;

#ifdef CLOCK_MONOTONIC_RAW
	if (clock_gettime(CLOCK_MONOTONIC_RAW, &ts) != 0) {
#else
	if (clock_gettime(CLOCK_MONOTONIC, &ts) != 0) {
#endif
		perror("clock_gettime");
		exit(EXIT_FAILURE);
	}

	return (uint64_t)ts.tv_sec * 1000000000ULL + (uint64_t)ts.tv_nsec;
}

static void check_doca(doca_error_t err, const char *what)
{
	if (err != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: %s: %s\n", what, errstr(err));
		exit(EXIT_FAILURE);
	}
}

static size_t parse_size_or_die(const char *s)
{
	char *tmp;
	char *end = NULL;
	uint64_t v;
	uint64_t mul = 1;
	size_t len;

	len = strlen(s);
	if (len == 0) {
		fprintf(stderr, "ERROR: invalid size: %s\n", s);
		exit(EXIT_FAILURE);
	}

	tmp = strdup(s);
	if (tmp == NULL) {
		fprintf(stderr, "ERROR: strdup failed\n");
		exit(EXIT_FAILURE);
	}

	char last = tmp[len - 1];

	if (last == 'K' || last == 'k') {
		mul = 1024ULL;
		tmp[len - 1] = '\0';
	} else if (last == 'M' || last == 'm') {
		mul = 1024ULL * 1024ULL;
		tmp[len - 1] = '\0';
	} else if (last == 'G' || last == 'g') {
		mul = 1024ULL * 1024ULL * 1024ULL;
		tmp[len - 1] = '\0';
	}

	errno = 0;
	v = strtoull(tmp, &end, 10);

	if (errno != 0 || end == tmp || *end != '\0' || v == 0) {
		fprintf(stderr, "ERROR: invalid size: %s\n", s);
		free(tmp);
		exit(EXIT_FAILURE);
	}

	free(tmp);

	if (v > SIZE_MAX / mul) {
		fprintf(stderr, "ERROR: size overflow: %s\n", s);
		exit(EXIT_FAILURE);
	}

	return (size_t)(v * mul);
}

static uint8_t *read_file_or_die(const char *path, size_t *len)
{
	FILE *f;
	struct stat st;
	uint8_t *buf;
	size_t nread;

	if (stat(path, &st) != 0) {
		fprintf(stderr, "ERROR: stat(%s): %s\n", path, strerror(errno));
		exit(EXIT_FAILURE);
	}

	if (!S_ISREG(st.st_mode)) {
		fprintf(stderr, "ERROR: not a regular file: %s\n", path);
		exit(EXIT_FAILURE);
	}

	if (st.st_size < 0) {
		fprintf(stderr, "ERROR: invalid file size: %s\n", path);
		exit(EXIT_FAILURE);
	}

	*len = (size_t)st.st_size;

	buf = malloc(*len == 0 ? 1 : *len);
	if (buf == NULL) {
		fprintf(stderr, "ERROR: malloc input failed\n");
		exit(EXIT_FAILURE);
	}

	f = fopen(path, "rb");
	if (f == NULL) {
		fprintf(stderr, "ERROR: fopen(%s): %s\n", path, strerror(errno));
		free(buf);
		exit(EXIT_FAILURE);
	}

	nread = fread(buf, 1, *len, f);
	if (nread != *len) {
		fprintf(stderr, "ERROR: fread(%s): expected %zu, got %zu\n",
			path, *len, nread);
		fclose(f);
		free(buf);
		exit(EXIT_FAILURE);
	}

	fclose(f);
	return buf;
}

static void write_file_or_die(const char *path, const void *buf, size_t len)
{
	FILE *f;
	size_t nwritten;

	f = fopen(path, "wb");
	if (f == NULL) {
		fprintf(stderr, "ERROR: fopen(%s): %s\n", path, strerror(errno));
		exit(EXIT_FAILURE);
	}

	nwritten = fwrite(buf, 1, len, f);
	if (nwritten != len) {
		fprintf(stderr, "ERROR: fwrite(%s): expected %zu, wrote %zu\n",
			path, len, nwritten);
		fclose(f);
		exit(EXIT_FAILURE);
	}

	fclose(f);
}

static int open_first_doca_device(const char *pci_addr, struct doca_dev **dev)
{
	struct doca_devinfo **dev_list = NULL;
	uint32_t nb_devs = 0;
	doca_error_t result;

	/*
	 * DOCA 2.0.2027:
	 *   doca_devinfo_list_create()
	 *   doca_devinfo_list_destroy()
	 *
	 * Keep -p for CLI compatibility, but open the first available local DOCA device.
	 */
	(void)pci_addr;

	result = doca_devinfo_list_create(&dev_list, &nb_devs);
	if (result != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: doca_devinfo_list_create: %s\n", errstr(result));
		return -1;
	}

	if (nb_devs == 0) {
		fprintf(stderr, "ERROR: no DOCA devices found\n");
		doca_devinfo_list_destroy(dev_list);
		return -1;
	}

	result = doca_dev_open(dev_list[0], dev);
	if (result != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: doca_dev_open: %s\n", errstr(result));
		doca_devinfo_list_destroy(dev_list);
		return -1;
	}

	doca_devinfo_list_destroy(dev_list);
	return 0;
}

static struct doca_mmap *create_mmap_or_die(struct doca_dev *dev,
					    void *addr,
					    size_t len,
					    uint32_t permissions)
{
	struct doca_mmap *mmap = NULL;
	union doca_data user_data;
	doca_error_t result;

	memset(&user_data, 0, sizeof(user_data));

	result = doca_mmap_create(&user_data, &mmap);
	check_doca(result, "doca_mmap_create");

	result = doca_mmap_set_memrange(mmap, addr, len);
	check_doca(result, "doca_mmap_set_memrange");

	result = doca_mmap_set_permissions(mmap, permissions);
	check_doca(result, "doca_mmap_set_permissions");

	result = doca_mmap_dev_add(mmap, dev);
	check_doca(result, "doca_mmap_dev_add");

	result = doca_mmap_start(mmap);
	check_doca(result, "doca_mmap_start");

	return mmap;
}

static void join_path(const char *dir, const char *file, char *out, size_t out_len)
{
	if (file[0] == '/')
		snprintf(out, out_len, "%s", file);
	else
		snprintf(out, out_len, "%s/%s", dir, file);
}

static int should_skip_dirent(const struct dirent *ent)
{
	if (strcmp(ent->d_name, ".") == 0 || strcmp(ent->d_name, "..") == 0)
		return 1;

	if (strcmp(ent->d_name, "manifest.csv") == 0)
		return 1;

	return 0;
}

static int run_one_decompress(struct doca_ctx *ctx,
			      struct doca_dev *dev,
			      struct doca_workq *workq,
			      const char *input_path,
			      size_t dst_len,
			      int write_output,
			      const char *output_path,
			      struct one_result *out)
{
	uint8_t *src_mem = NULL;
	uint8_t *dst_mem = NULL;
	size_t src_len = 0;
	size_t final_len = 0;

	struct doca_mmap *src_mmap = NULL;
	struct doca_mmap *dst_mmap = NULL;
	struct doca_buf_inventory *inventory = NULL;
	struct doca_buf *src_buf = NULL;
	struct doca_buf *dst_buf = NULL;

	union doca_data inventory_user_data;
	struct doca_compress_deflate_job job;
	struct doca_event event;

	uint64_t checksum = 0;
	uint64_t start_ns;
	uint64_t end_ns;
	uint64_t elapsed_ns;
	double elapsed_us;
	double mibps;

	doca_error_t result;
	int ret = -1;
	uint16_t refcnt = 0;

	memset(out, 0, sizeof(*out));
	memset(&inventory_user_data, 0, sizeof(inventory_user_data));
	memset(&job, 0, sizeof(job));
	memset(&event, 0, sizeof(event));

	src_mem = read_file_or_die(input_path, &src_len);

	dst_mem = calloc(1, dst_len);
	if (dst_mem == NULL) {
		fprintf(stderr, "ERROR: calloc destination failed\n");
		goto cleanup;
	}

	src_mmap = create_mmap_or_die(dev, src_mem, src_len, DOCA_ACCESS_LOCAL_READ_ONLY);
	dst_mmap = create_mmap_or_die(dev, dst_mem, dst_len, DOCA_ACCESS_LOCAL_READ_WRITE);

	result = doca_buf_inventory_create(&inventory_user_data, NUM_BUFS, 0, &inventory);
	if (result != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: doca_buf_inventory_create: %s\n", errstr(result));
		goto cleanup;
	}

	result = doca_buf_inventory_start(inventory);
	if (result != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: doca_buf_inventory_start: %s\n", errstr(result));
		goto cleanup;
	}

	result = doca_buf_inventory_buf_by_args(
		inventory,
		src_mmap,
		src_mem,
		src_len,
		src_mem,
		src_len,
		&src_buf);
	if (result != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: doca_buf_inventory_buf_by_args(src): %s\n", errstr(result));
		goto cleanup;
	}

	result = doca_buf_inventory_buf_by_args(
		inventory,
		dst_mmap,
		dst_mem,
		dst_len,
		dst_mem,
		0,
		&dst_buf);
	if (result != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: doca_buf_inventory_buf_by_args(dst): %s\n", errstr(result));
		goto cleanup;
	}

	job.base.type = DOCA_DECOMPRESS_DEFLATE_JOB;
	job.base.flags = DOCA_JOB_FLAGS_NONE;
	job.base.ctx = ctx;
	job.src_buff = src_buf;
	job.dst_buff = dst_buf;
	job.output_chksum = &checksum;

	/*
	 * Timing starts here.
	 * This excludes file read, malloc/calloc, mmap, inventory, buffer setup, and job construction.
	 * It includes only job submit + workq progress until completion.
	 */
	start_ns = now_nsec();

	result = doca_workq_submit(workq, &job.base);
	if (result != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: doca_workq_submit: %s\n", errstr(result));
		goto cleanup;
	}

	do {
		result = doca_workq_progress_retrieve(workq, &event, DOCA_WORKQ_RETRIEVE_FLAGS_NONE);
	} while (result == DOCA_ERROR_AGAIN);

	end_ns = now_nsec();

	if (result != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: doca_workq_progress_retrieve: %s\n", errstr(result));
		goto cleanup;
	}

	if ((doca_error_t)event.result.u64 != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: decompress job failed: %s\n",
			errstr((doca_error_t)event.result.u64));
		goto cleanup;
	}

	result = doca_buf_get_data_len(dst_buf, &final_len);
	if (result != DOCA_SUCCESS) {
		fprintf(stderr, "ERROR: doca_buf_get_data_len(dst): %s\n", errstr(result));
		goto cleanup;
	}

	elapsed_ns = end_ns - start_ns;
	if (elapsed_ns == 0)
		elapsed_ns = 1;

	elapsed_us = (double)elapsed_ns / 1000.0;

	mibps = ((double)final_len / (1024.0 * 1024.0)) /
		(elapsed_us / 1000000.0);

	if (write_output && output_path != NULL)
		write_file_or_die(output_path, dst_mem, final_len);

	out->compressed_bytes = (uint64_t)src_len;
	out->decompressed_bytes = (uint64_t)final_len;
	out->elapsed_us = elapsed_us;
	out->bw_mibps = mibps;
	out->crc = (uint32_t)(checksum & 0xffffffffu);
	out->adler = (uint32_t)((checksum >> 32) & 0xffffffffu);

	ret = 0;

cleanup:
	if (src_buf != NULL)
		(void)doca_buf_refcount_rm(src_buf, &refcnt);
	if (dst_buf != NULL)
		(void)doca_buf_refcount_rm(dst_buf, &refcnt);
	if (inventory != NULL)
		(void)doca_buf_inventory_destroy(inventory);
	if (src_mmap != NULL)
		(void)doca_mmap_destroy(src_mmap);
	if (dst_mmap != NULL)
		(void)doca_mmap_destroy(dst_mmap);

	free(src_mem);
	free(dst_mem);

	return ret;
}

static int run_directory(struct doca_ctx *ctx,
			 struct doca_dev *dev,
			 struct doca_workq *workq,
			 const char *dir_path,
			 size_t dst_len)
{
	struct dirent **namelist = NULL;
	int n;
	int count = 0;

	uint64_t total_compressed_bytes = 0;
	uint64_t total_decompressed_bytes = 0;
	double total_elapsed_us = 0.0;
	double total_bw_mibps;

	n = scandir(dir_path, &namelist, NULL, alphasort);
	if (n < 0) {
		fprintf(stderr, "ERROR: scandir(%s): %s\n", dir_path, strerror(errno));
		return -1;
	}

	for (int i = 0; i < n; i++) {
		char path[PATH_MAX];
		struct stat st;
		struct one_result r;

		if (should_skip_dirent(namelist[i])) {
			free(namelist[i]);
			continue;
		}

		join_path(dir_path, namelist[i]->d_name, path, sizeof(path));

		if (stat(path, &st) != 0) {
			fprintf(stderr, "ERROR: stat(%s): %s\n", path, strerror(errno));
			free(namelist[i]);
			goto fail;
		}

		if (!S_ISREG(st.st_mode)) {
			free(namelist[i]);
			continue;
		}

		if (run_one_decompress(ctx,
				       dev,
				       workq,
				       path,
				       dst_len,
				       0,
				       NULL,
				       &r) != 0) {
			fprintf(stderr, "ERROR: failed on file %s\n", path);
			free(namelist[i]);
			goto fail;
		}

		count++;
		total_compressed_bytes += r.compressed_bytes;
		total_decompressed_bytes += r.decompressed_bytes;
		total_elapsed_us += r.elapsed_us;

		free(namelist[i]);
	}

	free(namelist);

	if (count == 0) {
		fprintf(stderr, "ERROR: no regular chunk files found in directory: %s\n", dir_path);
		return -1;
	}

	total_bw_mibps = ((double)total_decompressed_bytes / (1024.0 * 1024.0)) /
			(total_elapsed_us / 1000000.0);

	printf("Decompress complete\n");
	printf("  chunks: %d\n", count);
	printf("  compressed_bytes: %" PRIu64 "\n", total_compressed_bytes);
	printf("  bytes: %" PRIu64 "\n", total_decompressed_bytes);
	printf("  time:  %.3f usec\n", total_elapsed_us);
	printf("  BW:    %.6f MiB/s\n", total_bw_mibps);

	return 0;

fail:
	for (int j = 0; j < n; j++) {
		if (namelist[j] != NULL)
			free(namelist[j]);
	}
	free(namelist);
	return -1;
}

int main(int argc, char **argv)
{
	const char *pci = NULL;
	const char *input_path = NULL;
	const char *output_path = DEFAULT_OUTPUT;
	const char *dir_path = NULL;

	int write_output = 1;
	int directory_mode = 0;

	size_t dst_len = 0;

	struct doca_dev *dev = NULL;
	struct doca_compress *compress = NULL;
	struct doca_ctx *ctx = NULL;
	struct doca_workq *workq = NULL;

	union doca_data compress_user_data;
	doca_error_t result;

	int opt;
	int option_index = 0;
	int ret = EXIT_FAILURE;

	static struct option long_options[] = {
		{"no-write", no_argument, 0, 1000},
		{"help", no_argument, 0, 'h'},
		{0, 0, 0, 0}
	};

	memset(&compress_user_data, 0, sizeof(compress_user_data));

	while ((opt = getopt_long(argc, argv, "p:f:o:s:d:S:h", long_options, &option_index)) != -1) {
		switch (opt) {
		case 'p':
			pci = optarg;
			break;
		case 'f':
			input_path = optarg;
			break;
		case 'o':
			output_path = optarg;
			write_output = 1;
			break;
		case 's':
			dst_len = parse_size_or_die(optarg);
			break;
		case 'd':
			dir_path = optarg;
			directory_mode = 1;
			write_output = 0;
			break;
		case 'S':
			dst_len = parse_size_or_die(optarg);
			break;
		case 1000:
			write_output = 0;
			break;
		case 'h':
		default:
			usage(argv[0]);
			return EXIT_FAILURE;
		}
	}

	if (pci == NULL) {
		usage(argv[0]);
		return EXIT_FAILURE;
	}

	if (!directory_mode && (input_path == NULL || dst_len == 0)) {
		usage(argv[0]);
		return EXIT_FAILURE;
	}

	if (directory_mode && (dir_path == NULL || dst_len == 0)) {
		usage(argv[0]);
		return EXIT_FAILURE;
	}

	if (open_first_doca_device(pci, &dev) != 0)
		goto cleanup;

	result = doca_compress_create(&compress);
	check_doca(result, "doca_compress_create");

	ctx = doca_compress_as_ctx(compress);
	if (ctx == NULL) {
		fprintf(stderr, "ERROR: doca_compress_as_ctx returned NULL\n");
		goto cleanup;
	}

	result = doca_ctx_dev_add(ctx, dev);
	check_doca(result, "doca_ctx_dev_add");

	result = doca_ctx_start(ctx);
	check_doca(result, "doca_ctx_start");

	result = doca_workq_create(WORKQ_DEPTH, &workq);
	check_doca(result, "doca_workq_create");

	result = doca_ctx_workq_add(ctx, workq);
	check_doca(result, "doca_ctx_workq_add");

	if (directory_mode) {
		if (run_directory(ctx, dev, workq, dir_path, dst_len) != 0)
			goto cleanup;
	} else {
		struct one_result r;

		if (run_one_decompress(ctx,
				       dev,
				       workq,
				       input_path,
				       dst_len,
				       write_output,
				       output_path,
				       &r) != 0)
			goto cleanup;

		printf("Decompress complete\n");
		printf("  compressed_bytes: %" PRIu64 "\n", r.compressed_bytes);
		printf("  bytes: %" PRIu64 "\n", r.decompressed_bytes);
		printf("  time:  %.3f usec\n", r.elapsed_us);
		printf("  BW:    %.6f MiB/s\n", r.bw_mibps);
		printf("  crc:   0x%08" PRIx32 "\n", r.crc);
		printf("  adler: 0x%08" PRIx32 "\n", r.adler);
		if (write_output)
			printf("  output: %s\n", output_path);
	}

	ret = EXIT_SUCCESS;

cleanup:
	if (workq != NULL)
		(void)doca_workq_destroy(workq);
	if (ctx != NULL)
		(void)doca_ctx_stop(ctx);
	if (compress != NULL)
		(void)doca_compress_destroy(compress);
	if (dev != NULL)
		(void)doca_dev_close(dev);

	return ret;
}