#include "glslang/Public/ShaderLang.h"
#include "glslang/Public/ResourceLimits.h"
#include "SPIRV/GlslangToSpv.h"
#include "includer.hpp"

extern "C" {
#include "cfgparse.h"
#include "plat.h"
#include "sh_enums.h"
#include "str.h"
}

#include <sstream>
#include <stdint.h>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>

const char* glsl_version_s = "#version 440 core";
const int glsl_version = 440;
const char* builtin_src = "";
static const auto client_version = glslang::EShTargetVulkan_1_0;
static const auto target_version = glslang::EShTargetSpv_1_0;
const int desc_parse_memory = 4096 * 1024;

std::string get_desc(const char* src) {
	const char* srcs[] = {
		glsl_version_s, "\n",
		"#define DESC", "\n",
		builtin_src,
		src
	};
	glslang::TShader shader(EShLangVertex);
	Includer inc;
	std::string prepr;
	shader.setStrings(srcs, sizeof srcs / sizeof *srcs);
	shader.setEnvClient(glslang::EShClientVulkan, client_version);
	shader.setEnvTarget(glslang::EShTargetSpv, target_version);
	if (!shader.preprocess(
		GetDefaultResources(),
		glsl_version,
		ENoProfile,
		false,
		false,
		(EShMessages)(EShMsgSpvRules | EShMsgVulkanRules),
		&prepr,
		inc
	)) {
		print_err("Preprocessing failed.\n");
		print_err("%s\n", shader.getInfoLog());
		exit(1);
	}
	return prepr;
}

bool rf(const char* n, char*& buf, size_t& size) {
	FILE* f = fopen(n, "r");
	if (!f) return false;
	fseek(f, 0, SEEK_END);
	size = ftell(f);
	rewind(f);
	buf = new char[size + 1];
	buf[size] = 0;
	buf[fread(buf, 1, size, f)] = 0;
	return true;
}

cfg_Object* parse_desc(void*& mem, const std::string& s) {
	Arena a;
	mem = malloc(desc_parse_memory);
	if (!mem) {
		print_err("Out of memory.\n");
		pbreak(30);
	}
	init_arena(&a, mem, desc_parse_memory);
	return cfg_parse(s.c_str(), &a);
}

static const char* program_type_strings[] = {
#define x(n) #n,
	sprogram_type_xmacro()
#undef x
};

static const char* stage_strings[] = {
#define x(n) #n,
	shader_type_xmacro()
#undef x
};

static const char* type_strings[] = {
#define x(n) #n,
	svariable_type_xmacro()
#undef x
};

static const char* rate_strings[] = {
#define x(n) #n,
	sbinding_rate_xmacro()
#undef x
};

int get_program_type(cfg_Object* desc) {
	int i;
	const char* s = find_string_default(desc, "type", 0);
	if (!s) {
		print_err("No program type\n", s);
		pbreak(35);
	}
	for (i = 0; i < sprogram_type_count; i++) {
		if (!strcmp(s, program_type_strings[i]))
			return i;
	}
	print_err("Invalid program type %s\n", s);
	pbreak(36);
	return sprogram_type_graphics;
}

int type_from_string(const char* s) {
	int i;
	for (i = 0; i < svariable_type_count; i++)
		if (!strcmp(type_strings[i], s))
			return i;
	print_err("Invalid variable type %s\n", s);
	pbreak(302);
	return 0;
}

int stage_from_string(const char* s) {
	int i;
	for (i = 0; i < shader_type_count; i++)
		if (!strcmp(stage_strings[i], s))
			return i;
	print_err("Invalid stage type %s\n", s);
	pbreak(800);
	return 0;
}

SBinding_Rate rate_from_string(const char* s) {
	int i;
	for (i = 0; i < sbinding_rate_count; i++) {
		if (!strcmp(rate_strings[i], s))
			return (SBinding_Rate)i;
	}
	print_err("Invalid binding rate %s\n", s);
	pbreak(315);
	return (SBinding_Rate)0;
}

struct Desc {
	struct Variable {
		int type;
		std::string name;
		std::string tname;

		std::pair<int, int> size() const {
			switch (type) {
				case svariable_type_float:
					return { 4, 4 };
				case svariable_type_vec2:
					return { 8, 8 };
				case svariable_type_vec3:
					return { 12, 16 };
				case svariable_type_vec4:
					return { 16, 16 };
				case svariable_type_mat2:
					return { 16, 16 };
				case svariable_type_mat3:
					return { 36, 16 };
				case svariable_type_mat4:
					return { 64, 16 };
			}
			assert(0);
			return { 0, 0 };
		}

