
U32 D_ImageHandle(D_Context *draw, Str8 name) {
	U32 result = 0;

	U64 hash  = Str8_Hash(name);
	U64 index = hash & (D_ASSET_HASH_COUNT - 1);

	D_AssetHash *lookup = draw->lookup[index];
	while (lookup != 0) {
		if (lookup->hash == hash && Str8_Equal(lookup->value, name, 0)) {
			result = lookup->id;
			break;
		}

		lookup = lookup->next;
	}

	return result;
}


void D_AnimationInit(D_Animation *a, U32 id, U32 rows, U32 cols, F32 time) {
	a->id = id;

	a->rows = rows;
	a->cols = cols;

	a->time = 0;
	a->frame_time = time;

	a->frame = V2F(1.0f / (F32) cols, 1.0f / (F32) rows);

	a->index = 0;
}

R2f D_AnimationFrame(D_Animation *a) {
	R2f result = { 0 };

	U32 row = a->index / a->cols;
	U32 col = a->index % a->cols;

	result.min = V2F(a->frame.w * col, a->frame.h * row);
	result.max = V2F(result.min.x + a->frame.w, result.min.y + a->frame.h);
	return result;
}

void D_AnimationUpdate(D_Animation *a, F32 dt) {
	a->time += dt;
	if (a->time >= a->frame_time) {
		a->time   = 0;
		a->index += 1;

		if (a->index >= (a->rows * a->cols)) { a->index = 0; }
	}
}

void D_Begin(D_Context *draw, Vk_Frame *frame, U32 max_rects) {
	Vk_Buffer *rbo   = &frame->rbo;

	draw->rbo       = rbo;
	draw->n_rects   = 0;
	draw->max_rects = max_rects;
	draw->rects     = rbo->data;
}

void D_End(D_Context *draw, Vk_Frame *frame) {
	VkCommandBuffer cmd = frame->cmd;
	Vk_Pipeline *basic  = &draw->pipelines[0];

	// We can probably stop doing this at some point
	VkDescriptorSet set;
	VkDescriptorSetAllocateInfo alloc_info = { 0 };
	alloc_info.sType              = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
	alloc_info.descriptorPool     = frame->descriptors;
	alloc_info.descriptorSetCount = 1;
	alloc_info.pSetLayouts        = &basic->layout.set;

	vk.AllocateDescriptorSets(vk.device, &alloc_info, &set);

	// 'update' the descriptor sets for binding
	M_TempScope(0, 0) {
		VkWriteDescriptorSet writes[2] = { 0 };

		VkDescriptorBufferInfo rbo_info = { 0 };
		rbo_info.buffer = draw->rbo->handle;
		rbo_info.offset = 0;
		rbo_info.range  = VK_WHOLE_SIZE;

		U32 total = draw->n_images + draw->n_fonts;
		VkDescriptorImageInfo *image_info = M_ArenaPush(temp.arena, VkDescriptorImageInfo, .count = total);

		for (U32 it = 0; it < draw->n_images; ++it) {
			image_info[it].imageView   = draw->images[it].image.view;
			image_info[it].sampler     = vk.sampler;
			image_info[it].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
		}

		U32 idx = draw->n_images;
		for (D_Font *it = draw->fonts; it != 0; it = it->next) {
			image_info[idx].imageView   = it->image.view;
			image_info[idx].sampler     = vk.sampler; // @Todo: probably want linear filtering on fonts
			image_info[idx].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

			idx += 1;
		}

		writes[0].sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		writes[0].dstSet          = set;
		writes[0].dstBinding      = 0;
		writes[0].descriptorCount = 1;
		writes[0].descriptorType  = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
		writes[0].pBufferInfo     = &rbo_info;

		writes[1].sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
		writes[1].dstSet          = set;
		writes[1].dstBinding      = 1;
		writes[1].descriptorCount = total;
		writes[1].descriptorType  = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
		writes[1].pImageInfo      = image_info;

		vk.UpdateDescriptorSets(vk.device, ArraySize(writes), writes, 0, 0);
	}

	vk.CmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, basic->handle);
	vk.CmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, basic->layout.pipeline, 0, 1, &set, 0, 0);
	vk.CmdPushConstants(cmd, basic->layout.pipeline, VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(Mat4x4F), &draw->camera->proj.fwd);

	VkViewport viewport = { 0, 0, (F32) draw->window_width, (F32) draw->window_height, 0.0f, 1.0f };
	VkRect2D   scissor  = { 0, 0, draw->window_width, draw->window_height };

	vk.CmdSetViewport(cmd, 0, 1, &viewport);
	vk.CmdSetScissor(cmd, 0, 1, &scissor);

	vk.CmdDraw(cmd, 6, draw->n_rects, 0, 0);
}

