#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 <sstream>
#include <tuple>
#include <stdint.h>

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;

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;
	if (size != fread(buf, 1, size, f)) {
		delete[] buf;
		return false;
	}
	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* type_strings[] = {
#define x(n) #n,
	svariable_type_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;
}

struct Desc {
	enum {
		VERTEX,
		FRAGMENT,
		stage_count
	};
	struct Variable {
		int type;
		std::string name;
		std::string tname;
	};
	int type;
	std::vector<Variable> attrs;
	std::vector<Variable> trgts;
	std::string entrypoints[stage_count];
	void read_var(Variable& d, cfg_Object* desc) {
		const char* stype = find_string_default(desc, "type", 0);
		if (!stype) {
			print_err("%s needs to have a type\n", desc->name);
			pbreak(303);
		}
		const char* sname = find_string_default(desc, "name", 0);
		if (!sname) {
			print_err("%s needs to have a name\n", desc->name);
			pbreak(304);
		}
		d.type = type_from_string(stype);
		d.tname = std::string(stype);
		d.name = std::string(sname);
		if (d.name.size() > 28) {
			print_err("variable name %s is too long (max 28 chars).\n", sname);
			pbreak(305);
		}
	}
	void build(cfg_Object* desc) {
		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[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[FRAGMENT] = sf;
		desc = desc->next;
		while (desc) {
			Variable v;
			if (!strcmp(desc->name, "attribute")) {
				read_var(v, desc);
				attrs.push_back(v);
			} else if (!strcmp(desc->name, "target")) {
				read_var(v, desc);
				trgts.push_back(v);
			}
			desc = desc->next;
		}
	}

	std::string build_vs() const {
		std::stringstream ss;
		size_t i, l = attrs.size();
		for (i = 0; i < l; i++) {
			auto& attr = attrs[i];
			ss << "layout (location = 0) in ";
			ss << attr.tname << " ";
			ss << attr.name << ";\n";
		}
		return ss.str();
	}

	std::string build_fs() const {
		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";
		}
		return ss.str();
	}
};

std::vector<uint32_t> compile_shader(
	const Desc& d,
	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;
	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());
	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(
	const Desc& d,
	int stage,
	const char*& dfn,
	EShLanguage& l,
	std::string& ps
) {
	switch (stage) {
		case Desc::VERTEX:
			dfn = "VERTEX_SHADER";
			l = EShLangVertex;
			ps += d.build_vs();
			break;
		case Desc::FRAGMENT:
			dfn = "FRAGMENT_SHADER";
			l = EShLangFragment;
			ps += d.build_fs();
			break;
	}
}

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

void write_csh(
	const char* fname,
	const Desc& d,
	const std::vector<uint32_t>* stages
) {
	int hsize = 8, i, coff;
	FILE* f = fopen(fname, "wb");
	if (!f) {
		print_err("Failed to open %s\n", fname);
		pbreak(500);
	}
	fwrite("CSH2", 4, 1, f);
	fwrite(&d.type, 4, 1, f);
	for (const auto& a : d.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;
	}
	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 (i = 0, coff = 0; i < Desc::stage_count; i++) {
		int o = 0;
		if (d.entrypoints[i].empty())
			fwrite(&o, 4, 1, f);
		else {
			o = hsize + coff;
			fwrite(&o, 4, 1, f);
			coff += stages[i].size() * sizeof(uint32_t);
		}
	}
	for (i = 0; i < Desc::stage_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[Desc::stage_count];
	cfg_Object* cdesc;
	void* dp_mem;
	Desc desc;
	if (argc < 3) {
		print_err("Usage: %s infile outfile.\n", argv[0]);
	}
	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(spv, src, desc);
	write_csh(argv[2], desc, spv);
	return 0;
}