#include "Rts.h"
#include <string.h>
#include <time.h>

// needs C11
#include <threads.h>

static size_t min_sz(size_t a, size_t b) {
	return a < b ? a : b;
}

extern RtsConfig __attribute__((weak)) rtsConfig;

// A copy of GCDetails_ with known structure that can be depended on by the Haskell code.
struct ShadowDetails {
	int64_t timestamp_sec;
	int64_t timestamp_nsec;

	// The generation number of this GC
	uint32_t gen;
	// Number of threads used in this GC
	uint32_t threads;
	// Number of bytes allocated since the previous GC
	uint64_t allocated_bytes;
	// Total amount of live data in the heap (incliudes large + compact data).
	// Updated after every GC. Data in uncollected generations (in minor GCs)
	// are considered live.
	uint64_t live_bytes;
	// Total amount of live data in large objects
	uint64_t large_objects_bytes;
	// Total amount of live data in compact regions
	uint64_t compact_bytes;
	// Total amount of slop (wasted memory)
	uint64_t slop_bytes;
	// Total amount of memory in use by the RTS
	uint64_t mem_in_use_bytes;
	// Total amount of data copied during this GC
	uint64_t copied_bytes;
	// In parallel GC, the max amount of data copied by any one thread
	uint64_t par_max_copied_bytes;
	// In parallel GC, the amount of balanced data copied by all threads
	uint64_t par_balanced_copied_bytes;
	// The time elapsed during synchronisation before GC
	// NOTE: nanoseconds!
	uint64_t sync_elapsed_ns;
	// The CPU time used during GC itself
	// NOTE: nanoseconds!
	uint64_t cpu_ns;
	// The time elapsed during GC itself
	// NOTE: nanoseconds!
	uint64_t elapsed_ns;

	// Concurrent garbage collector

	// The CPU time used during the post-mark pause phase of the concurrent
	// nonmoving GC.
	// NOTE: nanoseconds!
	uint64_t nonmoving_gc_sync_cpu_ns;
	// The time elapsed during the post-mark pause phase of the concurrent
	// nonmoving GC.
	// NOTE: nanoseconds!
	uint64_t nonmoving_gc_sync_elapsed_ns;
	// The CPU time used during the post-mark pause phase of the concurrent
	// nonmoving GC.
	// NOTE: nanoseconds!
	uint64_t nonmoving_gc_cpu_ns;
	// The time elapsed during the post-mark pause phase of the concurrent
	// nonmoving GC.
	// NOTE: nanoseconds!
	uint64_t nonmoving_gc_elapsed_ns;
};

static void shadow_copy(struct ShadowDetails *dst, const struct GCDetails_ *src) {
#define COPY(field) dst->field = src->field;
#define COPYTIME(field) dst->field = TimeToNS(src->field);
	COPY(gen);
	COPY(threads);
	COPY(allocated_bytes);
	COPY(live_bytes);
	COPY(large_objects_bytes);
	COPY(compact_bytes);
	COPY(slop_bytes);
	COPY(mem_in_use_bytes);
	COPY(copied_bytes);
	COPY(par_max_copied_bytes);
	COPY(par_balanced_copied_bytes);
	COPYTIME(sync_elapsed_ns);
	COPYTIME(cpu_ns);
	COPYTIME(elapsed_ns);

	COPYTIME(nonmoving_gc_sync_cpu_ns);
	COPYTIME(nonmoving_gc_sync_elapsed_ns);
	COPYTIME(nonmoving_gc_cpu_ns);
	COPYTIME(nonmoving_gc_elapsed_ns);
#undef COPY
#undef COPYTIME
}

// --------
// GLOBAL VARIABLES
// --------

static bool constructor_worked = false;
static bool hook_initialised = false;
static bool logging_enabled = false;
static void (*hook_c_delegate)(const struct GCDetails_*) = NULL;

static mtx_t state_mutex;

static void (*old_hook)(const struct GCDetails_ *details) = NULL;
static size_t detlog_capacity = 0, detlog_length = 0;
static struct ShadowDetails *detlog = NULL;

// --------
// END OF GLOBAL VARIABLES
// --------

