#include <cassert>
#include <cstddef>
#include <cstdlib>
#include <cstring>
#include <string>
#include <fstream>
#include <iostream>

#include <glm/gtc/matrix_transform.hpp>

#include "asset.h"
#define GL_DEBUG_IMPLEMENTATION
#include "GLDebug.h"
#include "shader.h"
#include "util.h"


// NOTE: forward declarations

const std::string dumpTextFile(const char* filepath);
bool parseShader(MemoryArena* arena, GLContext* gl_ctx, shader_program* s);
void initCTXSizes(GLContext* gl_ctx);
bool compileAndLinkShader(shader_program* shader,
	const char* vert_src,
	const char* frag_src,
	GLuint& vs_id,
	GLuint& fs_id);
u32 getGLTypeSize(GLenum e);
GLTexture* getFreeGLTexture(GLContext* gl_ctx);
bool loadGLTexture(util_image* image, GLuint& tex_id);


// NOTE: interface

bool
addShaderProgram(MemoryArena* arena,
	GLContext* gl_ctx,
	const char* vs,
	const char* fs,
	const char* name)
{
	// TODO: replace std::string w/ utilAllocCStr() ?
	std::string hash_str = std::string(vs) + std::string(fs);
	const u32 max_len = 256;
	u64 hash = utilFNV64a_str(hash_str.substr(0, max_len).c_str());

	if (getShaderByHash(gl_ctx, hash)) {
		printf("%s(), shader is already loaded\n", __FUNCTION__);
		return false;
	}

	shader_program* s = getFreeShader(gl_ctx);

	if (s) {
		s->name = arenaCopyCStr(arena, name);
		s->hash = hash;

		// TODO: replace std::string w/ utilAllocateCStr() ?
		std::string vert = dumpTextFile(vs);
		std::string frag = dumpTextFile(fs);
		GLuint vs_id, fs_id;

		if (compileAndLinkShader(s, vert.c_str(), frag.c_str(), vs_id, fs_id)) {
			// NOTE: we call this here because we can only make the GL_MAX*
			// 	queries after at least one shader has been linked
			initCTXSizes(gl_ctx);
			s->model_xform_id = glGetUniformLocation(s->prog_id, "model_xform");
			glDetachShader(s->prog_id, vs_id);
			glDetachShader(s->prog_id, fs_id);
			glDeleteShader(vs_id);
			glDeleteShader(fs_id);

			return parseShader(arena, gl_ctx, s);
		}

		printf("%s(), Error linking shader\n", __FUNCTION__);
		return false;
	}

	printf("%s(), error loading shader\n", __FUNCTION__);
	return false;
}

shader_program*
getFreeShader(GLContext* gl_ctx)
{
	if (gl_ctx->num_shaders >= gl_ctx->max_shaders) {
		printf("%s(), GLContext->shaders full\n", __FUNCTION__);
		return nullptr;
	}

	shader_program* s = &gl_ctx->shaders[gl_ctx->num_shaders++];
	return s;
}

shader_program*
getShaderByHash(GLContext* gl_ctx, u64 hash)
{
	for (u32 i; i < gl_ctx->num_shaders; i++) {
		if (gl_ctx->shaders[i].hash == hash)
			return &gl_ctx->shaders[i];
	}

	return nullptr;
}

shader_program*
getShaderByName(const char* name, GLContext* gl_ctx)
{
	// FIXME: no reason to use hashes here. should use std::strncmp instead
	u64 hash = utilFNV64a_str(name);

	for (u32 i = 0; i < gl_ctx->num_shaders; i++) {
		if (utilFNV64a_str(gl_ctx->shaders[i].name) == hash)
			return &gl_ctx->shaders[i];
	}

	printf("%s(), shader not found, %s\n", __FUNCTION__, name);
	return nullptr;
}

shader_program*
getShaderByID(GLContext* gl_ctx, GLuint prog_id)
{
	for (u32 i = 0; i < gl_ctx->num_shaders; i++) {
		if (gl_ctx->shaders[i].prog_id)
			return &gl_ctx->shaders[i];
	}

	printf("%s(), shader not found, %d\n", __FUNCTION__, prog_id);
	return nullptr;
}

