#include "asset-manager.hpp"

#include "embedded-shaders.h"
#include "pk.h"
#include "thread-pool.hpp"

#include <cstring>
#include <filesystem>
#include <fstream>
#include <future>

struct Asset_Master {
	pk_membucket *bkt;
	pk_bkt_arr_t<Asset> bc{};
	ThreadPoolHandle thread_pool = ThreadPoolHandle_MAX;
} asset_mstr;

void AM_Init() {
	int i;
	asset_mstr.bkt = pk_mem_bucket_create("pk_bkt_arr AM", 1024 * 1024, PK_MEMBUCKET_FLAG_NONE);
	new (&asset_mstr.bc) pk_bkt_arr_t<Asset>{ pk_bkt_arr_handle_MAX_constexpr, asset_mstr.bkt, asset_mstr.bkt };
	asset_mstr.thread_pool = PkeThreads_Init(2, 255);
	for (i = 0; i < embedded_shader_index_count; ++i) {
		AM_Register_Static(embedded_shaders[i].name, PKE_ASSET_TYPE_SHADER, embedded_shaders[i].data, embedded_shaders[i].size, nullptr);
	}
}

void AM_Load_Task(Asset &asset) {
	std::ifstream file(asset.basePath, std::ios::ate | std::ios::binary);
	if (!file.is_open()) {
		asset.state = PKE_ASSET_LOADING_STATE_FAILED;
		return;
	}
	asset.size = std::filesystem::file_size(asset.basePath);
	void *ptr = pk_new_base(asset.size, alignof(max_align_t));
	file.seekg(0, std::ios::beg);
	file.read(static_cast<char *>(ptr), asset.size);
	file.close();
	asset.ptr = ptr;
	asset.state = PKE_ASSET_LOADING_STATE_LOADED;
}

inline Asset *AM_Get_Inner(const char (&key)[AssetKeyLength]) {
	auto asset_find_cb = [](void *user_data, const void *user_obj_data, const void *arr_obj_data) {
		(void)user_data;
		const char (&inner_key)[16] = *reinterpret_cast<const char(*&)[16]>(user_obj_data);
		const Asset *asset = reinterpret_cast<const Asset *>(arr_obj_data);
		return strncmp(inner_key, asset->key, 16) == 0;
	};
	AssetHandle handle { pk_bkt_arr_find_first_handle(&asset_mstr.bc, asset_find_cb, NULL, &key[0]) };
	if (handle == AssetHandle_MAX) return nullptr;
	return &asset_mstr.bc[handle];
}

AssetHandle AM_Register(const AssetKey &key, AssetType type, const void *data, int64_t size, std::size_t alignment, union pke_asset_details *details) {
	assert(data != nullptr && "Attempt to register invalid asset data");
	assert(data != CAFE_BABE(void) && "Attempt to register invalid asset data");
	assert(size != 0 && "Attempt to register asset data of size 0");
	Asset *searchedAsset = AM_Get_Inner(key);
	if (searchedAsset != nullptr) {
		return searchedAsset->handle;
	}
	AssetHandle assetHandle{pk_bkt_arr_new_handle(&asset_mstr.bc)};
	Asset &asset = asset_mstr.bc[assetHandle];
	new (&asset) Asset{};
	memset(&asset.details, 0, sizeof(union pke_asset_details));
	if (details != nullptr) asset.details = *details;
	asset.handle = assetHandle;
	strncpy(asset.key, key, AssetKeyLength);
	asset.basePath = nullptr;
	asset.size = size;
	void *ptr = pk_new_base(size, alignment);
	memcpy(ptr, data, size);
	asset.ptr = ptr;
	asset.state = PKE_ASSET_LOADING_STATE_LOADED;
	asset.type = type;
	return assetHandle;
}

AssetHandle AM_Register(const AssetKey &key, AssetType type, const char *path, union pke_asset_details *details) {
	Asset *searchedAsset = AM_Get_Inner(key);
	if (searchedAsset != nullptr) {
		return searchedAsset->handle;
	}
	AssetHandle assetHandle{pk_bkt_arr_new_handle(&asset_mstr.bc)};
	Asset &asset = asset_mstr.bc[assetHandle];
	new (&asset) Asset{};
	memset(&asset.details, 0, sizeof(union pke_asset_details));
	if (details != nullptr) asset.details = *details;
	asset.handle = assetHandle;
	strncpy(asset.key, key, AssetKeyLength);
	int64_t pathLen = strlen(path);
	auto *copiedPath = pk_new_arr<char>(pathLen + 1);
	copiedPath[pathLen] = '\0';
	strncpy(copiedPath, path, pathLen);
	asset.basePath = copiedPath;
	asset.state = PKE_ASSET_LOADING_STATE_LOADING;
	asset.type = type;
	std::packaged_task<void()> *task = pk_new<std::packaged_task<void()>>();
	new (task) std::packaged_task<void()>( [&asset] {
		AM_Load_Task(asset);
	});
	asset.future = task->get_future();
	PkeThreads_Enqueue(asset_mstr.thread_pool, task);
	return assetHandle;
}

