#include "asset-manager.hpp"

#include "bucketed-array.hpp"
#include "thread-pool.hpp"

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

const pk_handle_item_index_T maxAssetItemsPerBucket = 64;

BucketContainer<Asset, AssetHandle> Asset_BucketContainer{};

ThreadPoolHandle assetThreadPool = ThreadPoolHandle_MAX;

AssetKey EngineDefinedAssets[EngineDefinedAssetCount] = {
	"pke_prsnt_vrt\0\0",
	"pke_prsnt_frg\0\0",
	"pke_txtr_vrt\0\0\0",
	"pke_txtr_frg\0\0\0",
	"pke_glyph_vrt\0\0",
	"pke_glyph_frg\0\0",
};

void AM_Init() {
	Buckets_Init(Asset_BucketContainer, maxAssetItemsPerBucket);
	assetThreadPool = PkeThreads_Init(2, 255);
	AM_Register(EngineDefinedAssets[0], PKE_ASSET_TYPE_SHADER, "assets/shaders/present.vert.spv");
	AM_Register(EngineDefinedAssets[1], PKE_ASSET_TYPE_SHADER, "assets/shaders/present.frag.spv");
	AM_Register(EngineDefinedAssets[2], PKE_ASSET_TYPE_SHADER, "assets/shaders/vertex.vert.spv");
	AM_Register(EngineDefinedAssets[3], PKE_ASSET_TYPE_SHADER, "assets/shaders/texture.frag.spv");
	AM_Register(EngineDefinedAssets[4], PKE_ASSET_TYPE_SHADER, "assets/shaders/glyph.vert.spv");
	AM_Register(EngineDefinedAssets[5], PKE_ASSET_TYPE_SHADER, "assets/shaders/glyph.frag.spv");
}

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);
	asset.ptr = pk_new_base(asset.size, 64);
	file.seekg(0, std::ios::beg);
	file.read(static_cast<char *>(asset.ptr), asset.size);
	file.close();
	asset.state = PKE_ASSET_LOADING_STATE_LOADED;
}

inline Asset *AM_Get_Inner(AssetKey key) {
	for (pk_handle_bucket_index_T b = 0; b <= Asset_BucketContainer.pkeHandle.bucketIndex; ++b) {
		pk_handle_item_index_T count = 0;
		if (b == Asset_BucketContainer.pkeHandle.bucketIndex) {
			count = Asset_BucketContainer.pkeHandle.itemIndex;
		} else {
			count = maxAssetItemsPerBucket;
		}
		for (pk_handle_item_index_T i = 0; i < count; ++i) {
			if (strncmp(key, Asset_BucketContainer.buckets[b][i].key, 16) == 0) {
				return &Asset_BucketContainer.buckets[b][i];
			}
		}
	}
	return nullptr;
}

AssetHandle AM_Register(AssetKey key, AssetType type, const void *data, int64_t size, std::size_t alignment) {
	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{Buckets_NewHandle(Asset_BucketContainer)};
	Asset &asset = Asset_BucketContainer.buckets[assetHandle.bucketIndex][assetHandle.itemIndex];
	new (&asset) Asset{};
	asset.handle = assetHandle;
	strncpy(asset.key, key, AssetKeyLength);
	asset.basePath = nullptr;
	asset.size = size;
	asset.ptr = pk_new_base(size, alignment);
	memcpy(asset.ptr, data, size);
	asset.state = PKE_ASSET_LOADING_STATE_LOADED;
	asset.type = type;
	return assetHandle;
}

AssetHandle AM_Register(AssetKey key, AssetType type, const char *path) {
	Asset *searchedAsset = AM_Get_Inner(key);
	if (searchedAsset != nullptr) {
		return searchedAsset->handle;
	}
	AssetHandle assetHandle{Buckets_NewHandle(Asset_BucketContainer)};
	Asset &asset = Asset_BucketContainer.buckets[assetHandle.bucketIndex][assetHandle.itemIndex];
	new (&asset) Asset{};
	asset.handle = assetHandle;
	strncpy(asset.key, key, AssetKeyLength);
	int64_t pathLen = strlen(path);
	auto *copiedPath = pk_new<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(assetThreadPool, task);
	return assetHandle;
}

AssetHandle AM_Register(const char *path, AssetType type) {
	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{Buckets_NewHandle(Asset_BucketContainer)};
	Asset &asset = Asset_BucketContainer.buckets[assetHandle.bucketIndex][assetHandle.itemIndex];
	new (&asset) Asset{};
	asset.handle = assetHandle;
	strncpy(asset.key, assetKey, AssetKeyLength);
	auto *copiedPath = pk_new<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(assetThreadPool, task);
	return assetHandle;
}

void AM_Release(AssetHandle assetHandle) {
	assert(assetHandle != AssetHandle_MAX);
	Asset &asset = Asset_BucketContainer.buckets[assetHandle.bucketIndex][assetHandle.itemIndex];
	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_handle_validate(assetHandle, Asset_BucketContainer.pkeHandle, maxAssetItemsPerBucket);
	assert(validationResult == 0);
	auto &asset = Asset_BucketContainer.buckets[assetHandle.bucketIndex][assetHandle.itemIndex];
	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(assetThreadPool, 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(AssetKey key) {
	Asset *searchedAsset = AM_Get_Inner(key);
	if (searchedAsset != nullptr) {
		return searchedAsset->handle;
	}
	return AssetHandle_MAX;
}

uint64_t AM_GetBucketCount() {
	return Asset_BucketContainer.pkeHandle.bucketIndex + 1;
}

Asset *AM_GetAssets(pk_handle_bucket_index_T bucketIndex, pk_handle_item_index_T &itemCount) {
	if (bucketIndex == Asset_BucketContainer.pkeHandle.bucketIndex) {
		itemCount = Asset_BucketContainer.pkeHandle.itemIndex;
	} else {
		itemCount = maxAssetItemsPerBucket;
	}
	return Asset_BucketContainer.buckets[bucketIndex];
}

void AM_DebugPrint() {
	printf("Asset Manager printout:\n");
	for (uint64_t b = 0; b <= Asset_BucketContainer.pkeHandle.bucketIndex; ++b) {
		auto &bkt = Asset_BucketContainer.buckets[b];
		uint64_t counter = b == Asset_BucketContainer.pkeHandle.bucketIndex ? Asset_BucketContainer.pkeHandle.itemIndex : maxAssetItemsPerBucket;
		for (uint64_t i = 0; i < counter; ++i) {
			auto &asset = bkt[i];
			/*
			if (asset.size == 0)
				continue;
			*/
			printf("-Asset: 0x%016lX\n", (b << 32) + 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));
		}
	}
}

void AM_GC() {
	for (long b = 0; b <= Asset_BucketContainer.pkeHandle.bucketIndex; ++b) {
		for (long i = 0; i < Asset_BucketContainer.pkeHandle.itemIndex; ++i) {
			Asset &asset = Asset_BucketContainer.buckets[b][i];
			if (asset.referenceCount > 0)
				continue;
			if (asset.ptr != nullptr && asset.ptr != CAFE_BABE(void)) {
				pk_delete_base(asset.ptr, asset.size);
			}
			asset.size = 0;
			asset.ptr = CAFE_BABE(void);
			asset.future = std::future<void>{};
			asset.state = PKE_ASSET_LOADING_STATE_UNLOADED;
		}
	}
}

void AM_Teardown() {
	AM_GC();
	Buckets_Destroy(Asset_BucketContainer);
	PkeThreads_Teardown(assetThreadPool);
}