GLTexture*
getGLTexture(GLContext* gl_ctx, util_image* diffuse_img)
{
	u64 fp_hash = utilFNV64a_str(diffuse_img->file_path);

	for (u32 i = 0; i < gl_ctx->num_textures; i++) {
		GLTexture* glt = &gl_ctx->textures[i];

		if (glt->filepath_hash == fp_hash)
			return glt;
	}

	GLTexture* glt = getFreeGLTexture(gl_ctx);
	if (!glt) return nullptr;

	glt->pixel_format = (diffuse_img->num_channels == 3) ? GL_RGB : GL_RGBA;
	glt->width = diffuse_img->w;
	glt->height = diffuse_img->h;
	glt->filepath_hash = diffuse_img->filepath_hash;

	if (loadGLTexture(diffuse_img, glt->id))
		return glt;

	printf("%s(), Error, unable to load texture\n", __FUNCTION__);
	return nullptr;
}

void
renderVAO(GLmesh* glmesh, glm::mat4* model_xform, shader_program* shader)
{
	glUseProgram(shader->prog_id);
	glBindVertexArray(glmesh->vao_id);
	glUniformMatrix4fv(shader->model_xform_id, 1, GL_FALSE,
		(float*) model_xform);
	glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, glmesh->element_buf->id);
	glDrawElements(
		glmesh->draw_mode, glmesh->num_indices, GL_UNSIGNED_SHORT, 0);
	glBindVertexArray(0);
}

void
initTransforms(MemoryArena* arena,
	transforms* xforms,
	gl_buffer* xform_ubo,
	GLContext* gl_ctx,
	float fov,
	float near_clip_plane,
	float aspect_ratio,
	glm::vec3 cam_pos,
	glm::vec3 look_pos)
{
	xforms->view_xform = glm::lookAt(cam_pos, look_pos, glm::vec3(0, 1, 0));
	xforms->proj_xform = glm::infinitePerspective(
		glm::radians(fov), aspect_ratio, near_clip_plane);

	glGenBuffers(1, &xform_ubo->id);
	xform_ubo->target = GL_UNIFORM_BUFFER;
	xform_ubo->data_type = GL_FLOAT;
	xform_ubo->name = arenaCopyCStr(arena, "matrices");

	glBindBuffer(xform_ubo->target, xform_ubo->id);
	glBufferData(xform_ubo->target, sizeof(*xforms), xforms,
		GL_DYNAMIC_DRAW);
	// bindbufferbase
	xform_ubo->binding_idx = gl_ctx->binding_count++;
	glBindBufferBase(xform_ubo->target, xform_ubo->binding_idx, xform_ubo->id);
	glBindBuffer(xform_ubo->target, 0);
}

GLVertexAttrib*
getVertexAttribByName(shader_program* shader, const char* name)
{
	for (u32 i = 0; i < shader->num_vertex_attribs; i++) {
		if (strncmp(shader->vertex_attribs[i].name, name, 256) == 0)
			return &shader->vertex_attribs[i];
	}

	printf("%s, attribute: %s, not found on shader: %s\n",
			__FUNCTION__, name, shader->name);
	return nullptr;
}

GLmesh
initGLmesh(const mesh& m, u32 num_mappings, GLenum draw_mode, glm::vec3 pos)
{
	GLmesh glm = {0};
	glm.num_indices = m.num_indices;
	glm.draw_mode = draw_mode;
	glm.usage = GL_STATIC_DRAW; // TODO: logic for updating dynamic meshes
	glm.num_vertex_attrib_buffers = num_mappings;
	glm.vertex_attrib_buffers = UTIL_ALLOC(num_mappings, gl_buffer);
	glm.model_xform = UTIL_ALLOC(1, glm::mat4);
	glm.element_buf = UTIL_ALLOC(1, gl_buffer);
	*glm.model_xform = glm::translate(glm::mat4(1), pos);

	return glm;
}

void
initGLAttribBuffer(gl_buffer* buf, GLenum target, GLVertexAttrib* attrib)
{
	glGenBuffers(1, &buf->id);
	buf->target = target;
	buf->data_type = attrib->data_type;
	buf->location = attrib->location;
	buf->name = utilAllocateCStr(attrib->name);
}