AssetHandle AM_Register(const char *path, AssetType type, union pke_asset_details *details) {
	NULL_CHAR_ARR(assetKey, AssetKeyLength);
	int64_t pathLen = strlen(path);
	int64_t pathOffset = (pathLen > AssetKeyLength ? pathLen - AssetKeyLength : 0);
	strncpy(assetKey, path + pathOffset, AssetKeyLength);

	Asset *searchedAsset = AM_Get_Inner(assetKey);
	if (searchedAsset != nullptr) {
		return searchedAsset->handle;
	}

	AssetHandle assetHandle{pk_bkt_arr_new_handle(&asset_mstr.bc)};
	Asset &asset = asset_mstr.bc[assetHandle];
	new (&asset) Asset{};
	memset(&asset.details, 0, sizeof(union pke_asset_details));
	if (details != nullptr) asset.details = *details;
	asset.handle = assetHandle;
	strncpy(asset.key, assetKey, AssetKeyLength);
	auto *copiedPath = pk_new_arr<char>(pathLen + 1);
	copiedPath[pathLen] = '\0';
	strncpy(copiedPath, path, pathLen);
	asset.basePath = copiedPath;
	asset.state = PKE_ASSET_LOADING_STATE_LOADING;
	asset.type = type;
	std::packaged_task<void()> *task = pk_new<std::packaged_task<void()>>();
	new (task) std::packaged_task<void()>( [&asset] {
		AM_Load_Task(asset);
	});
	asset.future = task->get_future();
	PkeThreads_Enqueue(asset_mstr.thread_pool, task);
	return assetHandle;
}

AssetHandle AM_Register_Static(const AssetKey &key, AssetType type, const void *data, int64_t size, union pke_asset_details *details) {
	assert(data != nullptr && "Attempt to register invalid asset data");
	assert(data != CAFE_BABE(void) && "Attempt to register invalid asset data");
	assert(size != 0 && "Attempt to register asset data of size 0");
	Asset *searchedAsset = AM_Get_Inner(key);
	if (searchedAsset != nullptr) {
		return searchedAsset->handle;
	}
	AssetHandle assetHandle{pk_bkt_arr_new_handle(&asset_mstr.bc)};
	Asset &asset = asset_mstr.bc[assetHandle];
	new (&asset) Asset{};
	memset(&asset.details, 0, sizeof(union pke_asset_details));
	if (details != nullptr) asset.details = *details;
	asset.handle = assetHandle;
	strncpy(asset.key, key, AssetKeyLength);
	asset.basePath = nullptr;
	asset.size = size;
	asset.ptr = data;
	asset.state = PKE_ASSET_LOADING_STATE_LOADED;
	asset.type = type;
	asset.flags = PKE_ASSET_FLAGS_MEM_STATIC;
	return assetHandle;
}

void AM_Release(AssetHandle assetHandle) {
	assert(assetHandle != AssetHandle_MAX);
	Asset &asset = asset_mstr.bc[assetHandle];
	asset.referenceCount -= 1;
	assert(asset.referenceCount >= 0 && "asset was unloaded more times than it was retrieved");
}

