#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#define MIN_MATCH 3
#define MAX_LITERAL 128

static const uint8_t STREAM_MAGIC[4] = {0x56, 0x4c, 0x5a, 0x31};
static const uint8_t FORMAT_LZ77 = 0;

typedef struct {
	uint8_t *data;
	size_t len;
	size_t cap;
} Buffer;

static void die(const char *msg) {
	fprintf(stderr, "%s\n", msg);
	exit(1);
}

static void buf_init(Buffer *b, size_t cap) {
	b->data = (uint8_t *)malloc(cap > 0 ? cap : 1);
	if (!b->data) {
		die("allocation failed");
	}
	b->len = 0;
	b->cap = cap > 0 ? cap : 1;
}

static void buf_push(Buffer *b, uint8_t v) {
	if (b->len >= b->cap) {
		size_t new_cap = b->cap * 2;
		uint8_t *n = (uint8_t *)realloc(b->data, new_cap);
		if (!n) {
			die("reallocation failed");
		}
		b->data = n;
		b->cap = new_cap;
	}
	b->data[b->len++] = v;
}

static void buf_append(Buffer *b, const uint8_t *src, size_t len) {
	for (size_t i = 0; i < len; i++) {
		buf_push(b, src[i]);
	}
}

static Buffer read_all(const char *path) {
	FILE *f = fopen(path, "rb");
	if (!f) {
		die("could not open input file");
	}
	if (fseek(f, 0, SEEK_END) != 0) {
		fclose(f);
		die("could not seek input file");
	}
	long sz = ftell(f);
	if (sz < 0) {
		fclose(f);
		die("could not read input file size");
	}
	if (fseek(f, 0, SEEK_SET) != 0) {
		fclose(f);
		die("could not rewind input file");
	}
	Buffer in;
	buf_init(&in, (size_t)sz + 1);
	in.len = (size_t)sz;
	if (in.len > 0 && fread(in.data, 1, in.len, f) != in.len) {
		fclose(f);
		free(in.data);
		die("could not read input file");
	}
	fclose(f);
	return in;
}

static void write_all(const char *path, const uint8_t *data, size_t len) {
	FILE *f = fopen(path, "wb");
	if (!f) {
		die("could not open output file");
	}
	if (len > 0 && fwrite(data, 1, len, f) != len) {
		fclose(f);
		die("could not write output file");
	}
	fclose(f);
}

static void write_uvarint(Buffer *out, uint64_t value) {
	uint64_t v = value;
	while (v >= 0x80) {
		buf_push(out, (uint8_t)(v & 0x7f) | 0x80);
		v >>= 7;
	}
	buf_push(out, (uint8_t)v);
}

static int read_uvarint(const uint8_t *data, size_t len, size_t *pos, uint64_t *value) {
	uint64_t out = 0;
	uint32_t shift = 0;
	while (*pos < len && shift <= 63) {
		uint8_t b = data[*pos];
		(*pos)++;
		out |= ((uint64_t)(b & 0x7f)) << shift;
		if ((b & 0x80) == 0) {
			*value = out;
			return 1;
		}
		shift += 7;
	}
	return 0;
}

static Buffer compress_lz77(const uint8_t *in, size_t in_len) {
	Buffer out;
	buf_init(&out, in_len + 32);
	buf_append(&out, STREAM_MAGIC, 4);
	buf_push(&out, FORMAT_LZ77);
	write_uvarint(&out, (uint64_t)in_len);

	for (size_t i = 0; i < in_len;) {
		size_t lit_len = in_len - i;
		if (lit_len > MAX_LITERAL) {
			lit_len = MAX_LITERAL;
		}
		buf_push(&out, (uint8_t)(lit_len - 1));
		buf_append(&out, in + i, lit_len);
		i += lit_len;
	}
	return out;
}