// FIXME: might as well pass in pointer to GLmesh, since that's how we're
// 	going to use this, can void copying GLmesh twice that way
GLmesh
loadGLMesh(const mesh& m,
	GLenum draw_mode,
	GLuint diffuse_texture_id,
	const glm::vec3& pos,
	u32 num_mappings,
	GLBufferToAttribMapping mappings[])
{
	// NOTE: need to use freeGLMesh() when freeing because we're not storing
	// 	the gl_buffers on a memory arean (yet)
	GLmesh glm = initGLmesh(m, num_mappings, draw_mode, pos);
	glm.tex_id = diffuse_texture_id;
	glGenVertexArrays(1, &glm.vao_id);
	glBindVertexArray(glm.vao_id);

	for (u32 i = 0; i < num_mappings; i++) {
		gl_buffer& buf = glm.vertex_attrib_buffers[i];
		GLVertexAttrib* attrib = mappings[i].attrib;
		u32 type_size = getGLTypeSize(attrib->data_type);
		assert(type_size > 0);

		initGLAttribBuffer(&buf, GL_ARRAY_BUFFER, attrib);
		glBindBuffer(buf.target, buf.id);
		glBufferData(buf.target,
			m.num_vertices * type_size,
			mappings[i].mesh_buf,
			glm.usage);
		glVertexAttribPointer(attrib->location, attrib->num_components,
							attrib->component_type, GL_FALSE, 0, 0);
		glEnableVertexAttribArray(attrib->location);
	}

	glGenBuffers(1, &glm.element_buf->id);
	glm.element_buf->target = GL_ELEMENT_ARRAY_BUFFER;
	glm.element_buf->data_type = GL_UNSIGNED_SHORT;
	glBindBuffer(glm.element_buf->target, glm.element_buf->id);
	glBufferData(glm.element_buf->target,
		m.num_indices * sizeof(u16),
		m.indices,
		glm.usage);

	// TODO: many of these GL functions can set an error state
	// TODO: return error status
	glBindVertexArray(0);

	return glm;
}

void
freeGLMesh(GLmesh* glm)
{
	if (glm) {
		for (u32 i = 0; i < glm->num_vertex_attrib_buffers; i++)
			utilSafeFree(glm->vertex_attrib_buffers[i].name);

		utilSafeFree(glm->vertex_attrib_buffers);
		utilSafeFree(glm->element_buf);
		// NOTE: don't free the GLmesh because it is part of a memory arena
		//utilSafeFree(glm);
	}
}


// NOTE: internal

GLTexture*
getFreeGLTexture(GLContext* gl_ctx)
{
	if (gl_ctx->num_textures < gl_ctx->max_textures)
		return &gl_ctx->textures[gl_ctx->num_textures++];

	printf("%s(), no free textures\n", __FUNCTION__);
	return nullptr;
}

bool
loadGLTexture(util_image* image, GLuint& tex_id)
{
	glGenTextures(1, &tex_id);
	glBindTexture(GL_TEXTURE_2D, tex_id);
	glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
	GLenum pixel_format = (image->num_channels == 3) ? GL_RGB : GL_RGBA;
	glTexImage2D(GL_TEXTURE_2D, 0, pixel_format, image->w, image->h, 0,
				pixel_format, GL_UNSIGNED_BYTE, image->pixels);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);

	return (glGetError() == GL_NO_ERROR);
}

bool
compileAndLinkShader(shader_program* shader,
	const char* vert_src,
	const char* frag_src,
	GLuint& vs_id,
	GLuint& fs_id)
{
	if (strlen(vert_src) > 0 && strlen(frag_src) > 0) {
		vs_id = glCreateShader(GL_VERTEX_SHADER);
		fs_id = glCreateShader(GL_FRAGMENT_SHADER);
		glShaderSource(vs_id, 1, &vert_src, NULL);
		glShaderSource(fs_id, 1, &frag_src, NULL);

		// TODO: add a glError to string function in GLDebug.h
		glCompileShader(vs_id);
		assert(glGetError() == GL_NO_ERROR);
		glCompileShader(fs_id);
		assert(glGetError() == GL_NO_ERROR);

		shader->prog_id = glCreateProgram();
		glAttachShader(shader->prog_id, vs_id);
		glAttachShader(shader->prog_id, fs_id);

		glLinkProgram(shader->prog_id);
		GLint is_linked = 0;
		glGetProgramiv(shader->prog_id, GL_LINK_STATUS, &is_linked);

		return (is_linked == GL_TRUE);
	}

	// TODO: make another logging macro for printf style logging
	printf("%s(), empty shader source\n", __FUNCTION__);
	return false;
}

const std::string
dumpTextFile(const char* filepath)
{
	std::ifstream fs { filepath };
	std::string s {
		std::istreambuf_iterator<char>(fs),
		std::istreambuf_iterator<char>()};

	if (!fs || !fs.good())
		std::cout << "error reading file, " << filepath << "\n";

	return s;
}