const Asset *AM_Get(AssetHandle assetHandle) {
	auto validationResult = pk_bkt_arr_handle_validate(&asset_mstr.bc, assetHandle);
	assert(validationResult == PK_BKT_ARR_HANDLE_VALIDATION_VALID);
	auto &asset = asset_mstr.bc[assetHandle];
	if (asset.state == PKE_ASSET_LOADING_STATE_LOADED) {
		asset.referenceCount += 1;
		return &asset;
	}
	if (asset.state == PKE_ASSET_LOADING_STATE_UNLOADED) {
		if (asset.basePath == nullptr) {
			return nullptr;
		}
		asset.state = PKE_ASSET_LOADING_STATE_LOADING;
		std::packaged_task<void()> *task = pk_new<std::packaged_task<void()>>();
		new (task) std::packaged_task<void()>( [&asset] {
			AM_Load_Task(asset);
		});
		asset.future = task->get_future();
		PkeThreads_Enqueue(asset_mstr.thread_pool, task);
	}
	if (asset.state == PKE_ASSET_LOADING_STATE_LOADING) {
		if (asset.future.valid()) {
			asset.future.wait();
		} else {
			char buf[256];
			buf[255] = '\0';
			snprintf(buf, 255, "[AM_Get] Attempting to retrieve asset: '%.16s' which had a state of '%hhu', but did not have a valid future and thus cannot wait for completion", asset.key, static_cast<AssetLoadingState_T>(asset.state));
			throw buf;
		}
	}
	if (asset.state == PKE_ASSET_LOADING_STATE_LOADED)
		asset.referenceCount += 1;
	return &asset;
}

const AssetHandle AM_GetHandle(const AssetKey &key) {
	Asset *searchedAsset = AM_Get_Inner(key);
	if (searchedAsset != nullptr) {
		return searchedAsset->handle;
	}
	return AssetHandle_MAX;
}

pk_bkt_arr *AM_GetAssets() {
	return &asset_mstr.bc;
}

void AM_DebugPrint() {
	bool b;
	pk_iter_t<Asset> iter_asset{};
	fprintf(stdout, "Asset Manager printout:\n");
	b = pk_bkt_arr_iter_begin(&asset_mstr.bc, &iter_asset);
	while (b == true) {
		Asset &asset = *iter_asset;
		printf("-Asset: 0x%.08X 0x%.08X\n", asset.handle.b, asset.handle.i);
		printf("\tkey:  %.16s\n", asset.key);
		if (asset.basePath != nullptr) {
			printf("\tbasePath:  %s\n", asset.basePath);
		} else {
			printf("\tbasePath:  %p\n", (void *)nullptr);
		}
		printf("\tsize: %ld\n", asset.size);
		printf("\tptr   %p\n", asset.ptr);
		printf("\tfuture: %i\n", asset.future.valid());
		printf("\treferenceCount: %i\n", asset.referenceCount);
		printf("\tAssetLoadingState: %hhu\n", static_cast<AssetLoadingState_T>(asset.state));
		b = pk_bkt_arr_iter_increment(&asset_mstr.bc, &iter_asset);
	}
}

void AM_GC() {
	bool b;
	pk_iter_t<Asset> iter_asset{};
	b = pk_bkt_arr_iter_begin(&asset_mstr.bc, &iter_asset);
	while (b == true) {
		Asset &asset = *iter_asset;
		if (PK_HAS_FLAG(asset.flags, PKE_ASSET_FLAGS_MEM_STATIC)) {
			fprintf(stdout, "[AM_GC] Asset '%.16s' is static, skipping.\n", asset.key);
			b = pk_bkt_arr_iter_increment(&asset_mstr.bc, &iter_asset);
			continue;
		}
		switch (asset.state) {
			case PKE_ASSET_LOADING_STATE_LOADING:
				fprintf(stdout, "[AM_GC] Asset '%.16s' is still loading.\n", asset.key);
				b = pk_bkt_arr_iter_increment(&asset_mstr.bc, &iter_asset);
				continue;
			default:
				void(0);
		}
		if (asset.referenceCount > 0) {
			fprintf(stdout, "[AM_GC] Asset '%.16s' still in use, count: %i\n", asset.key, asset.referenceCount);
			b = pk_bkt_arr_iter_increment(&asset_mstr.bc, &iter_asset);
			continue;
		}
		if (asset.ptr != nullptr && asset.ptr != CAFE_BABE(void)) {
			pk_delete_base(asset.ptr, asset.size);
			fprintf(stdout, "[AM_GC] Asset '%.16s' not in use, freeing.\n", asset.key);
		}
		asset.size = 0;
		asset.ptr = CAFE_BABE(void);
		new (&asset.future) std::future<void>{};
		asset.state = PKE_ASSET_LOADING_STATE_UNLOADED;
		b = pk_bkt_arr_iter_increment(&asset_mstr.bc, &iter_asset);
	}
}

void AM_Teardown() {
	PkeThreads_Teardown(asset_mstr.thread_pool);
	AM_GC();
	pk_bkt_arr_teardown(&asset_mstr.bc);
	pk_mem_bucket_destroy(asset_mstr.bkt);
}