#include <utility>
#include <vector>
#include "LLVMVisitor.h"

LLVMRet LLVMVisitor::f_defun(const std::vector<Value*>& args) {
	Name* fname = static_cast<Name*>(args[0]);
	List* fargs = static_cast<List*>(args[1]);

	/* Transform into IR of structure
	 * define @name(T arg1, T arg2) {
	 * ...
	 *     ret T %reg
	 * }
	 */

	lreg.reset();

	LLVMRet funcRet;
	unsigned outbodySize = outBuf.size();
	for (unsigned i = 2; i < args.size(); i++) {
		funcRet = args[i]->Visit(*this);
	}
	
	std::map<std::string, LLVMRet> localVars;
	locals.push(&localVars);

	outFile << "define " << funcRet.type << " @" << fname->get_value() << "(";

	for (Value* v : fargs->get_value()) {
		std::string name = "%_" + static_cast<Name*>(v)->get_value();
		// TODO: figure out argument type, based on usage
		outFile << "i32 " << name;
		localVars.insert({ name, { "i32", name } });
		if (v != fargs->get_value().back())
			outFile << ", ";
	}
	outFile << ") {\n";

	for (unsigned i = outbodySize; i < outBuf.size(); i++) {
		outFile << "\t" << outBuf[i] << "\n";
	}
	while (outbodySize < outBuf.size()) {
		outBuf.pop_back();
	}

	outFile << "\tret " << funcRet << "\n}\n\n";
	locals.pop();
	LLVMRet retType = { funcRet.type, "@" + fname->get_value() };
	funcs_user.insert({ retType.addr, retType });
	return retType;
}

LLVMRet LLVMVisitor::f_defvar(funcargs args) {
	std::string name = "@" + static_cast<Name*>(args[0])->get_value();
	LLVMRet retVal = args[1]->Visit(*this);
	outFile << name << " = global " << retVal << "\n";

	LLVMRet val = { retVal.type, name };
	vars_user.insert({ name, val });
	return val;
}

LLVMRet LLVMVisitor::f_if(funcargs args) {
	LLVMRet boolExpr = args[0]->Visit(*this);
	Register rResult = lreg++,
	         rTrue = lreg++,
	         rFalse = lreg++,
	         rCont = (args.size() == 2) ? rFalse : lreg++;
	
	unsigned reservedI = outBuf.size();
	outBuf.push_back("");
	outBuf.push_back("");
	outBuf.push_back("");

	LLVMRet tExpr = args[1]->Visit(*this);

	outBuf[reservedI  ] = *rResult + " = alloca " + tExpr.type;
	outBuf[reservedI+1] = "br " + boolExpr.toStr() + ", label " + *rTrue + ", label " + *rFalse + "\n";
	outBuf[reservedI+2] = "v" + rTrue.headless() + ":";

	outBuf.push_back("store " + tExpr.toStr() + ", ptr " + *rResult);
	outBuf.push_back("br label " + *rCont + "\n");

	if (args.size() > 2) {
		outBuf.push_back("v" + rFalse.headless() + ":");
		LLVMRet fExpr = args[2]->Visit(*this);
		outBuf.push_back("store " + fExpr.toStr() + ", ptr " + *rResult);
		outBuf.push_back("br label " + *rCont + "\n");
	}

	outBuf.push_back("v" + rCont.headless() + ":");
	Register finalRegister = lreg++;
	outBuf.push_back(*finalRegister + " = load " + tExpr.type + ", ptr " + *rResult);
	return { tExpr.type, *finalRegister };
}

LLVMRet LLVMVisitor::nameAsVar(const LLVMRet& ret) {
	if (!locals.empty()) {
		auto asLoc = locals.top()->find("%_" + ret.addr);
		if (asLoc != locals.top()->end())
			return asLoc->second;
	}

	auto asVar = vars_user.find("@" + ret.addr);
	if (asVar != vars_user.end())
		return asVar->second;

	auto asFunc = funcs_user.find("@" + ret.addr);
	if (asFunc != funcs_user.end())
		return asFunc->second;

	// TODO: try to figure out type (f_defun)
	LLVMRet newLocal = { "i32", "%_" + ret.addr };
	if (!locals.empty()) {
		locals.top()->insert({ newLocal.addr, newLocal });
	}
	return newLocal;
}

LLVMRet LLVMVisitor::aritop(const std::string name, const std::vector<Value*>& args, unsigned astart) {
	LLVMRet leftRet = args[astart]->Visit(*this);
	LLVMRet rightRet = (args.size() - astart > 2) ? aritop(name, args, astart+1)
	                                              : args[astart+1]->Visit(*this);

	if (leftRet.type == "name")  leftRet  = nameAsVar(leftRet);
	if (rightRet.type == "name") rightRet = nameAsVar(rightRet);

	Register r = lreg++;
	outBuf.push_back(*r + " = " + ((leftRet.type == "double") ? "f" : "") + name + " " + leftRet.toStr() + ", " + rightRet.addr);

	return { leftRet.type, *r };
}