static void hook_callback(const struct GCDetails_ *details) {
	static bool fatal_failure = false;

	if (fatal_failure) goto cleanup_no_mutex;

	// Do this now already, before waiting on the mutex
	struct timespec now;
	if (logging_enabled && clock_gettime(CLOCK_MONOTONIC, &now) != 0) {
		perror("clock_gettime");
		fatal_failure = true;
		goto cleanup_no_mutex;
	}

	if (mtx_lock(&state_mutex) != thrd_success) {
		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex lock failed\n");
		fatal_failure = true;
		goto cleanup_no_mutex;
	}

	// mutex is locked from here

	if (logging_enabled) {
		if (detlog_length == detlog_capacity) {
			detlog_capacity = detlog_capacity == 0 ? 128 : 2 * detlog_capacity;
			detlog = realloc(detlog, detlog_capacity * sizeof(detlog[0]));
			if (detlog == NULL || detlog_capacity == 0) {  // also check for overflow here
				fprintf(stderr, "ghc-gc-hook: ERROR: Could not allocate memory for GC log hook\n");
				fatal_failure = true;
				goto cleanup;
			}
		}

		struct ShadowDetails *dst = &detlog[detlog_length];
		dst->timestamp_sec = now.tv_sec;
		dst->timestamp_nsec = now.tv_nsec;
		shadow_copy(dst, details);
		detlog_length++;
	}

	if (hook_c_delegate) hook_c_delegate(details);

cleanup:
	mtx_unlock(&state_mutex);  // ignore return value

cleanup_no_mutex:
	if (old_hook) old_hook(details);
}

__attribute__((constructor))
static void constructor(void) {
	if (mtx_init(&state_mutex, mtx_plain) != thrd_success) {
		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex initialisation failed\n");
		return;
	}

	constructor_worked = true;
}

// --------
// EXPORTED FUNCTIONS
// --------

// Only works if logging is enabled.
void copy_log_to_buffer(size_t space_available, char *buffer, size_t *unit_size, size_t *num_stored) {
	*unit_size = sizeof(detlog[0]);

	if (mtx_lock(&state_mutex) != thrd_success) {
		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex lock failed\n");
		*num_stored = 0;
		return;
	}

	if (detlog_length == 0) {
		*num_stored = 0;
		goto unlock_return;
	}

	const size_t n = min_sz(space_available / sizeof(detlog[0]), detlog_length);

	// First copy over the fitting items
	memcpy(buffer, detlog, n * sizeof(detlog[0]));
	*unit_size = sizeof(detlog[0]);
	*num_stored = n;

	// Then shift back the remaining items
	memmove(detlog, detlog + n, (detlog_length - n) * sizeof(detlog[0]));
	detlog_length -= n;

unlock_return:
	mtx_unlock(&state_mutex);
}

// Sets the GC hook, logging or C hook delegate not yet enabled. Returns success.
bool set_gchook(void) {
	if (mtx_lock(&state_mutex) != thrd_success) {
		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex lock failed\n");
		return false;
	}

	bool retval = false;

	if (!constructor_worked) {
		fprintf(stderr, "ghc-gc-hook: ERROR: Cannot set hook, system does not allow initialisation\n");
		goto unlock_return_retval;
	}

	if (hook_initialised) {
		fprintf(stderr, "ghc-gc-hook: ERROR: Hook already initialised\n");
		goto unlock_return_retval;
	}

	if (&rtsConfig == NULL) {
		fprintf(stderr, "ghc-gc-hook: ERROR: rtsConfig not defined; the GC hook cannot be used from within a TemplateHaskell splice\n");
		goto unlock_return_retval;
	}

	old_hook = rtsConfig.gcDoneHook;
	rtsConfig.gcDoneHook = hook_callback;

	hook_initialised = true;
	retval = true;

unlock_return_retval:
	mtx_unlock(&state_mutex);
	return retval;
}

// Enable logging on the GC hook.
void gchook_enable_logging(bool yes) {
	if (!hook_initialised) {
		if (!set_gchook()) exit(1);  // meh
	}

	if (mtx_lock(&state_mutex) != thrd_success) {
		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex lock failed\n");
		return;
	}

	if (logging_enabled && !yes) {
		detlog_length = 0;
		detlog_capacity = 0;
		free(detlog);
		detlog = NULL;
	}

	logging_enabled = yes;

	mtx_unlock(&state_mutex);
}

// Set a C function to be called after every GC with the GCDetails_ structure
// from `rts/include/RtsAPI.h`. Returns success.
bool gchook_set_c_delegate(void (*delegate)(const struct GCDetails_*)) {
	if (!hook_initialised) {
		if (!set_gchook()) exit(1);  // meh
	}

	if (mtx_lock(&state_mutex) != thrd_success) {
		fprintf(stderr, "ghc-gc-hook: ERROR: Mutex lock failed\n");
		return false;
	}

	bool retval = false;

	if (hook_c_delegate != NULL) {
		fprintf(stderr, "ghc-gc-hook: ERROR: C hook delegate already set\n");
		goto unlock_return_retval;
	}

	hook_c_delegate = delegate;
	retval = true;

unlock_return_retval:
	mtx_unlock(&state_mutex);
	return retval;
}

// vim: set noet sw=4 ts=4: