#define _GNU_SOURCE
#define _FILE_OFFSET_BITS	64

#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdbool.h>
#include <stdint.h>
#include <limits.h>
#include <string.h>

#include <libdeflate.h>
#include <zstd.h>
#include <lzo/lzo1x.h>
#include <lz4.h>
#include <lz4hc.h>

#define DEFAULT_BLOCK_SIZE	4096
#define DEFAULT_CHUNK_SIZE	65536

#define __le16	uint16_t
#define __le32	uint32_t
#define __le64	uint64_t

#include "dm-update.h"

static bool dio;
static unsigned long block_size;
static unsigned long chunk_size;
static uint8_t block_size_bits;

static const char *device[2];
static int handle[2];
static off_t size[2];
static char *block_buffer[2];

static const char *output_file;
static int output_handle;

static char *chunk_buffer;
static unsigned chunk_offset;

static struct update_entry *update_entries;
static size_t n_update_entries;

static uint64_t src_pos;

static struct update_superblock superblock;

static void compress_none(void *src, size_t src_len, void **dst, size_t *dst_len)
{
	*dst = malloc(src_len);
	if (!*dst)
		perror("malloc"), exit(1);
	memcpy(*dst, src, src_len);
	*dst_len = src_len;
}

static void compress_deflate(void *src, size_t src_len, void **dst, size_t *dst_len)
{
	struct libdeflate_compressor *cmp;
	size_t max_size;

	cmp = libdeflate_alloc_compressor(12);
	if (!cmp)
		perror("libdeflate_alloc_compressor"), exit(1);
	max_size = libdeflate_deflate_compress_bound(cmp, src_len);

	*dst = malloc(max_size);
	if (!*dst)
		perror("malloc"), exit(1);
	
	*dst_len = libdeflate_deflate_compress(cmp, src, src_len, *dst, max_size);
	if (!*dst_len)
		fprintf(stderr, "libdeflate_deflate_compress overrun\n"), exit(1);

	libdeflate_free_compressor(cmp);
}

static void compress_zstd(void *src, size_t src_len, void **dst, size_t *dst_len)
{
	size_t max_size;

	max_size = ZSTD_compressBound(src_len);

	*dst = malloc(max_size);
	if (!*dst)
		perror("malloc"), exit(1);
	
	*dst_len = ZSTD_compress(*dst, max_size, src, src_len, ZSTD_maxCLevel());

	if (ZSTD_isError(*dst_len))
		fprintf(stderr, "%s\n", ZSTD_getErrorName(*dst_len)), exit(1);
}

static void *lzo_work_mem;

static void init_lzo(void)
{
	if (lzo_init() != LZO_E_OK)
		fprintf(stderr, "lzo_init() failed\n"), exit(1);

	lzo_work_mem = malloc(LZO1X_999_MEM_COMPRESS);
	if (!lzo_work_mem)
		perror("malloc"), exit(1);
}

static void compress_lzo(void *src, size_t src_len, void **dst, size_t *dst_len)
{
	size_t max_size;
	int r;

	max_size = (src_len + src_len / 16 + 64 + 3);

	*dst = malloc(max_size);
	if (!*dst)
		perror("malloc"), exit(1);

	r = lzo1x_999_compress(src, src_len, *dst, dst_len, lzo_work_mem);
	if (r != LZO_E_OK)
		fprintf(stderr, "lzo compress failed: %d\n", r), exit(1);
}

static void compress_lz4(void *src, size_t src_len, void **dst, size_t *dst_len)
{
	size_t max_size;

	max_size = LZ4_compressBound(src_len);

	*dst = malloc(max_size);
	if (!*dst)
		perror("malloc"), exit(1);
	
	*dst_len = LZ4_compress_default(src, *dst, src_len, max_size);
	if (!*dst_len)
		fprintf(stderr, "LZ4_compress_default failed\n"), exit(1);
}

static void compress_lz4hc(void *src, size_t src_len, void **dst, size_t *dst_len)
{
	size_t max_size;

	max_size = LZ4_compressBound(src_len);

	*dst = malloc(max_size);
	if (!*dst)
		perror("malloc"), exit(1);
	
	*dst_len = LZ4_compress_HC(src, *dst, src_len, max_size, LZ4HC_CLEVEL_MAX);
	if (!*dst_len)
		fprintf(stderr, "LZ4_compress_HC failed\n"), exit(1);
}

