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

#define CODEC_COUNT 4
#define MOUNT_BUDGET (256U * 1024)
#define GLOBAL_BUDGET (512U * 1024)
#define MINIMUM_WORK (128U * 1024)

enum cache_state {
	CACHE_EMPTY,
	CACHE_INFLIGHT,
	CACHE_READY,
};

enum cache_lookup {
	CACHE_BYPASS,
	CACHE_HIT,
	CACHE_OWNER,
};

struct metric {
	uint64_t hits;
	uint64_t misses;
	uint64_t bypasses;
	uint64_t evictions;
	uint64_t reclaims;
	size_t resident_bytes;
};

struct budget {
	size_t limit;
	size_t used;
	size_t peak;
};

struct key {
	uint64_t identity;
	uint8_t codec;
};

struct cache {
	struct budget *global;
	struct metric metrics[CODEC_COUNT];
	struct key key;
	unsigned char *data;
	size_t charged_bytes;
	size_t budget_bytes;
	size_t peak_bytes;
	enum cache_state state;
	bool enabled;
};

static void
fail(const char *message)
{
	fprintf(stderr, "B33 cache oracle failure: %s\n", message);
	exit(1);
}

static uint64_t
hash_bytes(const unsigned char *data, size_t size)
{
	uint64_t hash;
	size_t index;

	hash = UINT64_C(1469598103934665603);
	for (index = 0; index < size; ++index) {
		hash ^= data[index];
		hash *= UINT64_C(1099511628211);
	}
	return (hash);
}

static void
fill(unsigned char *data, size_t size, const struct key *key)
{
	size_t index;

	for (index = 0; index < size; ++index)
		data[index] = (unsigned char)(key->identity + key->codec * 31 +
		    index * 17);
}

static bool
key_equal(const struct key *left, const struct key *right)
{

	return (left->identity == right->identity && left->codec == right->codec);
}

static bool
reserve(struct budget *budget, size_t bytes)
{

	if (bytes > budget->limit || budget->used > budget->limit - bytes)
		return (false);
	budget->used += bytes;
	if (budget->used > budget->peak)
		budget->peak = budget->used;
	return (true);
}

static void
release(struct budget *budget, size_t bytes)
{

	if (bytes > budget->used)
		fail("global budget underflow");
	budget->used -= bytes;
}

static void
cache_init(struct cache *cache, struct budget *global, bool enabled)
{

	memset(cache, 0, sizeof(*cache));
	cache->global = global;
	cache->budget_bytes = MOUNT_BUDGET;
	cache->enabled = enabled;
}

static bool
policy_admit(const struct cache *cache, uint8_t codec, size_t decoded_size,
    size_t compressed_size)
{

	return (cache->enabled && codec < CODEC_COUNT &&
	    decoded_size <= cache->budget_bytes &&
	    compressed_size + decoded_size >= MINIMUM_WORK);
}

static void
drop(struct cache *cache, bool evicted, bool reclaimed)
{
	struct metric *metric;

	if (cache->charged_bytes == 0)
		return;
	if (cache->key.codec >= CODEC_COUNT)
		fail("charged unknown codec");
	metric = &cache->metrics[cache->key.codec];
	if (cache->state == CACHE_READY) {
		if (metric->resident_bytes != cache->charged_bytes)
			fail("codec resident accounting mismatch");
		metric->resident_bytes = 0;
		if (evicted)
			++metric->evictions;
		if (reclaimed)
			++metric->reclaims;
	}
	release(cache->global, cache->charged_bytes);
	cache->charged_bytes = 0;
	free(cache->data);
	cache->data = NULL;
}

static enum cache_lookup
claim(struct cache *cache, const struct key *key, size_t decoded_size,
    size_t compressed_size, unsigned char *output)
{
	struct metric *metric;

	if (key->codec >= CODEC_COUNT)
		fail("unknown requested codec");
	metric = &cache->metrics[key->codec];
	if (!policy_admit(cache, key->codec, decoded_size, compressed_size)) {
		++metric->bypasses;
		return (CACHE_BYPASS);
	}
	if (cache->state == CACHE_READY && key_equal(&cache->key, key)) {
		memcpy(output, cache->data, decoded_size);
		++metric->hits;
		return (CACHE_HIT);
	}
	if (cache->state == CACHE_INFLIGHT) {
		++metric->bypasses;
		return (CACHE_BYPASS);
	}
	if (cache->state == CACHE_READY)
		drop(cache, true, false);
	cache->state = CACHE_EMPTY;
	if (!reserve(cache->global, decoded_size)) {
		++metric->bypasses;
		return (CACHE_BYPASS);
	}
	cache->charged_bytes = decoded_size;
	if (decoded_size > cache->peak_bytes)
		cache->peak_bytes = decoded_size;
	cache->key = *key;
	cache->state = CACHE_INFLIGHT;
	++metric->misses;
	return (CACHE_OWNER);
}

