#include "linker.hpp"
#include "image_util.hpp"
#include "std.hpp"

using namespace std;

class BBModule : public Module {
	public:
	BBModule();
	BBModule(istream& in);
	~BBModule();

	void* link(Module* libs);
	bool  createExe(const char* exe_file, const char* dll_file);

	int getPC();

	void emit(int byte);
	void emitw(int word);
	void emitd(int dword);
	void emitx(void* mem, int sz);
	bool addSymbol(const char* sym, int pc);
	bool addReloc(const char* dest_sym, int pc, bool pcrel);

	bool findSymbol(const char* sym, int* pc);

	private:
	char* data;
	int   data_sz, pc;
	bool  linked;

	map<string, int> symbols;
	map<int, string> rel_relocs, abs_relocs;

	bool findSym(const string& t, Module* libs, int* n)
	{
		if (findSymbol(t.c_str(), n))
			return true;
		if (libs->findSymbol(t.c_str(), n))
			return true;
		string err = "Symbol '" + t + "' not found";
		MessageBox(GetDesktopWindow(), err.c_str(), "Blitz Linker Error", MB_TOPMOST | MB_SETFOREGROUND);
		return false;
	}

	void ensure(int n)
	{
		if (pc + n <= data_sz)
			return;
		data_sz = data_sz / 2 + data_sz;
		if (data_sz < pc + n)
			data_sz = pc + n;
		char* old_data = data;
		data           = new char[data_sz];
		memcpy(data, old_data, pc);
		delete old_data;
	}
};

BBModule::BBModule() : data(0), data_sz(0), pc(0), linked(false) {}

BBModule::~BBModule()
{
	if (linked)
		VirtualFree(data, 0, MEM_RELEASE);
	else
		delete[] data;
}

void* BBModule::link(Module* libs)
{
	if (linked)
		return data;

	int                        dest;
	map<int, string>::iterator it;

	char* p = (char*)VirtualAlloc(0, pc, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
	memcpy(p, data, pc);
	delete[] data;
	data = p;

	linked = true;

	for (it = rel_relocs.begin(); it != rel_relocs.end(); ++it) {
		if (!findSym(it->second, libs, &dest))
			return 0;
		int* p = (int*)(data + it->first);
		*p += (dest - (int)p);
	}

	for (it = abs_relocs.begin(); it != abs_relocs.end(); ++it) {
		if (!findSym(it->second, libs, &dest))
			return 0;
		int* p = (int*)(data + it->first);
		*p += dest;
	}

	return data;
}

int BBModule::getPC()
{
	return pc;
}

void BBModule::emit(int byte)
{
	ensure(1);
	data[pc++] = byte;
}

void BBModule::emitw(int word)
{
	ensure(2);
	*(short*)(data + pc) = word;
	pc += 2;
}

void BBModule::emitd(int dword)
{
	ensure(4);
	*(int*)(data + pc) = dword;
	pc += 4;
}

void BBModule::emitx(void* mem, int sz)
{
	ensure(sz);
	memcpy(data + pc, mem, sz);
	pc += sz;
}

bool BBModule::addSymbol(const char* sym, int pc)
{
	string t(sym);
	if (symbols.find(t) != symbols.end())
		return false;
	symbols[t] = pc;
	return true;
}

bool BBModule::addReloc(const char* dest_sym, int pc, bool pcrel)
{
	map<int, string>& rel = pcrel ? rel_relocs : abs_relocs;
	if (rel.find(pc) != rel.end())
		return false;
	rel[pc] = string(dest_sym);
	return true;
}

bool BBModule::findSymbol(const char* sym, int* pc)
{
	string                     t  = string(sym);
	map<string, int>::iterator it = symbols.find(t);
	if (it == symbols.end())
		return false;
	*pc = it->second + (int)data;
	return true;
}

int Linker::version()
{
	return VERSION;
}

bool Linker::canCreateExe()
{
#ifdef DEMO
	return false;
#else
	return true;
#endif
}

Module* Linker::createModule()
{
	return new BBModule();
}

void Linker::deleteModule(Module* mod)
{
	delete mod;
}

Linker* _cdecl linkerGetLinker()
{
	static Linker linker;
	return &linker;
}

bool BBModule::createExe(const char* exe_file, const char* dll_file)
{
#ifdef DEMO
	return false;
#else

	//find proc address of bbWinMain
	HMODULE hmod = LoadLibrary(dll_file);
	if (!hmod)
		return false;
	int proc  = (int)GetProcAddress(hmod, "_bbWinMain@0");
	int entry = proc - (int)hmod;
	FreeLibrary(hmod);
	if (!proc)
		return false;

	if (!CopyFile(dll_file, exe_file, false))
		return false;

	if (!openImage(exe_file))
		return false;

	makeExe(entry);

	//create module
	//code size: code...
	//num_syms:  name,val...
	//num_rels:  name,val...
	//num_abss:  name,val...
	//
	qstreambuf buf;
	iostream   out(&buf);

	map<string, int>::iterator it;
	map<int, string>::iterator rit;

	//write the code
	int sz = pc;
	out.write((char*)&sz, 4);
	out.write(data, pc);

	//write symbols
	sz = symbols.size();
	out.write((char*)&sz, 4);
	for (it = symbols.begin(); it != symbols.end(); ++it) {
		string t = it->first + '\0';
		out.write(t.data(), t.size());
		sz = it->second;
		out.write((char*)&sz, 4);
	}

	//write relative relocs
	sz = rel_relocs.size();
	out.write((char*)&sz, 4);
	for (rit = rel_relocs.begin(); rit != rel_relocs.end(); ++rit) {
		string t = rit->second + '\0';
		out.write(t.data(), t.size());
		sz = rit->first;
		out.write((char*)&sz, 4);
	}

	//write absolute relocs
	sz = abs_relocs.size();
	out.write((char*)&sz, 4);
	for (rit = abs_relocs.begin(); rit != abs_relocs.end(); ++rit) {
		string t = rit->second + '\0';
		out.write(t.data(), t.size());
		sz = rit->first;
		out.write((char*)&sz, 4);
	}

	replaceRsrc(10, 1111, 1033, buf.data(), buf.size());

	closeImage();

	return true;
#endif
}