u32
getGLTypeSize(GLenum e)
{
	switch (e) {
		case GL_FLOAT_VEC3: return 3 * sizeof(GLfloat);
		case GL_FLOAT_VEC4: return 4 * sizeof(GLfloat);
		case GL_FLOAT_MAT4: return 16 * sizeof(GLfloat);
		default:
			printf("%s(), unknown GLenum\n", __FUNCTION__);
			return 0;
	}
}

// NOTE: returns sizes based on GLSL layout std140
// 	https://www.khronos.org/opengl/wiki/Interface_Block_(GLSL)#Memory_layout
u32
getGLTypeSizeStd140(GLenum e)
{
	switch (e) {
		case GL_FLOAT_VEC3: return 4 * sizeof(GLfloat);
		case GL_FLOAT_VEC4: return 4 * sizeof(GLfloat);
		case GL_FLOAT_MAT4: return 16 * sizeof(GLfloat);
		default:
			printf("%s(), unknown GLenum\n", __FUNCTION__);
			return 0;
	}
}

const gl_uniform
parseUniform(MemoryArena* arena, shader_program* s, u32 uniform_idx)
{
	gl_uniform unif = {0};
	GLchar unif_name[256] = {0};
	GLsizei name_len = 0;

	glGetActiveUniform(s->prog_id,
		uniform_idx,
		sizeof(unif_name),
		&name_len,
		&unif.num_elements,
		&unif.data_type,
		unif_name);
	glGetActiveUniformsiv(s->prog_id, 1, &uniform_idx,
			GL_UNIFORM_BLOCK_INDEX, &unif.block_idx);

	unif.idx = uniform_idx;
	unif.name = arenaCopyCStr(arena, unif_name);
	unif.location = glGetUniformLocation(s->prog_id, unif.name);

	return unif;
}

bool
parseShaderUniforms(MemoryArena* arena, shader_program* s, GLContext* gl_ctx)
{
	// NOTE: only add uniforms in the default block to the base uniform array
	GLint num_uniforms_total = 0;
	glGetProgramiv(s->prog_id, GL_ACTIVE_UNIFORMS, &num_uniforms_total);
	GLint indices[num_uniforms_total];

	for (u32 i = 0; i < (u32) num_uniforms_total; i++) {
		GLint block_idx = 0;
		glGetActiveUniformsiv(s->prog_id, 1, &i,
				GL_UNIFORM_BLOCK_INDEX, &block_idx);

		if (block_idx == -1) {
			indices[s->num_uniforms] = i;
			s->num_uniforms++;
		}
	}

	s->uniforms = ARENA_ALLOC(arena, gl_uniform, s->num_uniforms);

	for (u32 i = 0; i < s->num_uniforms; i++) {
		const gl_uniform unif = parseUniform(arena, s, indices[i]);
		std::memcpy(&s->uniforms[i], &unif, sizeof(unif));
	}

	return true;
}

void
initCTXSizes(GLContext* gl_ctx)
{
	// NOTE: see https://docs.gl/gl3/glGet for other useful context info
	if (gl_ctx->max_binding_points == 0) {
		glGetIntegerv(GL_MAX_UNIFORM_BUFFER_BINDINGS,
				&gl_ctx->max_binding_points);
		glGetIntegerv(GL_MAX_VERTEX_UNIFORM_BLOCKS, &gl_ctx->max_vertex_blocks);
		glGetIntegerv(GL_MAX_FRAGMENT_UNIFORM_BLOCKS,
				&gl_ctx->max_fragment_blocks);
		glGetIntegerv(GL_MAX_UNIFORM_BLOCK_SIZE, &gl_ctx->max_ublock_size);
		glGetIntegerv(GL_MAX_VERTEX_ATTRIBS, &gl_ctx->max_vertex_attribs);

#if 1
		printf("%s(), context size info set\n", __FUNCTION__);
		printf("GL_MAX_UNIFORM_BUFFER_BINDINGS: %d\n",
				gl_ctx->max_binding_points);
		printf("GL_MAX_VERTEX_UNIFORM_BLOCKS: %d\n",
				gl_ctx->max_vertex_blocks);
		printf("GL_MAX_FRAGMENT_UNIFORM_BLOCKS: %d\n",
				gl_ctx->max_fragment_blocks);
		printf("GL_MAX_UNIFORM_BLOCK_SIZE: %d\n",
				gl_ctx->max_ublock_size);
		printf("GL_MAX_VERTEX_ATTRIBS: %d\n", gl_ctx->max_vertex_attribs);
#endif
	}
}

