#include "mem.h"

#include <uv.h>

#if !defined(_WIN32)
#include <openssl/crypto.h>
#endif

#include <quickjs.h>

#include <string.h>

static int64_t s_tf_malloc_size;
static int64_t s_uv_malloc_size;
static int64_t s_tls_malloc_size;
static int64_t s_js_malloc_size;

static void* _tf_alloc(int64_t* total, size_t size)
{
	void* ptr = malloc(size + sizeof(size_t));
	if (ptr)
	{
		__atomic_add_fetch(total, size, __ATOMIC_RELAXED);
		memcpy(ptr, &size, sizeof(size_t));
		return (void*)((intptr_t)ptr + sizeof(size_t));
	}
	else
	{
		return NULL;
	}
}

static void* _tf_realloc(int64_t* total, void* ptr, size_t size)
{
	void* old_ptr = ptr ? (void*)((intptr_t)ptr - sizeof(size_t)) : NULL;
	size_t old_size = 0;
	if (old_ptr)
	{
		memcpy(&old_size, old_ptr, sizeof(size_t));
	}
	void* new_ptr = NULL;
	if (old_ptr && !size)
	{
		free(old_ptr);
	}
	else
	{
		new_ptr = realloc(old_ptr, size + sizeof(size_t));
	}
	if (new_ptr)
	{
		__atomic_add_fetch(total, (int64_t)size - (int64_t)old_size, __ATOMIC_RELAXED);
		memcpy(new_ptr, &size, sizeof(size_t));
		return (void*)((intptr_t)new_ptr + sizeof(size_t));
	}
	else
	{
		__atomic_sub_fetch(total, old_size, __ATOMIC_RELAXED);
		return NULL;
	}
}

static void _tf_free(int64_t* total, void* ptr)
{
	if (ptr)
	{
		void* old_ptr = (void*)((intptr_t)ptr - sizeof(size_t));
		size_t size = 0;
		memcpy(&size, old_ptr, sizeof(size_t));
		__atomic_sub_fetch(total, size, __ATOMIC_RELAXED);
		free(old_ptr);
	}
}

static void* _tf_uv_alloc(size_t size)
{
	return _tf_alloc(&s_uv_malloc_size, size);
}

static void* _tf_uv_realloc(void* ptr, size_t size)
{
	return _tf_realloc(&s_uv_malloc_size, ptr, size);
}

static void* _tf_uv_calloc(size_t nmemb, size_t size)
{
	void* ptr = calloc(1, nmemb * size + sizeof(size_t));
	if (ptr)
	{
		size_t total_size = nmemb * size;
		__atomic_add_fetch(&s_uv_malloc_size, total_size, __ATOMIC_RELAXED);
		memcpy(ptr, &total_size, sizeof(size_t));
		return (void*)((intptr_t)ptr + sizeof(size_t));
	}
	else
	{
		return NULL;
	}
}

static void _tf_uv_free(void* ptr)
{
	_tf_free(&s_uv_malloc_size, ptr);
}

void tf_mem_replace_uv_allocator()
{
	uv_replace_allocator(_tf_uv_alloc, _tf_uv_realloc, _tf_uv_calloc, _tf_uv_free);
}

size_t tf_mem_get_uv_malloc_size()
{
	return s_uv_malloc_size;
}

#if !defined(_WIN32)
void* _tf_tls_alloc(size_t size, const char* file, int line)
{
	return _tf_alloc(&s_tls_malloc_size, size);
}

void* _tf_tls_realloc(void* ptr, size_t size, const char* file, int line)
{
	return _tf_realloc(&s_tls_malloc_size, ptr, size);
}

void _tf_tls_free(void* ptr, const char* file, int line)
{
	_tf_free(&s_tls_malloc_size, ptr);
}
#endif

void tf_mem_replace_tls_allocator()
{
#if !defined(_WIN32)
	CRYPTO_set_mem_functions(_tf_tls_alloc, _tf_tls_realloc, _tf_tls_free);
#endif
}

size_t tf_mem_get_tls_malloc_size()
{
	return s_tls_malloc_size;
}

void* tf_malloc(size_t size)
{
	return _tf_alloc(&s_tf_malloc_size, size);
}

void* tf_realloc(void* ptr, size_t size)
{
	return _tf_realloc(&s_tf_malloc_size, ptr, size);
}

void tf_free(void* ptr)
{
	_tf_free(&s_tf_malloc_size, ptr);
}

char* tf_strdup(const char* string)
{
	size_t len = strlen(string);
	char* buffer = tf_malloc(len + 1);
	memcpy(buffer, string, len + 1);
	return buffer;
}

void* tf_resize_vec(void* ptr, size_t size)
{
	void* alloc_ptr = ptr ? (void*)((intptr_t)ptr - sizeof(size_t)) : NULL;
	size_t alloc_size = 0;
	if (alloc_ptr)
	{
		memcpy(&alloc_size, alloc_ptr, sizeof(size_t));
	}

	if ((alloc_size >= 16 * size + sizeof(size_t)) || !size)
	{
		/* If we've dropped significantly in size or are freeing, resize down. */
		return tf_realloc(ptr, size);
	}
	else if (alloc_size >= size + sizeof(size_t))
	{
		/* Otherwise, if we're big enough, stay the same size. */
		return ptr;
	}
	else
	{
		/* If we need to grow, overallocate 2x to give room to continue growing. */
		return tf_realloc(ptr, size * 2);
	}
}

size_t tf_mem_get_tf_malloc_size()
{
	return s_tf_malloc_size;
}

static void* _tf_js_malloc(JSMallocState* state, size_t size)
{
	int64_t delta = 0;
	void* ptr = _tf_alloc(&delta, size);
	if (ptr)
	{
		__atomic_add_fetch(&s_js_malloc_size, delta, __ATOMIC_RELAXED);
		state->malloc_count++;
		state->malloc_size += delta;
	}
	return ptr;
}

static void _tf_js_free(JSMallocState* state, void* ptr)
{
	if (ptr)
	{
		int64_t delta = 0;
		_tf_free(&delta, ptr);
		__atomic_add_fetch(&s_js_malloc_size, delta, __ATOMIC_RELAXED);
		state->malloc_count--;
		state->malloc_size += delta;
	}
}

static void* _tf_js_realloc(JSMallocState* state, void* ptr, size_t size)
{
	int64_t delta = 0;
	void* result = _tf_realloc(&delta, ptr, size);
	__atomic_add_fetch(&s_js_malloc_size, delta, __ATOMIC_RELAXED);
	state->malloc_count += (ptr ? -1 : 0) + (result ? 1 : 0);
	state->malloc_size += delta;
	return result;
}

static size_t _tf_js_malloc_usable_size(const void* ptr)
{
	void* old_ptr = ptr ? (void*)((intptr_t)ptr - sizeof(size_t)) : NULL;
	size_t old_size = 0;
	if (old_ptr)
	{
		memcpy(&old_size, old_ptr, sizeof(size_t));
	}
	return old_size;
}

void tf_get_js_malloc_functions(JSMallocFunctions* out)
{
	*out = (JSMallocFunctions)
	{
		.js_malloc = _tf_js_malloc,
		.js_free = _tf_js_free,
		.js_realloc = _tf_js_realloc,
		.js_malloc_usable_size = _tf_js_malloc_usable_size,
	};
}

size_t tf_mem_get_js_malloc_size()
{
	return s_js_malloc_size;
}