static void
complete(struct cache *cache, const struct key *key, unsigned char *data,
    int error)
{
	struct metric *metric;

	if (cache->state != CACHE_INFLIGHT || !key_equal(&cache->key, key))
		fail("owner lost key");
	metric = &cache->metrics[key->codec];
	if (error != 0) {
		if (data != NULL || error < 1)
			fail("failure publication is untyped");
		release(cache->global, cache->charged_bytes);
		cache->charged_bytes = 0;
		cache->state = CACHE_EMPTY;
		return;
	}
	if (data == NULL || metric->resident_bytes != 0)
		fail("success publication is untyped");
	cache->data = data;
	metric->resident_bytes = cache->charged_bytes;
	cache->state = CACHE_READY;
}

static bool
reclaim(struct cache *cache)
{

	if (cache->state == CACHE_INFLIGHT)
		return (false);
	if (cache->state == CACHE_READY) {
		drop(cache, true, true);
		cache->state = CACHE_EMPTY;
	}
	return (true);
}

static void
cache_fini(struct cache *cache)
{

	if (cache->state == CACHE_INFLIGHT)
		fail("unmount did not drain owner");
	drop(cache, false, false);
	cache->state = CACHE_EMPTY;
}

static unsigned char *
make_data(const struct key *key, size_t size)
{
	unsigned char *data;

	data = malloc(size);
	if (data == NULL)
		fail("decoded allocation");
	fill(data, size, key);
	return (data);
}

static void
test_codec_accounting(struct budget *global, struct metric output[CODEC_COUNT],
    uint64_t *eviction_hash)
{
	struct cache cache;
	struct key key;
	unsigned char *expected, *readback;
	uint64_t before, after;
	unsigned int codec;

	cache_init(&cache, global, true);
	expected = malloc(MOUNT_BUDGET);
	readback = malloc(MOUNT_BUDGET);
	if (expected == NULL || readback == NULL)
		fail("accounting buffers");
	before = 0;
	for (codec = 0; codec < CODEC_COUNT; ++codec) {
		key = (struct key){ .identity = 100 + codec, .codec = codec };
		fill(expected, MOUNT_BUDGET, &key);
		if (claim(&cache, &key, MOUNT_BUDGET, 4096, readback) != CACHE_OWNER)
			fail("codec-neutral request was not admitted");
		complete(&cache, &key, make_data(&key, MOUNT_BUDGET), 0);
		if (claim(&cache, &key, MOUNT_BUDGET, 4096, readback) != CACHE_HIT ||
		    memcmp(readback, expected, MOUNT_BUDGET) != 0)
			fail("codec cache hit changed bytes");
		if (codec == 0)
			before = hash_bytes(readback, MOUNT_BUDGET);
	}
	key = (struct key){ .identity = 100, .codec = 0 };
	if (claim(&cache, &key, MOUNT_BUDGET, 4096, readback) != CACHE_OWNER)
		fail("evicted codec key was not reusable");
	complete(&cache, &key, make_data(&key, MOUNT_BUDGET), 0);
	if (claim(&cache, &key, MOUNT_BUDGET, 4096, readback) != CACHE_HIT)
		fail("reused codec key did not hit");
	after = hash_bytes(readback, MOUNT_BUDGET);
	if (before != after)
		fail("eviction changed decoded hash");
	if (!reclaim(&cache) || global->used != 0)
		fail("ready reclaim did not release bytes");
	if (cache.peak_bytes != MOUNT_BUDGET)
		fail("mount peak did not equal hard budget");
	memcpy(output, cache.metrics, sizeof(cache.metrics));
	*eviction_hash = after;
	free(readback);
	free(expected);
}