i32
ctxGetUniformBlockBinding(GLContext* gl_ctx, const char* name)
{
	for (u32 i = 0; i < gl_ctx->num_ubos; i++) {
		gl_buffer& ubo = gl_ctx->uniform_buffers[i];

		if (std::strstr(ubo.name, name))
			return ubo.binding_idx;
	}

	printf("%s(), no buffer found with name: %s\n", __FUNCTION__, name);
	return -1;
}

bool
parseUniformBlocks(MemoryArena* arena, shader_program* s, GLContext* gl_ctx)
{
	glGetProgramiv(s->prog_id, GL_ACTIVE_UNIFORM_BLOCKS,
			(GLint*) &s->num_blocks);
	s->uniform_blocks = ARENA_ALLOC(arena, GLUniformBlock, s->num_blocks);

	for (u32 i = 0; i < s->num_blocks; i++) {
		GLUniformBlock& ub = s->uniform_blocks[i];
		ub.block_id = i;

		GLchar block_name[256] = {0};
		GLsizei name_len = 0;
		glGetActiveUniformBlockName(
				s->prog_id, i, 256, &name_len, block_name);
		ub.name = arenaCopyCStr(arena, block_name);

		glGetActiveUniformBlockiv(s->prog_id, i,
				GL_UNIFORM_BLOCK_ACTIVE_UNIFORMS, (GLint*) &ub.num_uniforms);
		ub.uniforms = ARENA_ALLOC(arena, gl_uniform, ub.num_uniforms);
		GLint indices[ub.num_uniforms] = {0};
		glGetActiveUniformBlockiv(s->prog_id, i,
				GL_UNIFORM_BLOCK_ACTIVE_UNIFORM_INDICES, (GLint*) &indices);

		for (u32 j = 0; j < ub.num_uniforms; j++) {
			const gl_uniform unif = parseUniform(arena, s, indices[j]);
			std::memcpy(&ub.uniforms[j], &unif, sizeof(unif));
		}

		ub.binding_idx = ctxGetUniformBlockBinding(gl_ctx, ub.name);

#if 0
		if (ub.binding_idx < 0)
			return false;

		glUniformBlockBinding(s->prog_id, i, ub.binding_idx);
#endif
	}


	return true;
}

u32
getNumAttribComponents(GLenum type)
{
	switch (type) {
		case GL_FLOAT_VEC3: return 3;
		default:
			printf("%s(), unknown GLenum\n", __FUNCTION__);
			return 0;
	}
}

GLenum
getAttribComponentType(GLenum type)
{
	switch (type) {
		case GL_FLOAT_VEC3: return GL_FLOAT;
		default:
			printf("%s(), unknown GLenum\n", __FUNCTION__);
			return 0;
	}
}

bool
parseAttributes(MemoryArena* arena, shader_program* s, GLContext* gl_ctx)
{
	GLint num_attribs;
	glGetProgramiv(s->prog_id, GL_ACTIVE_ATTRIBUTES, &num_attribs);
	s->num_vertex_attribs = num_attribs;
	s->vertex_attribs = ARENA_ALLOC(arena, GLVertexAttrib, num_attribs);

	GLchar attrib_name[256] = {0};
	GLsizei length;
	GLint size;
	GLenum type;

	for (int i = 0; i < num_attribs; i++) {
		glGetActiveAttrib(s->prog_id, i, sizeof(attrib_name),
							&length, &size, &type, attrib_name);
		GLint location = glGetAttribLocation(s->prog_id, attrib_name);
		GLVertexAttrib* attrib = &s->vertex_attribs[i];
		attrib->data_type = type;
		attrib->location = location;
		attrib->num_components = getNumAttribComponents(type);
		assert(attrib->num_components > 0);
		attrib->component_type = getAttribComponentType(type);
		assert(attrib->component_type > 0);
		attrib->name = arenaCopyCStr(arena, attrib_name, sizeof(attrib_name));
	}

	return true;
}

bool
parseShader(MemoryArena* arena, GLContext* gl_ctx, shader_program* s)
{
	if (parseShaderUniforms(arena, s, gl_ctx)
		&& parseUniformBlocks(arena, s, gl_ctx)
		&& parseAttributes(arena, s, gl_ctx))
	{
		return true;
	}

	printf("%s(), Error parsing shader\n", __FUNCTION__);
	return false;
}