internal U32 V4f_UnormColour(V4f c) {
	// @Todo: SRGB handling
	U32 result =
		((U8) (255.0f * c.a)) << 24 |
		((U8) (255.0f * c.r)) << 16 |
		((U8) (255.0f * c.g)) <<  8 |
		((U8) (255.0f * c.b)) <<  0;

	return result;
}

void _D_Rect(D_Context *draw, D_RectOpts *opts) {
	if (draw->n_rects < draw->max_rects) {
		D_Rect *rect = &draw->rects[draw->n_rects];

		rect->texture = opts->texture;

		F32 w, h;

		rect->x = opts->p.x;
		rect->y = opts->p.y;

		rect->uv[0] = opts->uv.min.x;
		rect->uv[1] = opts->uv.min.y;
		rect->uv[2] = opts->uv.max.x;
		rect->uv[3] = opts->uv.max.y;

		rect->angle = opts->angle;

		if (opts->flags & D_RECT_PER_VERTEX_COLOUR) {
			rect->c[0] = V4f_UnormColour(opts->vtxc[0]);
			rect->c[1] = V4f_UnormColour(opts->vtxc[1]);
			rect->c[2] = V4f_UnormColour(opts->vtxc[2]);
			rect->c[3] = V4f_UnormColour(opts->vtxc[3]);
		}
		else {
			U32 unorm = V4f_UnormColour(opts->c);

			rect->c[0] = unorm;
			rect->c[1] = unorm;
			rect->c[2] = unorm;
			rect->c[3] = unorm;
		}

		if (opts->flags & D_RECT_IGNORE_ASPECT) {
			w = opts->dim.w;
			h = opts->dim.h;
		}
		else {
			Vk_Image *image = 0;
			if (opts->texture >= draw->n_images) {
				// @hack: first font texture
				//
				image = &draw->fonts->image;
			}
			else {
				image = &draw->images[opts->texture].image;
			}

			F32 width  = cast(F32) image->width;
			F32 height = cast(F32) image->height;

			if (opts->flags & D_RECT_UV_ASPECT) {
				width  *= (opts->uv.max.x - opts->uv.min.x);
				height *= (opts->uv.max.y - opts->uv.min.y);
			}

			F32 aspect_w = (width > height) ? 1.0f : (width / height);
			F32 aspect_h = (width > height) ? (height / width) : 1.0f;

			w = opts->scale * aspect_w;
			h = opts->scale * aspect_h;
		}

		if (opts->flags & D_RECT_FLIP_X) { w *= -1.0f; }
		if (opts->flags & D_RECT_FLIP_Y) { h *= -1.0f; }

		rect->w = w;
		rect->h = h;

		draw->n_rects += 1;
	}
}

void D_Text(D_Context *draw, D_Font *font, Str8 text, F32 x, F32 y) {
	F32 xoff = x;

	// @Todo: control scale
	F32 scale = 1.0f / (font->ascent - font->descent);

	for (S64 it = 0; it < text.count; ++it) {
		if (text.data[it] >= ' ' && text.data[it] <= '~') {
			D_Glyph *glyph = &font->glyphs[text.data[it] - ' '];

			F32 image_scale = scale * Max(glyph->dim.w, glyph->dim.h);
			F32 advance     = scale * glyph->advance;

			V2f offset = V2f_Scale(glyph->offset, scale);

			D_Rect(draw, xoff + offset.x, y + offset.y, .texture = font->id, .uv = glyph->box, .scale = image_scale, .flags = D_RECT_UV_ASPECT);
			xoff += advance;
		}
	}
}