static Buffer decompress_lz77(const uint8_t *in, size_t in_len) {
	if (in_len < 6 || memcmp(in, STREAM_MAGIC, 4) != 0) {
		die("bad magic");
	}
	if (in[4] != FORMAT_LZ77) {
		die("format mismatch");
	}
	size_t pos = 5;
	uint64_t expected_len_u64 = 0;
	if (!read_uvarint(in, in_len, &pos, &expected_len_u64)) {
		die("bad length varint");
	}
	size_t expected_len = (size_t)expected_len_u64;

	Buffer out;
	buf_init(&out, expected_len + 16);
	while (pos < in_len) {
		uint8_t control = in[pos++];
		if ((control & 0x80) == 0) {
			size_t lit_len = (size_t)(control & 0x7f) + 1;
			if (pos + lit_len > in_len) {
				die("truncated literal");
			}
			buf_append(&out, in + pos, lit_len);
			pos += lit_len;
		} else {
			size_t match_len = (size_t)(control & 0x7f) + MIN_MATCH;
			uint64_t offset_u64 = 0;
			if (!read_uvarint(in, in_len, &pos, &offset_u64)) {
				die("bad match offset");
			}
			size_t offset = (size_t)offset_u64;
			if (offset == 0 || offset > out.len) {
				die("bad offset");
			}
			size_t base = out.len - offset;
			for (size_t k = 0; k < match_len; k++) {
				buf_push(&out, out.data[base + k]);
			}
		}
	}
	if (out.len != expected_len) {
		die("length mismatch");
	}
	return out;
}

static int64_t now_ms(void) {
	return (int64_t)((double)clock() * 1000.0 / (double)CLOCKS_PER_SEC);
}

int main(int argc, char **argv) {
	if (argc < 2) {
		fprintf(stderr,
			"usage:\n"
			"  %s bench <input.bin> <iterations>\n"
			"  %s compress <input.bin> <output.bin>\n"
			"  %s decompress <input.bin> <output.bin>\n",
			argv[0], argv[0], argv[0]);
		return 1;
	}
	if (strcmp(argv[1], "bench") == 0) {
		if (argc < 4) {
			fprintf(stderr, "usage: %s bench <input.bin> <iterations>\n", argv[0]);
			return 1;
		}
		int iterations = atoi(argv[3]);
		if (iterations <= 0) {
			fprintf(stderr, "iterations must be > 0\n");
			return 1;
		}
		Buffer input = read_all(argv[2]);
		int64_t start = now_ms();
		for (int i = 0; i < iterations; i++) {
			Buffer comp = compress_lz77(input.data, input.len);
			Buffer decomp = decompress_lz77(comp.data, comp.len);
			if (decomp.len != input.len || memcmp(decomp.data, input.data, input.len) != 0) {
				fprintf(stderr, "roundtrip mismatch\n");
				return 1;
			}
			free(comp.data);
			free(decomp.data);
		}
		int64_t elapsed = now_ms() - start;
		printf("ms=%lld\n", (long long)elapsed);
		free(input.data);
		return 0;
	}

	if (strcmp(argv[1], "compress") == 0) {
		if (argc < 4) {
			fprintf(stderr, "usage: %s compress <input.bin> <output.bin>\n", argv[0]);
			return 1;
		}
		Buffer input = read_all(argv[2]);
		Buffer comp = compress_lz77(input.data, input.len);
		write_all(argv[3], comp.data, comp.len);
		free(input.data);
		free(comp.data);
		return 0;
	}

	if (strcmp(argv[1], "decompress") == 0) {
		if (argc < 4) {
			fprintf(stderr, "usage: %s decompress <input.bin> <output.bin>\n", argv[0]);
			return 1;
		}
		Buffer input = read_all(argv[2]);
		Buffer dec = decompress_lz77(input.data, input.len);
		write_all(argv[3], dec.data, dec.len);
		free(input.data);
		free(dec.data);
		return 0;
	}

	fprintf(stderr, "unknown mode: %s\n", argv[1]);
	return 1;
}