		int size(int coff) const {
			auto [desired, align] = size();
			int pad = align_size(coff, align) - coff;
			return desired + pad;
		}
	};
	struct Binding {
		SBinding_Rate rate;
		std::string name;
		std::vector<Variable> attrs;
	};
	struct Texture {
		int stage;
		int dimension;
	};
	struct Descriptor {
		std::string name;
		int slot;
		int stage;
	};
	struct Struct {
		std::string name;
		std::vector<Variable> vars;
	};
	struct CBuffer {
		int size;
		int stage;
		int var_count;
		std::string strct;
		std::unordered_map<std::string, int> offsets;
	};
	int type;
	std::vector<Binding> bindings;
	std::vector<Variable> trgts;
	std::vector<Variable> interp;
	std::unordered_map<std::string, Struct> structs;
	std::unordered_map<std::string, Texture> textures;
	std::unordered_map<std::string, CBuffer> cbuffers;
	std::vector<Descriptor> descriptors;
	std::string entrypoints[shader_type_count];
	void read_var(Variable& d, cfg_Object* desc) {
		const char* sname = find_string_default(desc, "name", 0);
		if (!sname) {
			print_err("%s needs to have a name\n", desc->name);
			pbreak(303);
		}
		const char* stype = find_string_default(desc, "type", 0);
		if (!stype) {
			print_err("%s needs to have a type\n", sname);
			pbreak(304);
		}
		d.type = type_from_string(stype);
		d.tname = std::string(stype);
		d.name = std::string(sname);
		if (d.name.size() > 27) {
			print_err("variable name %s is too long (max 27 chars).\n", sname);
			pbreak(305);
		}
	}
	Binding* read_binding(cfg_Object* desc) {
		const char* sname = find_string_default(desc, "name", 0);
		if (!sname) {
			print_err("%s must have a name.\n", desc->name);
			pbreak(310);
		}
		const char* srate = find_string_default(desc, "rate", 0);
		if (!srate) {
			print_err("%s must specify a rate.\n", sname);
			pbreak(311);
		}
		Binding& b = bindings.emplace_back();
		b.rate = rate_from_string(srate);
		b.name = std::string(sname);
		if (b.name.size() > 23)  {
			print_err("Binding name %s is too long (max 23 chars).\n", sname);
			pbreak(312);
		}
		return &b;
	}
	void read_texture(cfg_Object* desc) {
		const char* sdem;
		const char* sname = find_string_default(desc, "name", 0);
		if (!sname) {
			print_err("%s must have a name.\n", desc->name);
			pbreak(801);
		}
		const char* sstage = find_string_default(desc, "stage", 0);
		if (!sstage) {
			print_err("%s must define a stage.\n", sname);
			pbreak(802);
		}
		std::string n = std::string(sname);
		if (n.size() > 23) {
			print_err("Texture name %s is too long (max 23 chars).\n", sname);

			pbreak(803);
		}
		if (!textures.contains(n)) {
			Texture& t = textures[n];
			t.stage = 0;
			t.dimension = 0;
		}
		sdem = find_string_default(desc, "dimension", 0);
		Texture& t = textures[n];
		t.stage |= 1 << stage_from_string(sstage);
		t.dimension =
			sdem && string_equal(sdem, "cube")? -6:
			find_int_default(desc, "dimension", 2);
	}
	cfg_Object* read_struct(cfg_Object* desc) {
		const char* sname = find_string_default(desc, "name", 0);
		if (!sname) {
			print_err("%s must have a name.\n", desc->name);
			pbreak(901);
		}
		std::string n = std::string(sname);
		if (structs.contains(n)) {
			print_err("Already a struct with the name %s.\n", sname);
			pbreak(903);
		}
		desc = desc->next;
		Struct& s = structs[n];
		while (desc && !strcmp(desc->name, "variable")) {
			Variable v;
			read_var(v, desc);
			s.vars.push_back(v);
			desc = desc->next;
		}
		return desc;
	}
	