void D_FontLoad(D_Context *draw, Str8 name, F32 size) {
	M_TempScope(0, 0) {
		(void) draw;

		Str8 exe_path = FS_SystemPath(temp.arena, FS_SYSTEM_PATH_EXE);
		Str8 path = Sf(temp.arena, "%.*s/assets/%.*s.ttf", Sv(exe_path), Sv(name));

		Str8 font_data = FS_ReadEntireFile(temp.arena, path);

		stbtt_fontinfo info = { 0 };
		stbtt_InitFont(&info, font_data.data, 0);

		F32 scale = stbtt_ScaleForPixelHeight(&info, 256);

		U32 w = 512;
		U32 h = 512;

		U8 *pixels = M_ArenaPush(temp.arena, U8, .count = (w * h));

		stbtt_pack_context pack = { 0 };
		stbtt_PackBegin(&pack, pixels, w, h, 0, 1, 0);

		U32 count = '~' - ' ';
		stbtt_packedchar *packed = M_ArenaPush(temp.arena, stbtt_packedchar, .count = count);
		stbtt_PackFontRange(&pack, font_data.data, 0, size, ' ', count, packed);

		D_Font *font = M_ArenaPush(draw->arena, D_Font);

		S32 asc, desc, line;
		stbtt_GetFontVMetrics(&info, &asc, &desc, &line);

		font->line_advance = cast(F32) (scale * line);
		font->ascent       = cast(F32) (scale * asc);
		font->descent      = cast(F32) (scale * desc);

		font->glyphs = M_ArenaPush(draw->arena, D_Glyph, .count = count);

		for (U32 it = 0; it < count; ++it) {
			D_Glyph *glyph = &font->glyphs[it];
			stbtt_packedchar *chr = &packed[it];

			S32 left, width;
			stbtt_GetCodepointHMetrics(&info, ' ' + it, &width, &left);

			glyph->box.min = V2F(chr->x0 / (F32) w, chr->y0 / (F32) h);
			glyph->box.max = V2F(chr->x1 / (F32) w, chr->y1 / (F32) h);

			glyph->advance  = chr->xadvance;
			glyph->dim.w    = cast(F32) (chr->x1 - chr->x0);
			glyph->dim.h    = cast(F32) (chr->y1 - chr->y0);
			glyph->offset.x = (0.5f * glyph->dim.w) + chr->xoff;
			glyph->offset.y = (0.5f * glyph->dim.h) + chr->yoff;
		}

		Vk_Buffer *staging = &draw->staging;
		U32 *px = (U32 *) staging->data;
		for (U32 y = 0; y < h; ++y) {
			for (U32 x = 0; x < w; ++x) {
				*px++ = 0x0 | ((U32) pixels[(y * w) + x] << 24);
			}
		}

		Vk_CommandBuffer *cmds = Vk_CommandBufferPush();

		font->id = draw->n_images + draw->n_fonts;

		font->image.width  = w;
		font->image.height = h;

		font->image.format = VK_FORMAT_R8G8B8A8_SRGB;
		font->image.usage  = VK_IMAGE_USAGE_SAMPLED_BIT;

		Vk_ImageCreate(&font->image);

		VkBufferImageCopy copy = { 0 };

		copy.bufferOffset      = 0;
		copy.bufferRowLength   = 0;
		copy.bufferImageHeight = 0;

		copy.imageSubresource.aspectMask     = VK_IMAGE_ASPECT_COLOR_BIT;
		copy.imageSubresource.mipLevel       = 0;
		copy.imageSubresource.baseArrayLayer = 0;
		copy.imageSubresource.layerCount     = 1;

		copy.imageExtent.width  = w;
		copy.imageExtent.height = h;
		copy.imageExtent.depth  = 1;

		VkImageMemoryBarrier2 transfer = { 0 };
		VkImageMemoryBarrier2 shader_read = { 0 };

		transfer.sType         = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2;
		transfer.srcStageMask  = VK_PIPELINE_STAGE_2_NONE;
		transfer.srcAccessMask = VK_ACCESS_2_NONE;
		transfer.dstStageMask  = VK_PIPELINE_STAGE_2_TRANSFER_BIT;
		transfer.dstAccessMask = VK_ACCESS_2_TRANSFER_WRITE_BIT;
		transfer.oldLayout     = VK_IMAGE_LAYOUT_UNDEFINED;
		transfer.newLayout     = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
		transfer.image         = font->image.handle;

		transfer.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
		transfer.subresourceRange.layerCount = 1;
		transfer.subresourceRange.levelCount = 1;

		shader_read.sType         = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2;
		shader_read.srcStageMask  = VK_PIPELINE_STAGE_2_TRANSFER_BIT;
		shader_read.srcAccessMask = VK_ACCESS_2_TRANSFER_WRITE_BIT;
		shader_read.dstStageMask  = VK_PIPELINE_STAGE_2_FRAGMENT_SHADER_BIT;
		shader_read.dstAccessMask = VK_ACCESS_2_SHADER_SAMPLED_READ_BIT;
		shader_read.oldLayout     = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
		shader_read.newLayout     = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
		shader_read.image         = font->image.handle;

		shader_read.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
		shader_read.subresourceRange.layerCount = 1;
		shader_read.subresourceRange.levelCount = 1;

		VkDependencyInfo dep = { 0 };
		dep.sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO;

		dep.imageMemoryBarrierCount = 1;
		dep.pImageMemoryBarriers    = &transfer;

		vk.CmdPipelineBarrier2(cmds->handle, &dep);

		vk.CmdCopyBufferToImage(cmds->handle, staging->handle, font->image.handle, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &copy);

		dep.pImageMemoryBarriers = &shader_read;

		vk.CmdPipelineBarrier2(cmds->handle, &dep);

		Vk_CommandBufferSubmit(cmds, true);

		SLL_PushN(draw->fonts, font, next);
		draw->n_fonts += 1;
	}
}
