#include <llvm/ADT/ArrayRef.h>
#include <llvm/LLVMContext.h>
#include <llvm/Module.h>
#include <llvm/Function.h>
#include <llvm/BasicBlock.h>
#include <llvm/IRBuilder.h>
#include <llvm/Support/raw_ostream.h>

#include <memory>
#include <vector>
#include <string>
#include <fstream>

using namespace llvm;

class Brainfuck{
	// Create a module
	LLVMContext & context {getGlobalContext()};
	Module module {"asdf", context};

	// Get some basic types and sizes
	Type* bool_type {IntegerType::getInt1Ty(context)};
	Type* char_type {IntegerType::getInt8Ty(context)};
	Type* int32_type {IntegerType::getInt32Ty(context)};
	Value* char_size {ConstantExpr::getTruncOrBitCast(ConstantExpr::getSizeOf(char_type), int32_type)};

	// Get some basic functions
	Function* getchar_func {cast<Function>(module.getOrInsertFunction("getchar", IntegerType::getInt32Ty(context), nullptr))};
	Function* putchar_func {cast<Function>(module.getOrInsertFunction("putchar", IntegerType::getInt32Ty(context), IntegerType::getInt32Ty(context), nullptr))};
	Function* main_function {cast<Function>(module.getOrInsertFunction("main", Type::getInt32Ty(context), nullptr))};

	// Builder for the main function
	BasicBlock* entry {BasicBlock::Create(context, "", main_function)};
	IRBuilder<> builder {entry};

	// Malloc memory
	const int memtotal = 100;
	Value* total_memory {ConstantInt::get(int32_type, memtotal)};
	Value* ptr_arr {CallInst::CreateMalloc(entry, int32_type, char_type, char_size, total_memory)};
	Value* curhead {nullptr};

public:
	Brainfuck(){
		// Do the malloc thing and memset it
		entry->getInstList().push_back(cast<Instruction>(ptr_arr));
		builder.CreateMemSet(ptr_arr, ConstantInt::get(char_type, 0), total_memory, 1);

		// Get head in the middle
		curhead = builder.CreateGEP(ptr_arr, ConstantInt::get(int32_type, memtotal/2));
	}

	void fuck_it(std::string const & program){
		// Parse brainfuck program
		fuck_it_rec(program, program.begin());

		// free and return 0
		entry->getInstList().push_back(CallInst::CreateFree(ptr_arr, entry));
		builder.CreateRet(ConstantInt::get(int32_type, 0));
	}

	void print(raw_ostream& os){
		module.print(os, nullptr);
	}

private:
	void dec(){
		auto tape_0 = builder.CreateLoad(curhead);
		auto tape_1 = builder.CreateSub(tape_0, ConstantInt::get(char_type, 1));
		builder.CreateStore(tape_1, curhead);
	}

	void inc(){
		auto tape_0 = builder.CreateLoad(curhead);
		auto tape_1 = builder.CreateAdd(tape_0, ConstantInt::get(char_type, 1));
		builder.CreateStore(tape_1, curhead);
	}

	void put_it(){
		auto tape_0 = builder.CreateLoad(curhead);
		auto tape_1 = builder.CreateSExt(tape_0, int32_type);
		builder.CreateCall(putchar_func, tape_1);
	}

	void get_it(){
		auto tape_0 = builder.CreateCall(getchar_func);
		auto tape_1 = builder.CreateTrunc(tape_0, char_type);
		builder.CreateStore(tape_1, curhead);
	}

	std::string::const_iterator loop_it(std::string const & program, std::string::const_iterator it){
		// I should merge some bits here... to inside the first label.
		auto tape_0 = builder.CreateLoad(curhead);
		auto cond = builder.CreateICmpNE(tape_0, ConstantInt::get(char_type, 0));
		auto true_bb = BasicBlock::Create(context, "", main_function);
		auto false_bb = BasicBlock::Create(context, "", main_function);
		builder.CreateCondBr(cond, true_bb, false_bb);

		builder.SetInsertPoint(true_bb);
		it = fuck_it_rec(program, it);
		tape_0 = builder.CreateLoad(curhead);
		cond = builder.CreateICmpNE(tape_0, ConstantInt::get(char_type, 0));
		builder.CreateCondBr(cond, true_bb, false_bb);

		entry = false_bb;
		builder.SetInsertPoint(false_bb);

		return it;
	}

	std::string::const_iterator fuck_it_rec(std::string const & program, std::string::const_iterator it){
		for(; it != program.end();){
			auto c = *it++;
			switch(c){
				case ']' : return it;
				case '>' : curhead = builder.CreateGEP(curhead, ConstantInt::get(int32_type, 1)); break;
				case '<' : curhead = builder.CreateGEP(curhead, ConstantInt::get(int32_type, -1)); break;
				case '+' : inc(); break;
				case '-' : dec(); break;
				case '.' : put_it(); break;
				case ',' : get_it(); break;
				case '[' : it = loop_it(program, it); break;
				default : break;
			}
		}
		return program.end();
	}
};

int main(int argc, char ** argv){
	assert(argc == 2);
	std::string filename = argv[1];

	std::ifstream file{filename, std::ios::binary};
    std::string source{std::istreambuf_iterator<char>(file), std::istreambuf_iterator<char>()};

	// Initialise compiler
	Brainfuck brain;

	// Compile!
	brain.fuck_it(source);

	std::string error_info;
	raw_fd_ostream os{(filename + ".ll").c_str(), error_info};
	brain.print(os);
}