#include "PlanschbeckenJIT.h"
#pragma warning(push, 0)
#include <llvm/IR/DiagnosticInfo.h>
#include <llvm/IRReader/IRReader.h>

#include <llvm/Transforms/Coroutines.h>
#include <llvm/Transforms/Scalar/GVN.h>
#include <llvm/Transforms/Instrumentation.h>
#include <llvm/Transforms/InstCombine/InstCombine.h>
#include <llvm/Transforms/IPO.h>
#include <llvm/Transforms/ObjCARC.h>
#include <llvm/Transforms/Scalar.h>
#include <llvm/Transforms/Utils.h>
#include <llvm/Transforms/Vectorize.h>
#pragma warning(pop)

#include <Windows.h>

const int value = 123456;

namespace llvm
{
namespace orc
{
llvm::Value* createAddress(uintptr_t adr, llvm::Type *type, llvm::LLVMContext &context)
{
	return llvm::ConstantExpr::getIntToPtr(
		llvm::ConstantInt::get(llvm::IntegerType::get(context, (sizeof(void*)*8)), adr), type);
}

std::unique_ptr<llvm::Module> resolveModule(std::unique_ptr<llvm::Module> module, llvm::LLVMContext &context)
{
#ifdef RESOLVE_IN_MODULE
	auto iter = module->global_begin();
	while(iter != module->global_end())
	{
		if(iter->getLinkage() == GlobalValue::LinkageTypes::ExternalLinkage)
		{
			const char*const symbolName = iter->getName().data();

			printf("Name: %s\n", symbolName);
			if(strcmp(symbolName, "?planschiValue@@3HB") == 0)
			{
				iter->replaceAllUsesWith(createAddress((uintptr_t)&value, iter->getType(), context));
			}
		}
		iter++;
	}
#endif // RESOLVE_IN_MODULE

	return module;
}	

SymbolNameSet resolveDylib(JITDylib &Parent, const SymbolNameSet &Names)
{
	HMODULE curProcess = GetModuleHandle(nullptr);

	orc::SymbolNameSet Added;
	orc::SymbolMap     NewSymbols;

	for(auto &Name : Names) 
	{
		const char*const functionName = (*Name).data();
		const uintptr_t  adr = (strcmp(functionName, "?planschiValue@@3HB") == 0) ?
								(uintptr_t)&value : (uintptr_t)GetProcAddress(curProcess, functionName);
		
		if(adr)
		{
			Added.insert(Name);
			NewSymbols[Name] = JITEvaluatedSymbol(adr, JITSymbolFlags::Exported);
		}

		else
		{
			printf("Not found: %s 0x%p\n", functionName, curProcess);
		}		
	}

	Parent.define(absoluteSymbols(std::move(NewSymbols)));
	return Added;
}

inline std::unique_ptr<llvm::Module> loadModule(const char *file, llvm::LLVMContext &context)
{
	llvm::SMDiagnostic error;
	return resolveModule(llvm::parseIRFile(file, error, context), context);
}

PlanschbeckenJIT ::	PlanschbeckenJIT(JITTargetMachineBuilder JTMB, DataLayout dl) 
		: objectLayer(this->es, [](){return std::make_unique<SectionMemoryManager>();}),
		  compileLayer(this->es, this->objectLayer, ConcurrentIRCompiler(std::move(JTMB))),
		  transformLayer(this->es, this->compileLayer, &PlanschbeckenJIT::optimizeModule),
		  dl(std::move(dl)), mangle(this->es, this->dl), ctx(std::make_unique<LLVMContext>())
{
	this->es.getMainJITDylib().setGenerator(
#ifdef RESOLVE_IN_MODULE
		cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(this->dl))
#else
		resolveDylib
#endif //RESOLVE_IN_MODULE
	);
}

void PlanschbeckenJIT :: addModule(const char *file)
{
	cantFail(this->transformLayer.add(this->es.getMainJITDylib(), ThreadSafeModule(std::move(loadModule(file, *(this->ctx).getContext())), this->ctx)));
}

Expected<ThreadSafeModule> PlanschbeckenJIT :: optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) 
{
	// Create a function pass manager.
	auto FPM = llvm::make_unique<legacy::FunctionPassManager>(TSM.getModule());

	// Add some optimizations.
	FPM->add(createInstructionCombiningPass());
	FPM->add(createReassociatePass());
	FPM->add(createGVNPass());
	FPM->add(createCFGSimplificationPass());
	FPM->doInitialization();

	// Run the optimizations over all functions in the module being added to
	// the JIT.
	for (auto &F : *TSM.getModule())
		FPM->run(F);

	return TSM;
}
}
}