#include "res.h"
#include "raygui.h"

typedef struct {
	map_t(Asset) assets;
	map_t(FunShader) sdr;
} Resources;

static Resources res = {0};

static void
asset_set_tag(Asset *asset, cstr path) {
	cstr ext = GetFileExtension(path);
	if (strcmp(ext, ".glb") == 0) {
		asset->tag = ASSET_TAG_MODEL;
	} else if (strcmp(ext, ".png") == 0) {
		asset->tag = ASSET_TAG_TEXTURE;
	} else if (strcmp(ext, ".ttf") == 0 ||
	           strcmp(ext, ".otf") == 0) {
		asset->tag = ASSET_TAG_FONT;
	} else if (strcmp(ext, ".ogg") == 0) {
		asset->tag = ASSET_TAG_SOUND;
	}
}

static bool
asset_load_data(cstr key) {
	Asset *asset = map_get(&res.assets, key);
	cstr ext = GetFileExtension(key);
	bool success = false;

	if (!strcmp(ext, ".glb")) {
		Model mdl = LoadModel(key);

		if (!IsModelValid(mdl)) {
			UnloadModel(mdl);
			return false;
		}

		bool skinned = mdl.skeleton.boneCount > 0 &&
			mdl.skeleton.bindPose != NULL;

		cstr sdr_key = skinned ? "skinned" : "default";

		for (int i = 0; i < mdl.materialCount; ++i) {
			mdl.materials[i].shader = shader_get(sdr_key)->sdr;
		}

		if (skinned) {
			int anim_count = 0;
			ModelAnimation *anims =
				LoadModelAnimations(key, &anim_count);

			if (anim_count > 0) {
				UpdateModelAnimation(mdl, anims[0], 0.0f);
			}

			UnloadModelAnimations(anims, anim_count);
		}

		for (int i = 0; i < mdl.materialCount; ++i) {
			mdl.materials[i].shader = shader_get(sdr_key)->sdr;
		}

		asset->tag = ASSET_TAG_MODEL;
		asset->mdl = mdl;
		success = true;
	}

	if (!strcmp(ext, ".png")) {
		Texture tex = LoadTexture(key);

		if (!IsTextureValid(tex)) {
			UnloadTexture(tex);
			return false;
		}

		GenTextureMipmaps(&tex);
		SetTextureWrap(tex, TEXTURE_WRAP_REPEAT);
		// SetTextureFilter(tex, TEXTURE_FILTER_ANISOTROPIC_16X);
		SetTextureFilter(tex, TEXTURE_FILTER_TRILINEAR);
		// SetTextureFilter(tex, TEXTURE_FILTER_BILINEAR);

		asset->tag = ASSET_TAG_TEXTURE;
		asset->tex = tex;
		success = true;
	}

	if (!strcmp(ext, ".ttf") || !strcmp(ext, ".otf")) {
		Font fnt = LoadFont(key);

		if (!IsFontValid(fnt)) {
			UnloadFont(fnt);
			return false;
		}

		asset->tag = ASSET_TAG_FONT;
		asset->fnt = fnt;
		success = true;
	}

	if (!strcmp(ext, ".ogg")) {
		Sound snd = LoadSound(key);

		if (!IsSoundValid(snd)) {
			UnloadSound(snd);
			return false;
		}

		asset->tag = ASSET_TAG_SOUND;
		asset->snd = snd;
		success = true;
	}

	if (!success) {
		TraceLog(LOG_WARNING, "'%s' no valid file ext found.", key);
		return false;
	}

	asset->loaded = true;
	return true;
}

static void
asset_unload(cstr key) {
	Asset *asset = map_get(&res.assets, key);
	if (!asset) {
		TraceLog(LOG_WARNING,
		         "asset_unload: invalid key '%s'\n", key);
		return;
	}

	if (!asset->loaded) {
		TraceLog(LOG_INFO,
		         "asset_unload: nothing to unload '%s'", key);
		return;
	}

	free(asset->path);
	asset->path_strlen = 0;

	switch(asset->tag) {
	case ASSET_TAG_MODEL:
		UnloadModel(asset->mdl);
	break;
	case ASSET_TAG_TEXTURE:
		UnloadTexture(asset->tex);
	break;
	case ASSET_TAG_SOUND:
		UnloadSound(asset->snd);
	break;
	case ASSET_TAG_FONT:
		UnloadFont(asset->fnt);
	break;
	case ASSET_TAG_INVALID:
		TraceLog(LOG_WARNING,
		         "asset_unload: invalid tag '%s'\n", key);
	break;
	}
}


