/* Mathias Pedersen 2019 */

#include <Serializer.hpp>
#include <VectorStream.hpp>
#include <Compiler.hpp>
#include <InstructionConverter.hpp>
#include <unordered_map>
#include <iostream>

extern "C"
{
#include <Lua/lauxlib.h>
#include <Lua/lobject.h>
#include <Lua/lstring.h>
}

void writeProto(VectorStream& vs, std::unordered_map<TString*, int>& stringTable, std::unordered_map<Proto*, int>& protoTable, Proto* p);

class StringTable
{
	VectorStream& vs;

	size_t stringCount = 1;
	std::unordered_map<TString*, int> stringMap;
	std::vector<TString*> strings;

	void getProtoStrings(Proto* p)
	{
		for (int i = 0; i < p->sizek; i++)
		{
			if (p->k[i].tt == LUA_TSTRING)
				insert(&p->k[i].value.gc->ts);
		}

		for (int i = 0; i < p->sizep; i++)
			getProtoStrings(p->p[i]);
	}

public:
	void insert(TString* ts)
	{
		if (stringMap.find(ts) != stringMap.end())
			return;
		stringMap[ts] = stringCount++;
		strings.push_back(ts);
	}

	std::unordered_map<TString*, int> write(Proto* p)
	{
		getProtoStrings(p);

		vs.writeCompressedInt(strings.size());
		for (TString* ts : strings)
			vs.writeString(std::string(getstr(ts), ts->tsv.len));

		return stringMap;
	}

	StringTable(VectorStream& vs) : vs(vs) {}
};

class ProtoTable
{
	VectorStream& vs;

	size_t protoCount = 0;
	std::unordered_map<Proto*, int> protoMap;
	std::vector<Proto*> protos;

	void addProto(Proto* p)
	{
		for (int i = 0; i < p->sizep; i++)
			addProto(p->p[i]);

		protoMap[p] = protoCount++;
		protos.push_back(p);
	}

public:
	std::unordered_map<Proto*, int> write(Proto* p, std::unordered_map<TString*, int>& stringTable)
	{
		addProto(p);

		vs.writeCompressedInt(protos.size());
		for (Proto* np : protos)
			writeProto(vs, stringTable, protoMap, np);

		return protoMap;
	}

	ProtoTable(VectorStream& vs) : vs(vs) {}
};

enum ConstantType
{
	ConstantNil,
	ConstantBoolean,
	ConstantNumber,
	ConstantString
};

void writeProto(VectorStream& vs, std::unordered_map<TString*, int>& stringTable, std::unordered_map<Proto*, int>& protoTable, Proto* p)
{
	vs.writeByte(p->maxstacksize);
	vs.writeByte(p->numparams);
	vs.writeByte(p->nups);
	vs.writeByte(p->is_vararg);

	InstructionConverter ic(p);
	std::vector<Instruction> instructions = ic.convertInstructions();

	/* Instructions */
	vs.writeCompressedInt(instructions.size());
	for (size_t i = 0; i < instructions.size(); i++)
		vs.writeInt(instructions[i]);

	/* Constants */
	vs.writeCompressedInt(p->sizek);
	for (int i = 0; i < p->sizek; i++)
	{
		TValue* k = &p->k[i];
		switch (k->tt)
		{
		case LUA_TNIL:
		{
			vs.writeByte(ConstantNil);
			break;
		}
		case LUA_TBOOLEAN:
		{
			vs.writeByte(ConstantBoolean);
			vs.writeByte(k->value.b);
			break;
		}
		case LUA_TNUMBER:
		{
			vs.writeByte(ConstantNumber);
			vs.writeDouble(k->value.n);
			break;
		}
		case LUA_TSTRING:
		{
			vs.writeByte(ConstantString);
			vs.writeCompressedInt(stringTable[rawtsvalue(k)]);
			break;
		}
		}
	}

	/* Prototypes */
	vs.writeCompressedInt(p->sizep);
	for (int i = 0; i < p->sizep; i++)
		vs.writeCompressedInt(protoTable[p->p[i]]);

	/* Some string index */
	vs.writeCompressedInt(1);

	/* Lineinfo */
	vs.writeCompressedInt(instructions.size());
	for (size_t i = 0; i < instructions.size(); i++)
		vs.writeCompressedInt(0);

	/* Write 0 for no debug info */
	vs.writeByte(0);
}

std::string serialize(lua_State* L, const std::string& source)
{
	VectorStream vs;

	const LClosure* cl = Compiler::compile(L, source);
	if (cl == nullptr)
	{
		vs.writeByte(0);
		vs.writeString(lua_tostring(L, -1));
		return vs.str();
	}

	vs.writeByte(1);

	/* Write string table */
	StringTable st(vs);
	st.insert(luaS_new(L, ""));
	std::unordered_map<TString*, int> stringTable = st.write(cl->p);

	/* Write proto table */
	ProtoTable pt(vs);
	std::unordered_map<Proto*, int> protoTable = pt.write(cl->p, stringTable);

	/* Entry point in proto table */
	vs.writeCompressedInt(protoTable[cl->p]);

	return vs.str();
}