	void read_cbuffer(cfg_Object* desc) {
		const char* sname = find_string_default(desc, "name", 0);
		if (!sname) {
			print_err("%s must have a name.\n", desc->name);
			pbreak(951);
		}
		const char* stype = find_string_default(desc, "type", 0);
		if (!stype) {
			print_err("%s must define a type.\n", desc->name);
		}
		std::string n = std::string(sname);
		std::string t = std::string(stype);
		if (n.size() > 23) {
			print_err("CBuffer name too long (max 23 chars).\n");
			pbreak(952);
		}
		if (!structs.contains(t)) {
			print_err("No such struct %s\n", stype);
			pbreak(953);
		}
		const char* sstage = find_string_default(desc, "stage", 0);
		if (!sstage) {
			print_err("%s must define a stage.\n", sname);
			pbreak(954);
		}
		Struct& strct = structs[t];
		CBuffer& buf = cbuffers[n];
		int offset = 0;
		buf.strct = t;
		buf.stage |= 1 << stage_from_string(sstage);
		buf.var_count = strct.vars.size();
		for (const auto& v : strct.vars) {
			buf.offsets[v.name] = offset;
			offset += v.size(offset);
		}
		buf.size = offset;
	}

	void build(cfg_Object* desc) {
		int i;
		Binding* cur_binding = 0;
		type = get_program_type(desc);
		if (type != sprogram_type_graphics) {
			assert(0); /* todo */
			return;
		}
		const char* sv = find_string_default(desc, "vertex", 0);
		if (!sv) {
			print_err("Graphics programs must define a vertex shader entry point.\n");
			pbreak(301);
		}
		entrypoints[shader_type_vertex] = sv;
		const char* sf = find_string_default(desc, "fragment", 0);
		if (!sf) {
			print_err("Graphics programs must define a fragment shader entry point.\n");
			pbreak(302);
		}
		entrypoints[shader_type_fragment] = sf;
		for (i = 0; i < (int)shader_type_count; i++) {
			if (entrypoints[i].size() > 23) {
				print_err("Entry point name is too long.\n");
				pbreak(3000);
			}
		}
		desc = desc->next;
		while (desc) {
			Variable v;
			if (!strcmp(desc->name, "binding")) {
				cur_binding = read_binding(desc);
			} else if (!strcmp(desc->name, "attribute")) {
				if (!cur_binding) {
					print_err("Can't specify a vertex attribute with no binding.\n");
					pbreak(303);
				}
				read_var(v, desc);
				cur_binding->attrs.push_back(v);
			} else if (!strcmp(desc->name, "target")) {
				read_var(v, desc);
				trgts.push_back(v);
			} else if (!strcmp(desc->name, "interpolator")) {
				read_var(v, desc);
				interp.push_back(v);
			} else if (!strcmp(desc->name, "texture")) {
				read_texture(desc);
			} else if (!strcmp(desc->name, "cbuffer")) {
				read_cbuffer(desc);
			} else if (!strcmp(desc->name, "struct")) {
				desc = read_struct(desc);
				continue;
			}
			desc = desc->next;
		}
		build_descriptors();
	}

	Descriptor* find_desc(const char* name) {
		int count = descriptors.size();
		int i, bucket = (int)(hash_string(name) % count);
		for (i = 0; i < count; i++) {
			Descriptor& d = descriptors[bucket];
			if (
				d.name.empty() ||
				!strcmp(name, d.name.c_str())
			) return &d;
			bucket = (bucket + 1) % count;
		}
		return 0;
	};

	void build_descriptors() {
		int slot = 0;
		descriptors.resize(textures.size() + cbuffers.size());
		for (auto& i : textures) {
			Descriptor* d = find_desc(i.first.c_str());
			d->name = i.first;
			d->slot = slot;
			d->stage = i.second.stage;
			slot++;
		}
		for (auto& i : cbuffers) {
			Descriptor* d = find_desc(i.first.c_str());
			d->name = i.first;
			d->slot = slot;
			d->stage = i.second.stage;
			slot++;
		}
	}

	void add_textures(std::stringstream& ss, Shader_Type stage) {
		for (const auto& it : textures) {
			const auto& texture = it.second;
			if (texture.stage & (1 << stage)) {
				Descriptor* d = find_desc(it.first.c_str());
				assert(d != 0);
				ss << "layout (binding = " << d->slot << ") ";
				ss << "uniform sampler";
				if (texture.dimension == -6)
					ss << "Cube ";
				else
					ss << texture.dimension << "D ";
				ss << it.first << ";\n";
			}
		}
	}

