#include "Hooks_SaveLoad.h"
#include "obse_common/SafeWrite.h"
#include "GameAPI.h"
#include "NiNodes.h"
#include "Serialization.h"
#include "GameTasks.h"
#include "PluginManager.h"

UInt32 g_gameLoaded = 0;

#if OBLIVION_VERSION == OBLIVION_VERSION_1_1

static const UInt32	kLoadGamePatchAddr =	0x0045FAA5;
static const UInt32	kLoadGameRetnAddr =		0x0045FAAA;
#define	kLoadGameEBPOffset					0x44

static const UInt32	kSaveGamePatchAddr =	0x00460813;
static const UInt32	kSaveGameRetnAddr =		0x00460818;

static const UInt32	kDeleteGamePatchAddr =	0x00453ED6;

static const UInt32	kRenameGamePatchAddr =	0x0045C1DA;

#elif OBLIVION_VERSION == OBLIVION_VERSION_1_2

static const UInt32	kLoadGamePatchAddr =	0x00466833;
static const UInt32	kLoadGameRetnAddr =		0x00466838;
#define	kLoadGameEBPOffset					0x40

static const UInt32	kSaveGamePatchAddr =	0x004656D8;
static const UInt32	kSaveGameRetnAddr =		0x004656DD;

static const UInt32	kDeleteGamePatchAddr =	0x00453496;

static const UInt32	kRenameGamePatchAddr =	0x0045F600;

#elif OBLIVION_VERSION == OBLIVION_VERSION_1_2_416

static const UInt32	kLoadGamePatchAddr =	0x00466995;
static const UInt32	kLoadGameRetnAddr =		0x0046699A;
#define	kLoadGameEBPOffset					0x40

static const UInt32	kSaveGamePatchAddr =	0x004657D6;
static const UInt32	kSaveGameRetnAddr =		0x004657DB;

static const UInt32	kDeleteGamePatchAddr =	0x004534A6;

static const UInt32	kRenameGamePatchAddr =	0x0045F5C0;

static const UInt32 kPreLoadPatchAddr =		0x00465BDD;
static const UInt32 kPreLoadRetnAddr =		0x00465BE3;

static const UInt32 kPostLoadPatchAddr =	0x00466AAA;
static const UInt32 kPostLoadRetnAddr =		0x00466AAF;

static const UInt32	kScript_RunAddr =		0x004FBE00;	// entry to Script::Run()

static const UInt8 Script_RunPatchInstructions[] =
{
	0xB0, 0x01,				// mov al, 1
	0xC2, 0x10, 0x00		// retn 0x10
};

static const UInt8 Script_RunOverwrittenInstructions[] =
{
	0x56, 0x8B, 0xF1, 0x8B, 0x4C
};

#else

#error unsupported oblivion version

#endif

/*
	When loading a saved game, scripts attached to references are executed as soon as the refr
	is loaded. This can cause CTD in mods which use token scripts.
	It can also cause errors with array/string variable access as the variable data has not yet been loaded
	Disable script execution during game load to prevent.
	Note that events which could potentially occur during the game load (OnLoad, OnReset) will remain
	marked until we re-enable script execution
*/
static bool s_scriptExecutionEnabled = true;

void __stdcall ToggleScriptExecution(bool bEnable)
{
	if (bEnable == s_scriptExecutionEnabled)
		return;

	DEBUG_PRINT("Script Execution Toggled %s", bEnable ? "on" : "off");

	s_scriptExecutionEnabled = bEnable;
	const UInt8* data = bEnable ? Script_RunOverwrittenInstructions : Script_RunPatchInstructions;
	for (UInt32 i = 0; i < sizeof(Script_RunPatchInstructions); i++)
		SafeWrite8(kScript_RunAddr + i, data[i]);
}

void __stdcall DoPreLoadGame(BSFile* file)
{
	PluginManager::Dispatch_Message(0, OBSEMessagingInterface::kMessage_PreLoadGame, file->m_path, strlen(file->m_path), NULL);
}

static __declspec(naked) void PreLoadGameHook(void)
{
	__asm {
		pushad

		push	0
		call	ToggleScriptExecution

		push	esi
		call	DoPreLoadGame

		popad
		mov		edx, g_ioManager
		jmp		[kPreLoadRetnAddr]
	}
}

static __declspec(naked) void PostLoadGameHook(void)
{
	__asm {
		pushad
		push	1
		call	ToggleScriptExecution
		popad
		
		pop ecx
		pop edi
		pop esi
		pop ebp
		pop ebx
		jmp		[kPostLoadRetnAddr]
	}
}

// stdcall to make it easier to work with in asm
static void __stdcall DoLoadGameHook(BSFile * file)
{
	g_gameLoaded = 1;

	_MESSAGE("DoLoadGameHook: %s", file->m_path);

	Serialization::HandleLoadGame(file->m_path);
}

static __declspec(naked) void LoadGameHook(void)
{
	__asm
	{
		pushad
		push		esi		// esi = BSFile *
		call		DoLoadGameHook
		popad

		// overwritten code
		mov			ecx, [ebp + kLoadGameEBPOffset]
		test		ecx, ecx
		jmp			[kLoadGameRetnAddr]
	}
}

static void __stdcall DoSaveGameHook(BSFile * file)
{
	_MESSAGE("DoSaveGameHook: %s", file->m_path);

	Serialization::HandleSaveGame(file->m_path);
}

static __declspec(naked) void SaveGameHook(void)
{
	__asm
	{
		pushad
		push		esi		// esi = BSFile *
		call		DoSaveGameHook
		popad

		// overwritten code
		mov			ecx, [ebp + kLoadGameEBPOffset]
		test		ecx, ecx	// not all versions do exactly this, but they are all equivalent
		jmp			[kSaveGameRetnAddr]
	}
}

// overwriting a call to DeleteFileA
static void __stdcall DeleteGameHook(const char * path)
{
	_MESSAGE("DeleteGameHook: %s", path);

	Serialization::HandleDeleteGame(path);

	// overwritten code
	DeleteFile(path);
}

static void RenameGameHook(const char * oldPath, const char * newPath)
{
	_MESSAGE("RenameGameHook: %s -> %s", oldPath, newPath);

	Serialization::HandleRenameGame(oldPath, newPath);

	rename(oldPath, newPath);
}

void Hook_SaveLoad_Init(void)
{
	WriteRelJump(kLoadGamePatchAddr, (UInt32)&LoadGameHook);
	WriteRelJump(kSaveGamePatchAddr, (UInt32)&SaveGameHook);
	WriteRelCall(kDeleteGamePatchAddr, (UInt32)&DeleteGameHook);
	SafeWrite8(kDeleteGamePatchAddr + 5, 0x90);	// nop out the 6th byte of the patched instruction
	WriteRelCall(kRenameGamePatchAddr, (UInt32)&RenameGameHook);
	WriteRelJump(kPreLoadPatchAddr, (UInt32)&PreLoadGameHook);
	WriteRelJump(kPostLoadPatchAddr, (UInt32)&PostLoadGameHook);
}