static uint64_t
test_global_exhaustion(struct budget *global)
{
	struct cache first, second, third;
	struct key first_key = { .identity = 201, .codec = 0 };
	struct key second_key = { .identity = 202, .codec = 2 };
	struct key third_key = { .identity = 203, .codec = 3 };
	unsigned char *scratch;
	uint64_t fallback_hash;

	cache_init(&first, global, true);
	cache_init(&second, global, true);
	cache_init(&third, global, true);
	scratch = malloc(MOUNT_BUDGET);
	if (scratch == NULL)
		fail("global fallback allocation");
	if (claim(&first, &first_key, MOUNT_BUDGET, 4096, scratch) != CACHE_OWNER ||
	    claim(&second, &second_key, MOUNT_BUDGET, 4096, scratch) != CACHE_OWNER)
		fail("two mounts did not fill global budget");
	complete(&first, &first_key, make_data(&first_key, MOUNT_BUDGET), 0);
	complete(&second, &second_key, make_data(&second_key, MOUNT_BUDGET), 0);
	if (global->used != GLOBAL_BUDGET || global->peak != GLOBAL_BUDGET)
		fail("global hard budget was not reached exactly");
	if (claim(&third, &third_key, MOUNT_BUDGET, 4096, scratch) != CACHE_BYPASS)
		fail("global exhaustion did not bypass");
	fill(scratch, MOUNT_BUDGET, &third_key);
	fallback_hash = hash_bytes(scratch, MOUNT_BUDGET);
	if (!reclaim(&first) ||
	    claim(&third, &third_key, MOUNT_BUDGET, 4096, scratch) != CACHE_OWNER)
		fail("reclaimed global bytes were not reusable");
	complete(&third, &third_key, make_data(&third_key, MOUNT_BUDGET), 0);
	cache_fini(&second);
	cache_fini(&third);
	if (global->used != 0)
		fail("unmount did not release global bytes");
	free(scratch);
	return (fallback_hash);
}

static void
test_failure_and_policy(struct budget *global)
{
	struct cache cache, disabled;
	struct key key = { .identity = 301, .codec = 1 };
	unsigned char output[MOUNT_BUDGET];
	unsigned int codec;

	cache_init(&cache, global, true);
	if (claim(&cache, &key, MOUNT_BUDGET, 4096, output) != CACHE_OWNER)
		fail("failure setup was not admitted");
	if (reclaim(&cache))
		fail("inflight reclaim did not report busy");
	complete(&cache, &key, NULL, EIO);
	if (global->used != 0 ||
	    claim(&cache, &key, MOUNT_BUDGET, 4096, output) != CACHE_OWNER)
		fail("failure did not release budget for retry");
	complete(&cache, &key, make_data(&key, MOUNT_BUDGET), 0);
	cache_fini(&cache);

	cache_init(&disabled, global, false);
	for (codec = 0; codec < CODEC_COUNT; ++codec) {
		key.codec = codec;
		if (claim(&disabled, &key, MOUNT_BUDGET, 4096, output) != CACHE_BYPASS)
			fail("disabled codec policy did not bypass");
		disabled.enabled = true;
		if (claim(&disabled, &key, 64U * 1024, 1, output) != CACHE_BYPASS)
			fail("low-work policy did not bypass");
		disabled.enabled = false;
	}
	disabled.enabled = true;
	if (claim(&disabled, &key, MOUNT_BUDGET + 1, 4096, output) != CACHE_BYPASS)
		fail("oversized policy did not bypass");
	cache_fini(&disabled);
	if (global->used != 0)
		fail("policy tests leaked global bytes");
}

int
main(void)
{
	struct budget global = { .limit = GLOBAL_BUDGET };
	struct metric metrics[CODEC_COUNT];
	uint64_t eviction_hash, fallback_hash;
	unsigned int codec;

	test_codec_accounting(&global, metrics, &eviction_hash);
	fallback_hash = test_global_exhaustion(&global);
	test_failure_and_policy(&global);
	for (codec = 0; codec < CODEC_COUNT; ++codec) {
		if (metrics[codec].hits == 0 || metrics[codec].misses == 0 ||
		    metrics[codec].evictions == 0 || metrics[codec].resident_bytes != 0)
			fail("four-codec metrics are incomplete");
	}
	printf("{\"status\":\"PASS\",\"global_limit\":%u,"
	    "\"global_peak\":%zu,\"global_remaining\":%zu,"
	    "\"mount_limit\":%u,\"eviction_hash\":\"%016" PRIx64 "\","
	    "\"fallback_hash\":\"%016" PRIx64 "\",\"metrics\":[",
	    GLOBAL_BUDGET, global.peak, global.used, MOUNT_BUDGET,
	    eviction_hash, fallback_hash);
	for (codec = 0; codec < CODEC_COUNT; ++codec) {
		if (codec != 0)
			printf(",");
		printf("{\"codec\":%u,\"hits\":%" PRIu64
		    ",\"misses\":%" PRIu64 ",\"bypasses\":%" PRIu64
		    ",\"evictions\":%" PRIu64 ",\"reclaims\":%" PRIu64
		    ",\"resident_bytes\":%zu}", codec, metrics[codec].hits,
		    metrics[codec].misses, metrics[codec].bypasses,
		    metrics[codec].evictions, metrics[codec].reclaims,
		    metrics[codec].resident_bytes);
	}
	printf("]}\n");
	return (0);
}