void
res_init(void) {
	map_init(&res.sdr);
	map_init(&res.assets);

	fun_shader_create("default",
	                "res/sdr/mdl_vtx_lit.vs",
	        	"res/sdr/mdl_vtx_lit.fs");

	fun_shader_create("skinned",
	                "res/sdr/mdl_skinned_vtx_lit.vs",
	        	"res/sdr/mdl_vtx_lit.fs");

	fun_shader_create("alpha_cutoff",
			"res/sdr/mdl_vtx_lit.vs",
			"res/sdr/mdl_vtx_lit_cutoff.fs");

	fun_shader_create("alpha_cutoff_nocull",
			"res/sdr/mdl_vtx_lit.vs",
			"res/sdr/mdl_two_sided_cutoff.fs");

	fun_shader_create("foliage",
			"res/sdr/mdl_foliage.vs",
			"res/sdr/mdl_vtx_lit_cutoff.fs");

	fun_shader_create("sky", NULL, "res/sdr/sky.fs");

	cstr filter = ".glb;.png;.ttf;.otf;.ogg";
	FilePathList res_list = LoadDirectoryFilesEx("res", filter, true);

	for (int i = 0; i < res_list.count; ++i) {
		cstr path = res_list.paths[i];
		Asset asset = {0};
		asset_set_tag(&asset, path);
		asset.path_strlen = strlen(path);
		asset.path = strdup(path); // freed in asset_unload
		map_set(&res.assets, asset.path, asset);
	}

	UnloadDirectoryFiles(res_list);
}

void
res_exit(void) {
	cstr key;
	map_iter_t iter;

	iter = map_iter(&res.assets);
	while ((key = map_next(&res.assets, &iter))) {
		asset_unload(key);
	}

	iter = map_iter(&res.sdr);
	while ((key = map_next(&res.sdr, &iter))) {
		FunShader *sdr = map_get(&res.sdr, key);
		if (!sdr) {
			TraceLog(LOG_ERROR, "res_exit: '%s' invalid\n", key);
		}
		fun_shader_delete(sdr);
	}

	map_deinit(&res.assets);
	map_deinit(&res.sdr);
}

void
res_update(void) {
}

char**
asset_type_keys(AssetTag tag, int *count) {
	*count = 0;

	if (!res.assets.base.nnodes) {
		return NULL;
	}

	char **keys = malloc(res.assets.base.nnodes * sizeof(*keys));
	if (!keys) {
		perror("asset_type_keys: mem allocation failed.");
		exit(1);
	}

	map_iter_t iter = map_iter(&res.assets);
	cstr key;

	while ((key = map_next(&res.assets, &iter))) {
		Asset *asset = map_get(&res.assets, key);
		if (asset->tag == tag) {
			keys[(*count)++] = asset->path;
		}
	}

	return keys;
}

Asset*
asset_get(cstr key) {
	Asset *asset = map_get(&res.assets, key);
	if (!asset) {
		TraceLog(LOG_ERROR, "asset '%s' not found.\n", key);
		return NULL;
	}

	if (!asset->loaded) {
		if (!asset_load_data(key)) {
			TraceLog(LOG_ERROR, "asset '%s' couldn't load.", key);
			return NULL;
		}
	}

	return asset;
}

FunShader*
shader_get(cstr key) {
	FunShader *sdr = map_get(&res.sdr, key);
	if (!sdr) {
		TraceLog(LOG_ERROR, "shader '%s' not found.", key);
		return NULL;
	}
	return sdr;
}

void
fun_shader_create(cstr key, cstr vpath, cstr fpath) {
	FunShader shader = {0};
	shader.sdr = LoadShader(vpath, fpath);
	if (!IsShaderValid(shader.sdr)) {
		printf("err: fun_shader_create('%s', '%s', '%s')\n",
			key, vpath, fpath);
		exit(1);
	}
	map_init(&shader.uniforms);
	map_set(&res.sdr, key, shader);
}

void
fun_shader_delete(FunShader *shader) {
	UnloadShader(shader->sdr);
	map_deinit(&shader->uniforms);
}

int
fun_shader_uniform(cstr skey, cstr ukey) {
	FunShader *sdr = shader_get(skey);

	int *ret = map_get(&sdr->uniforms, ukey);
	if (ret) {
		return *ret;
	}

	// BeginShaderMode(sdr->sdr);
	int loc = rlGetLocationUniform(sdr->sdr.id, ukey);
	if (loc < 0) {
		TraceLog(LOG_ERROR,
		         "shader '%s' uniform '%s' not found.\n", skey, ukey);
		exit(1);
	}
	// EndShaderMode();

	map_set(&sdr->uniforms, ukey, loc);
	ret = map_get(&sdr->uniforms, ukey);
	if (!ret) {
		TraceLog(LOG_ERROR,
		         "shader '%s' uniform '%s' not found", skey, ukey);
		return -1;
	}

	return *ret;
}