	void add_cbuffers(std::stringstream& ss, Shader_Type stage) {
		for (const auto& it : cbuffers) {
			const auto& cbuffer = it.second;
			if (cbuffer.stage & (1 << stage)) {
				const Struct& s = structs[cbuffer.strct];
				Descriptor* d = find_desc(it.first.c_str());
				assert(d != 0);
				ss << "layout (std140, binding = " << d->slot << ") ";
				ss << "uniform " << cbuffer.strct << "{\n";
				for (const auto& v : s.vars) {
					ss << v.tname << " " << v.name << ";\n";
				}
				ss << "} " << it.first << ";\n";
			}
		}
	}

	std::string build_vs() {
		std::stringstream ss;
		size_t i, li = bindings.size();
		size_t l = 0;
		for (i = 0; i < li; i++) {
			auto& binding = bindings[i];
			size_t j, lj = binding.attrs.size();
			for (j = 0; j < lj; j++, l++) {
				auto& attr = binding.attrs[j];
				ss << "layout (location = " << l << ") in ";
				ss << attr.tname << " ";
				ss << attr.name << ";\n";
			}
		}
		if (interp.size()) {
			ss << "layout (location = 0) out _Interpolator {\n";
			for (const auto& i : interp)
				ss << i.tname << " " << i.name << ";\n";
			ss << "} interpolator;\n";
		}
		add_textures(ss, shader_type_vertex);
		add_cbuffers(ss, shader_type_vertex);
		return ss.str();
	}

	std::string build_fs() {
		std::stringstream ss;
		size_t i, l = trgts.size();
		for (i = 0; i < l; i++) {
			auto& attr = trgts[i];
			ss << "layout (location = 0) out ";
			ss << attr.tname << " ";
			ss << attr.name << ";\n";
		}
		if (interp.size()) {
			ss << "layout (location = 0) in _Interpolator {\n";
			for (const auto& i : interp)
				ss << i.tname << " " << i.name << ";\n";
			ss << "} interpolator;\n";
		}
		add_textures(ss, shader_type_fragment);
		add_cbuffers(ss, shader_type_fragment);
		return ss.str();
	}
};

std::vector<uint32_t> compile_shader(
	Desc& d,
	const char* fname,
	const char* presrc,
	const char* src,
	const char* define,
	int stage,
	EShLanguage lang
) {
	std::string vars = d.build_vs();
	const char* srcs[] = {
		glsl_version_s, "\n",
		"#define ", define, "\n",
		builtin_src,
		presrc,
		src
	};
	glslang::TShader shader(lang);
	glslang::TProgram program;
	glslang::TIntermediate* ir;
	glslang::SpvOptions options;
	Includer inc;
	std::string prepr;
	std::vector<uint32_t> spv;
#ifdef DEBUG
	options.disableOptimizer = true;
	options.generateDebugInfo = true;
	options.emitNonSemanticShaderDebugInfo = true;
	options.emitNonSemanticShaderDebugSource = true;
	options.stripDebugInfo = false;
#else
	options.disableOptimizer = false;
	options.stripDebugInfo = true;
#endif
	EShMessages msg = (EShMessages)(EShMsgSpvRules | EShMsgVulkanRules);
	shader.setStrings(srcs, sizeof srcs / sizeof *srcs);
	shader.setEnvClient(glslang::EShClientVulkan, client_version);
	shader.setEnvTarget(glslang::EShTargetSpv, target_version);
	shader.setEntryPoint(d.entrypoints[stage].c_str());
#ifdef DEBUG
	shader.setDebugInfo(true);
	shader.setSourceEntryPoint(d.entrypoints[stage].c_str());
	shader.setSourceFile(fname);
	shader.addSourceText(src, strlen(src));
#else
	(void)fname;
#endif
	if (!shader.preprocess(
		GetDefaultResources(),
		glsl_version,
		ENoProfile,
		false,
		false,
		msg,
		&prepr,
		inc
	)) {
		print_err("%s\n", shader.getInfoLog());
		exit(400);
	}
	if (!shader.parse(
		GetDefaultResources(),
		glsl_version,
		ENoProfile,
		false,
		false,
		msg,
		inc
	)) {
		print_err("%s\n", shader.getInfoLog());
		exit(401);
	}
	program.addShader(&shader);
	if (!program.link(msg)) {
		print_err("%s\n", program.getInfoLog());
		exit(402);
	}
	ir = program.getIntermediate(lang);
	options.validate = true;
	glslang::GlslangToSpv(*ir, spv, &options);
	return spv;
}