static const struct compression_method {
	const char *name;
	void (*init)(void);
	void (*compress)(void *src, size_t src_len, void **dst, size_t *dst_len);
} methods[] = {
	{ "none", NULL, compress_none },
	{ "deflate", NULL, compress_deflate },
	{ "zstd", NULL, compress_zstd },
	{ "lzo", init_lzo, compress_lzo },
	{ "lz4", NULL, compress_lz4 },
	{ "lz4hc", NULL, compress_lz4hc },
};

static const struct compression_method *method;

static void rd(int h, char *data, size_t len, off_t off)
{
	ssize_t r;
again:
	r = pread(h, data, len, off);
	if (r == -1)
		perror("pread"), exit(1);
	if ((size_t)r < len) {
		if (!r)
			fprintf(stderr, "short read\n"), exit(1);
		data += r;
		len -= r;
		off += r;
		goto again;
	}
}

static void wr(const char *data, size_t len, off_t off)
{
	ssize_t w;
again:
	w = pwrite(output_handle, data, len, off);
	if (w == -1)
		perror("pwrite"), exit(1);
	if ((size_t)w < len) {
		data += w;
		len -= w;
		off += w;
		goto again;
	}
}

static struct update_entry *new_update_entry(void)
{
	if (!((n_update_entries + 1) & n_update_entries)) {
		update_entries = realloc(update_entries, (n_update_entries * 2 + 1) * sizeof(struct update_entry));
		if (!update_entries)
			perror("realloc"), exit(1);
	}
	return &update_entries[n_update_entries++];
}

static void flush_chunk(void)
{
	void *result;
	size_t result_len;
	method->compress(chunk_buffer, chunk_offset, &result, &result_len);
	wr(result, result_len, src_pos);
	src_pos += result_len;
	free(result);
	chunk_offset = 0;
}

static void calculate_difference(void)
{
	struct update_entry *e;
	off_t o;
	for (o = 0; o < size[0]; o += block_size) {
		int i;
		for (i = 0; i < 2; i++) {
			rd(handle[i], block_buffer[i], block_size, o);
		}
		if (memcmp(block_buffer[0], block_buffer[1], block_size)) {
			uint64_t blk;
			if (chunk_offset == chunk_size) {
				flush_chunk();
			}
			/*fprintf(stderr, "diff at %llx\n", (long long)o);*/

			e = new_update_entry();
			blk = o >> block_size_bits;
			e->dest_lo = htole32(blk);
			e->dest_hi = htole16(blk >> 32);
			e->src_lo = htole32(src_pos);
			e->src_hi = htole16(src_pos >> 32);
			e->offset = htole32(chunk_offset >> block_size_bits);

			memcpy(chunk_buffer + chunk_offset, block_buffer[1], block_size);
			chunk_offset += block_size;
		}
	}
	if (chunk_offset)
		flush_chunk();
	e = new_update_entry();
	e->dest_lo = htole32(-1);
	e->dest_hi = htole16(-1);
	e->src_lo = htole32(src_pos);
	e->src_hi = htole16(src_pos >> 32);
	e->offset = htole32(0);
}

static void write_directory(void)
{
	void *result;
	size_t result_len;
	void *padding;
	size_t padding_len;
	method->compress(update_entries, n_update_entries * sizeof(struct update_entry), &result, &result_len);
	wr(result, result_len, src_pos);
	free(result);

	superblock.dir_offset = htole64(src_pos);
	superblock.dir_compressed_size = htole64(result_len);
	superblock.dir_n = htole64(n_update_entries);

	src_pos += result_len;

	padding_len = -src_pos & (block_size - 1);
	padding = calloc(1, padding_len);
	if (!padding)
		perror("calloc"), exit(1);
	wr(padding, padding_len, src_pos);
	free(padding);
}

static void write_superblock(void)
{
	wr((const char *)&superblock, sizeof(struct update_superblock), 0);
}

