#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> locals;
	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;
		locals.insert({ name, { "i32", name } });
		if (v != fargs->get_value().back())
			outFile << ", ";
	}
	outFile << ") {\n";

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

	outFile << "\tret " << funcRet << "\n}\n\n";
	return { funcRet.type, "@" + fname->get_value() };
}

LLVMRet LLVMVisitor::nameAsVar(const Name* name, std::map<std::string, LLVMRet>& locals) {
	auto asLoc = locals.find("%_" + name->get_value());
	if (asLoc != locals.end())
		return asLoc->second;

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

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

	// TODO: try to figure out type (f_defun)
	LLVMRet newLocal = { "i32", "%_" + name->get_value() };
	locals.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);
	
	Register r = lreg++;
	outBuf.push_back(*r + " = " + ((leftRet.type == "double") ? "f" : "") + name + " " + leftRet.toStr() + ", " + rightRet.addr + "\n");

	return { leftRet.type, *r };
}

LLVMRet LLVMVisitor::f_add(const std::vector<Value*>& args) {
	return aritop("add", args);
}

LLVMRet LLVMVisitor::f_sub(const std::vector<Value*>& args) {
	return aritop("sub", 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 += ")\n");
	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({"defun", &LLVMVisitor::f_defun});
}

void LLVMVisitor::insertMain() {
	outFile << "define i32 @main() {\n";
	for (std::string& codeline : outBuf) {
		outFile << codeline;
	}
	outFile << "ret 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
}