#define _GNU_SOURCE

#include <doca_error.h>
#include <doca_dev.h>
#include <doca_mmap.h>

#include <errno.h>
#include <getopt.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

#define META_MAGIC 0x444f4341444d4131ULL /* "DOCADMA1" */
#define DEFAULT_DESC_PATH "descriptor.bin"
#define DEFAULT_BUF_PATH  "buffer.json"

static void die_errno(const char *what)
{
	fprintf(stderr, "%s: %s\n", what, strerror(errno));
	exit(EXIT_FAILURE);
}

static void die_doca(const char *what, doca_error_t err)
{
	fprintf(stderr, "%s: DOCA error code %d\n", what, (int)err);
	exit(EXIT_FAILURE);
}

static int write_file_all(const char *path, const void *data, size_t len)
{
	FILE *f = fopen(path, "wb");

	if (f == NULL)
		return -1;

	if (len > 0 && fwrite(data, 1, len, f) != len) {
		fclose(f);
		return -1;
	}

	if (fclose(f) != 0)
		return -1;

	return 0;
}

static int write_remote_buffer_json(const char *path, uintptr_t addr, size_t len)
{
	FILE *f = fopen(path, "w");

	if (f == NULL)
		return -1;

	fprintf(f,
		"{\n"
		"  \"magic\": \"0x%016" PRIx64 "\",\n"
		"  \"addr\": \"0x%016" PRIxPTR "\",\n"
		"  \"len\": %" PRIu64 "\n"
		"}\n",
		(uint64_t)META_MAGIC,
		addr,
		(uint64_t)len);

	if (fclose(f) != 0)
		return -1;

	return 0;
}

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

	result = doca_devinfo_create_list(&dev_list, &nb_devs);
	if (result != DOCA_SUCCESS)
		return result;

	for (i = 0; i < nb_devs; i++) {
		uint8_t is_equal = 0;

		if (pci_addr == NULL) {
			result = doca_dev_open(dev_list[i], dev);
			doca_devinfo_destroy_list(dev_list);
			return result;
		}

		result = doca_devinfo_is_equal_pci_addr(dev_list[i],
							pci_addr,
							&is_equal);
		if (result == DOCA_SUCCESS && is_equal) {
			result = doca_dev_open(dev_list[i], dev);
			doca_devinfo_destroy_list(dev_list);
			return result;
		}
	}

	doca_devinfo_destroy_list(dev_list);
	return DOCA_ERROR_NOT_FOUND;
}

static doca_error_t create_started_host_mmap(struct doca_dev *dev,
					     void *addr,
					     size_t len,
					     struct doca_mmap **mmap)
{
	doca_error_t result;

	/*
	 * DOCA 2.x commonly uses:
	 *   doca_mmap_create(const union doca_data *user_data, struct doca_mmap **mmap)
	 */
	result = doca_mmap_create(mmap);
	if (result != DOCA_SUCCESS)
		return result;

	result = doca_mmap_set_memrange(*mmap, addr, len);
	if (result != DOCA_SUCCESS)
		goto destroy;

	result = doca_mmap_set_permissions(*mmap, DOCA_ACCESS_FLAG_PCI_READ_WRITE);
	if (result != DOCA_SUCCESS)
		goto destroy;

	result = doca_mmap_add_dev(*mmap, dev);
	if (result != DOCA_SUCCESS)
		goto destroy;

	result = doca_mmap_start(*mmap);
	if (result != DOCA_SUCCESS)
		goto destroy;

	return DOCA_SUCCESS;

destroy:
	doca_mmap_destroy(*mmap);
	*mmap = NULL;
	return result;
}

static void usage(const char *prog)
{
	fprintf(stderr,
		"Usage: %s -p <host-pci-bdf> -s <bytes> -o <output-file> "
		"[-d descriptor.bin] [-b buffer.json]\n\n"
		"Example:\n"
		"  sudo %s -p 0000:3b:00.0 -s 268435456 -o received_256M.bin\n",
		prog, prog);
}

int main(int argc, char **argv)
{
	const char *pci = NULL;
	const char *out_path = NULL;
	const char *desc_path = DEFAULT_DESC_PATH;
	const char *buf_path = DEFAULT_BUF_PATH;
	size_t size = 0;

	struct doca_dev *dev = NULL;
	struct doca_mmap *mmap = NULL;

	void *host_buf = NULL;
	const void *export_desc = NULL;
	size_t export_desc_len = 0;

	int opt;
	doca_error_t result;

	while ((opt = getopt(argc, argv, "p:s:o:d:b:h")) != -1) {
		switch (opt) {
		case 'p':
			pci = optarg;
			break;
		case 's':
			size = strtoull(optarg, NULL, 0);
			break;
		case 'o':
			out_path = optarg;
			break;
		case 'd':
			desc_path = optarg;
			break;
		case 'b':
			buf_path = optarg;
			break;
		case 'h':
		default:
			usage(argv[0]);
			return EXIT_FAILURE;
		}
	}

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

	result = open_doca_device_by_pci(pci, &dev);
	if (result != DOCA_SUCCESS)
		die_doca("open_doca_device_by_pci", result);

	if (posix_memalign(&host_buf, 4096, size) != 0)
		die_errno("posix_memalign");

	memset(host_buf, 0, size);

	result = create_started_host_mmap(dev, host_buf, size, &mmap);
	if (result != DOCA_SUCCESS)
		die_doca("create_started_host_mmap", result);

	result = doca_mmap_export_pci(mmap, dev, &export_desc, &export_desc_len);
	if (result != DOCA_SUCCESS)
		die_doca("doca_mmap_export_pci", result);

	if (export_desc == NULL || export_desc_len == 0) {
		fprintf(stderr, "doca_mmap_export_pci returned empty descriptor\n");
		exit(EXIT_FAILURE);
	}

	if (write_file_all(desc_path, export_desc, export_desc_len) != 0)
		die_errno("write descriptor");

	if (write_remote_buffer_json(buf_path, (uintptr_t)host_buf, size) != 0)
		die_errno("write buffer metadata");

	printf("Host receive buffer exported.\n");
	printf("  descriptor: %s (%zu bytes)\n", desc_path, export_desc_len);
	printf("  buffer:     %s\n", buf_path);
	printf("  addr:       0x%016" PRIxPTR "\n", (uintptr_t)host_buf);
	printf("  len:        %zu\n", size);
	printf("\nCopy %s and %s to the DPU, run dpu sender, then press ENTER here.\n",
	       desc_path, buf_path);

	(void)getchar();

	if (write_file_all(out_path, host_buf, size) != 0)
		die_errno("write output file");

	printf("Wrote host buffer to %s\n", out_path);

	doca_mmap_destroy(mmap);
	doca_dev_close(dev);
	free(host_buf);

	return EXIT_SUCCESS;
}