LLVMRet LLVMVisitor::f_add(funcargs args) {
	return aritop("add", args);
}

LLVMRet LLVMVisitor::f_sub(funcargs args) {
	return aritop("sub", args);
}

LLVMRet LLVMVisitor::f_mul(funcargs args) {
	return aritop("mul", args);
}

LLVMRet LLVMVisitor::f_div(funcargs args) {
	return aritop("div", args);
}

LLVMRet LLVMVisitor::boolop(const std::string name, funcargs args) {
	LLVMRet leftArg = args[0]->Visit(*this),
	        rightArg = args[1]->Visit(*this);
	Register r = lreg++;

	outBuf.push_back(*r + " = " + ((leftArg.type == "double") ? "f" : "i") + "cmp " + ((leftArg.type == "double") ? "o" : (name != "eq") ? "s" : "") + name + " " + leftArg.toStr() + ", " + rightArg.addr);
	return { "i1", *r };
}

LLVMRet LLVMVisitor::f_eq(funcargs args) {
	return boolop("eq", args);
}

LLVMRet LLVMVisitor::f_gt(funcargs args) {
	return boolop("gt", args);
}

LLVMRet LLVMVisitor::f_ge(funcargs args) {
	return boolop("ge", args);
}

LLVMRet LLVMVisitor::f_lt(funcargs args) {
	return boolop("lt", args);
}

LLVMRet LLVMVisitor::f_le(funcargs args) {
	return boolop("le", args);
}

LLVMRet LLVMVisitor::Visit(const Funcall& vfuncall) {
	// Direct functions generate their own LLVM IR
	auto directFuncIt = funcs_direct.find(vfuncall.get_funname());
	if (directFuncIt != funcs_direct.end())
		return directFuncIt->second(*this, vfuncall.get_args());

	auto userFuncIt = funcs_user.find(vfuncall.get_funname());
	if (userFuncIt == funcs_user.end())
		throw std::runtime_error("Function isn't defined!");

	/* Transform into IR of structure:
	 * %0 = call T @func(T %a1, T %a2)
	 */

	std::vector<LLVMRet> pargs;
	for (Value* arg : vfuncall.get_args())
		pargs.push_back(arg->Visit(*this));

	std::string valReg = *(lreg++);
	std::string out = valReg + " = call " + userFuncIt->second.toStr() + "(";
	for (LLVMRet& arg : pargs) {
		out += arg.toStr();
		if (arg.addr != pargs.back().addr)
			out += ", ";
	}
	outBuf.push_back(out += ")");
	return { userFuncIt->second.type, valReg };
}

LLVMVisitor::LLVMVisitor(const char* fileName) : outFile(std::ofstream(fileName)), greg("@v"), lreg("%v"), funcs_direct(), funcs_user(), vars_user(), outBuf() {
	if (!outFile.is_open())
		throw std::runtime_error("Couldn't open file for writing!");

	funcs_direct.insert({"+", &LLVMVisitor::f_add});
	funcs_direct.insert({"-", &LLVMVisitor::f_sub});
	funcs_direct.insert({"*", &LLVMVisitor::f_mul});
	funcs_direct.insert({"/", &LLVMVisitor::f_div});
	funcs_direct.insert({"defun", &LLVMVisitor::f_defun});
	funcs_direct.insert({"defvar", &LLVMVisitor::f_defvar});
	funcs_direct.insert({"if", &LLVMVisitor::f_if});
	funcs_direct.insert({"=", &LLVMVisitor::f_eq});
	funcs_direct.insert({">", &LLVMVisitor::f_gt});
	funcs_direct.insert({">=", &LLVMVisitor::f_ge});
	funcs_direct.insert({"<", &LLVMVisitor::f_lt});
	funcs_direct.insert({"<=", &LLVMVisitor::f_le});
}

void LLVMVisitor::insertMain() {
	outFile << "define i32 @main() {\n";
	for (std::string& codeline : outBuf) {
		outFile << "\t" << codeline << "\n";
	}
	outFile << "\tret i32 0\n}\n";
}

LLVMRet LLVMVisitor::Visit(const Integer& vint) {
	return { "i32", std::to_string(vint.get_value()) };
}

LLVMRet LLVMVisitor::Visit(const Float& vfloat) {
	return { "double", std::to_string(vfloat.get_value()) };
}

LLVMRet LLVMVisitor::Visit(const Name& vname) {
	return { "name", vname.get_value() };
}

LLVMRet LLVMVisitor::Visit(const List& vlist) {
	return { "", "" }; // TODO
}