int main(int argc, const char * const argv[])
{
	bool no_args;
	int a, i;
	int dev_idx;
	off_t o;

	memset(&superblock, 0, sizeof(struct update_superblock));

	dio = false;
	block_size = DEFAULT_BLOCK_SIZE;
	chunk_size = 0;
	output_file = NULL;
	method = &methods[0];

	no_args = false;
	dev_idx = 0;
	for (a = 1; a < argc; a++) {
		char *end;
		const char *str = argv[a];
		if (!no_args && str[0] == '-') {
			if (!strcmp(str, "--")) {
				no_args = true;
				continue;
			}
			if (!strcmp(str, "--help")) {
				printf("Usage:\n\
update_diff old.img new.img -o diff.img\n\
	-b block size\n\
	-c compressed chunk size\n\
	-a compression algorithm (none, deflate, zstd, lzo, lz4, lz4hc)\n\
	-d use direct I/O\n\
	--help display help\n\
");
				return 0;
			}
			if (!strcmp(str, "-d")) {
				dio = true;
				continue;
			}
			if (a == argc - 1)
invalid_args:
				fprintf(stderr, "update-diff: invalid arguments\n"), exit(1);
			if (!strcmp(str, "-b")) {
				block_size = strtoul(argv[++a], &end, 10);
				if (*end ||
				    block_size & (block_size - 1) ||
				    block_size < 512 ||
				    block_size >= UINT_MAX)
					fprintf(stderr, "update-diff: invalid block size\n"), exit(1);
				continue;
			}
			if (!strcmp(str, "-c")) {
				chunk_size = strtoul(argv[++a], &end, 10);
				if (*end ||
				    chunk_size < 512 ||
				    chunk_size >= UINT_MAX)
					fprintf(stderr, "update-diff: invalid compressed chunk size\n"), exit(1);
				continue;
			}
			if (!strcmp(str, "-a")) {
				const char *m = argv[++a];
				unsigned u;
				for (u = 0; u < sizeof(methods) / sizeof(*methods); u++)
					if (!strcmp(m, methods[u].name))
						goto found_compression_method;
				fprintf(stderr, "unknown compression method\n");
				exit(1);
found_compression_method:
				method = &methods[u];
				continue;
			}
			if (!strcmp(str, "-o")) {
				output_file = argv[++a];
				continue;
			}
			goto invalid_args;
		}
		if (dev_idx >= 2)
			goto invalid_args;
		device[dev_idx++] = str;
	}

	if (dev_idx != 2)
		fprintf(stderr, "update-diff: block devices not specified\n"), exit(1);

	if (method->init)
		method->init();

	memcpy(superblock.magic, UPDATE_MAGIC, sizeof superblock.magic);
	strncpy(superblock.compression, method->name, sizeof superblock.compression - 1);

	if (!chunk_size) {
		chunk_size = DEFAULT_CHUNK_SIZE;
		if (chunk_size < block_size)
			chunk_size = block_size;
	} else {
		if (!chunk_size || chunk_size % block_size)
			fprintf(stderr, "update-diff: compressed chunk size must be multiple of block size\n"), exit(1);
	}
	superblock.block_bits = block_size_bits = ffsll(block_size) - 1;
	superblock.chunk_blocks = htole32(chunk_size >> block_size_bits);

	for (i = 0; i < 2; i++) {
		handle[i] = open(device[i], O_RDONLY | (dio ? O_DIRECT : 0));
		if (handle[i] == -1)
			perror(device[i]), exit(1);
		size[i] = lseek(handle[i], 0, SEEK_END);
		if (size[i] == -1)
			perror("lseek"), exit(1);
		if (lseek(handle[i], 0, SEEK_SET) == -1)
			perror("lseek"), exit(1);
		block_buffer[i] = valloc(block_size);
		if (!block_buffer[i])
			perror("valloc"), exit(1);
	}

	if (size[0] != size[1])
		fprintf(stderr, "update-diff: the devices must have the same size\n"), exit(1);

	if (size[0] & (block_size - 1))
		fprintf(stderr, "update-diff: size is not aligned on block boundary\n"), exit(1);

	if (!output_file)
		fprintf(stderr, "update-diff: output file not specified\n"), exit(1);
	output_handle = open(output_file, O_WRONLY | O_CREAT | O_TRUNC, 0644);
	if (output_handle  == -1)
		perror(output_file), exit(1);
	o = lseek(output_handle, sizeof(struct update_superblock), SEEK_SET);
	if (o == -1)
		perror("lseek"), exit(1);
	src_pos = o;

	chunk_buffer = malloc(chunk_size);
	if (!chunk_buffer)
		perror("malloc"), exit(1);
	chunk_offset = 0;

	update_entries = NULL;
	n_update_entries = 0;

	calculate_difference();
	write_directory();
	write_superblock();

	for (i = 0; i < 2; i++) {
		if (close(handle[i]))
			perror("close"), exit(1);
	}

	if (close(output_handle))
		perror("close"), exit(1);

	return 0;
}