void configure(
	Desc& d,
	int stage,
	const char*& dfn,
	EShLanguage& l,
	std::string& ps
) {
	switch (stage) {
		case shader_type_vertex:
			dfn = "VERTEX_SHADER";
			l = EShLangVertex;
			ps += d.build_vs();
			break;
		case shader_type_fragment:
			dfn = "FRAGMENT_SHADER";
			l = EShLangFragment;
			ps += d.build_fs();
			break;
	}
}

void compile_shaders(
	const char* fname,
	std::vector<uint32_t>* spv,
	const char* src,
	Desc& d
) {
	const char* define;
	EShLanguage lang;
	int i;
	for (i = 0; i < shader_type_count; i++) {
		if (!d.entrypoints[i].empty()) {
			std::string ps;
			configure(d, i, define, lang, ps);
			spv[i] = compile_shader(
				d,
				fname,
				ps.c_str(),
				src,
				define,
				i,
				lang
			);
		}
	}
}

void write_csh(
	const char* fname,
	const Desc& d,
	const std::vector<uint32_t>* stages
) {
	int hsize = 20, i, coff;
	FILE* f = fopen(fname, "wb");
	if (!f) {
		print_err("Failed to open %s\n", fname);
		pbreak(500);
	}
	int c;
	fwrite("CSH2", 4, 1, f);
	fwrite(&d.type, 4, 1, f);
	c = d.bindings.size(); fwrite(&c, 4, 1, f);
	c = d.trgts.size(); fwrite(&c, 4, 1, f);
	c = d.descriptors.size(); fwrite(&c, 4, 1, f);
	for (const auto& b : d.bindings) {
		char buf[24];
		int count = b.attrs.size();
		memset(buf, 0, sizeof buf);
		strcpy(buf, b.name.c_str());
		fwrite(buf, 1, sizeof buf, f);
		fwrite(&b.rate, 4, 1, f);
		fwrite(&count, 4, 1, f);
		for (const auto& a : b.attrs) {
			char buf[28];
			memset(buf, 0, sizeof buf);
			strcpy(buf, a.name.c_str());
			fwrite(buf, 1, sizeof buf, f);
			fwrite(&a.type, 4, 1, f);
			hsize += 32;
		}
		hsize += 32;
	}
	for (const auto& t : d.trgts) {
		char buf[28];
		memset(buf, 0, sizeof buf);
		strcpy(buf, t.name.c_str());
		fwrite(buf, 1, sizeof buf, f);
		fwrite(&t.type, 4, 1, f);
		hsize += 32;
	}
	for (const auto& d : d.descriptors) {
		char buf[24];
		memset(buf, 0, sizeof buf);
		strcpy(buf, d.name.c_str());
		fwrite(buf, 1, sizeof buf, f);
		fwrite(&d.slot,  4, 1, f);
		fwrite(&d.stage, 4, 1, f);
		hsize += 32;
	}
	hsize += shader_type_count * 32;
	for (i = 0, coff = 0; i < shader_type_count; i++) {
		int o = 0;
		char buf[24];
		memset(buf, 0, sizeof buf);
		if (d.entrypoints[i].empty()) {
			fwrite(&o, 4, 1, f);
			fwrite(&o, 4, 1, f);
		} else {
			int size = stages[i].size() * sizeof(uint32_t);
			strcpy(buf, d.entrypoints[i].c_str());
			o = hsize + coff;
			fwrite(&o, 4, 1, f);
			fwrite(&size, 4, 1, f);
			coff += size;
		}
		fwrite(buf, 1, sizeof buf, f);
	}
	for (i = 0; i < shader_type_count; i++)
		if (!d.entrypoints[i].empty()) {
			auto& stage = stages[i];
			fwrite(&stage[0], sizeof(uint32_t), stage.size(), f);
		}
}

int main(int argc, const char** argv) {
	char* src;
	size_t src_size;
	std::string desc_src;
	std::vector<uint32_t> spv[shader_type_count];
	cfg_Object* cdesc;
	void* dp_mem;
	Desc desc;
	if (argc < 3) {
		print_err("Usage: %s infile outfile.\n", argv[0]);
		return 1;
	}
	if (!rf(argv[1], src, src_size)) {
		print_err("Failed to read %s\n", argv[1]);
		return 1;
	}
	desc_src = get_desc(src);
	cdesc = parse_desc(dp_mem, desc_src.c_str());
	desc.build(cdesc);
	compile_shaders(argv[2], spv, src, desc);
	write_csh(argv[2], desc, spv);
	return 0;
}