#include "obse_common/SafeWrite.h"
#include "Hooks_Script.h"
#include "GameForms.h"
#include "Script.h"
#include "ScriptUtils.h"
#include "CommandTable.h"
#include <stack>
#include <string>

// a size of ~1KB should be enough for a single line of code
char s_ExpressionParserAltBuffer[0x500] = {0};

#if OBLIVION

#include "StringVar.h"

#if OBLIVION_VERSION == OBLIVION_VERSION_1_2_416
	const UInt32 ExtractStringPatchAddr = 0x004FB1EB;
	const UInt32 ExtractStringRetnAddr = 0x004FB1F4;

	static const UInt32 kResolveRefVarPatchAddr		= 0x004FA9F1;
	static const UInt32 kResolveNumericVarPatchAddr = 0x004FA11E;
	static const UInt32 kEndOfLineCheckPatchAddr	= 0;	// not yet supported at run-time

	// incremented on each recursive call to Activate, limit of 5 hard-coded
	static UInt32* kActivationRecurseDepth = (UInt32*)0x00B35F00;

	static const UInt32 kExpressionParserBufferOverflowHookAddr_1 = 0x004F40F9;
	static const UInt32 kExpressionParserBufferOverflowRetnAddr_1 = 0x004F40FE;

	static const UInt32 kExpressionParserBufferOverflowHookAddr_2 = 0x004F42A4;
	static const UInt32 kExpressionParserBufferOverflowRetnAddr_2 = 0x004F42A9;
#else
#error unsupported Oblivion version
#endif	// OBLIVION_VERSION

static void __stdcall DoExtractString(char* scriptData, UInt32 dataLen, char* dest, ScriptEventList* eventList)
{
	// copy the string
	memcpy(dest, scriptData, dataLen);
	dest[dataLen] = 0;
	
	if (dataLen && dest[0] == '$' && eventList && eventList->m_script)	// variable name
	{
		Script::VariableInfo* varInfo = NULL;
		varInfo = eventList->m_script->GetVariableByName(dest + 1);
		if (varInfo)
		{
			ScriptEventList::Var* var;
			var = eventList->GetVariable(varInfo->idx);
			if (var)
			{
				StringVar* strVar;
				strVar = g_StringMap.Get(var->data);
				if (strVar)
					if (strVar->GetLength() < 0x100)		// replace string with contents of string var
						strcpy_s(dest, strVar->GetLength() + 1, strVar->GetCString());
			}
		}
	}			// "%e" becomes an empty string
	else if (dataLen == 2 && dest[0] == '%' && toupper(dest[1]) == 'E')
		dest[0] = 0;
}

static __declspec(naked) void ExtractStringHook(void)
{
	// This hooks immediately before a call to memcpy().
	// Original code copies a string literal from script data to buffer
	// If string is of format $localVariableName, replace literal with contents of string var

	static char* scriptData;	// pointer to beginning of string in script data

	__asm {
		// Grab the args to memcpy()
		mov scriptData, eax
		mov eax, [esp+0x50]		// ScriptEventList
		pushad

		push eax
		push edi				// destination buffer
		push ebp				// data len
		mov eax, scriptData
		push eax
		call DoExtractString

		popad
		mov eax, scriptData
		jmp [ExtractStringRetnAddr]
	}
}

static __declspec(naked) void ExpressionParserBufferOverflowHook_1(void)
{
	__asm {
		lea	eax, s_ExpressionParserAltBuffer		// swap buffers
		push eax
		jmp	[kExpressionParserBufferOverflowRetnAddr_1]
	}
}

static __declspec(naked) void ExpressionParserBufferOverflowHook_2(void)
{
	__asm {
		push ebp
		lea edx, s_ExpressionParserAltBuffer
		jmp	[kExpressionParserBufferOverflowRetnAddr_2]
	}
}

void Hook_Script_Init()
{
	WriteRelJump(ExtractStringPatchAddr, (UInt32)&ExtractStringHook);

	// patch the "apple bug"
	// game caches information about the most recently retrieved RefVariable for the current executing script
	// if same refIdx requested twice in a row returns previously returned ref without
	// bothering to check if form stored in ref var has changed
	// this fixes it by overwriting a conditional jump with an unconditional one
	SafeWrite8(kResolveRefVarPatchAddr, 0xEB);

	// game also caches information about the most recently retrieved local numeric variable for
	// currently executing script. Causes issues with function scripts. As above, overwrite conditional jump with unconditional
	SafeWrite8(kResolveNumericVarPatchAddr, 0xEB);

	// hook code in the vanilla expression parser's subroutine to fix the buffer overflow
	WriteRelJump(kExpressionParserBufferOverflowHookAddr_1, (UInt32)&ExpressionParserBufferOverflowHook_1);
	WriteRelJump(kExpressionParserBufferOverflowHookAddr_2, (UInt32)&ExpressionParserBufferOverflowHook_2);
}

void ResetActivationRecurseDepth()
{
	// reset to circumvent the hard-coded limit
	*kActivationRecurseDepth = 0;
}

#else	// CS-stuff

#include "PluginManager.h"

#if CS_VERSION == CS_VERSION_1_2
	static const UInt32 kEndOfLineCheckPatchAddr = 0x00501E22;
	static const UInt32 kCopyStringArgHookAddr		 = 0x00501A8A;

	static const UInt32 kBeginScriptCompilePatchAddr = 0x005034A9;	// calls CompileScript()
	static const UInt32 kBeginScriptCompileCallAddr  = 0x00503330;	// bool __fastcall CompileScript(unk, unk, Script*, ScriptBuffer*)
	static const UInt32 kBeginScriptCompileRetnAddr	 = 0x005034AE;

	static const UInt32 kExpressionParserBufferOverflowHookAddr_1 = 0x004F9712;
	static const UInt32 kExpressionParserBufferOverflowRetnAddr_1 = 0x004F9717;

	static const UInt32 kExpressionParserBufferOverflowHookAddr_2 = 0x004F9863;
	static const UInt32 kExpressionParserBufferOverflowRetnAddr_2 = 0x004F986A;
#elif CS_VERSION == CS_VERSION_1_0
	static const UInt32 kEndOfLineCheckPatchAddr = 0x004F7E77;
	static const UInt32 kEndOfLineCheckJumpDelta = 0xFFFFFCF8;
	static const UInt32 kEndOfLineCheckJumpAddr		 = 0x004F7B75;
	
	static const UInt32 kBeginScriptCompilePatchAddr = 0x004F935E;	// calls CompileScript()
	static const UInt32 kBeginScriptCompileCallAddr  = 0x004F9200;	// bool __fastcall CompileScript(unk, unk, Script*, ScriptBuffer*)
	static const UInt32 kBeginScriptCompileRetnAddr  = 0x004F9363;
#else
#error unsupported CS version
#endif	// CS_VERSION

static __declspec(naked) void ExpressionParserBufferOverflowHook_1(void)
{
	__asm {
		lea	edx, s_ExpressionParserAltBuffer
		push edx
		jmp	[kExpressionParserBufferOverflowRetnAddr_1]
	}
}

static __declspec(naked) void ExpressionParserBufferOverflowHook_2(void)
{
	__asm {
		lea	edx, s_ExpressionParserAltBuffer
		mov	byte ptr [esi], 0x20
		jmp	[kExpressionParserBufferOverflowRetnAddr_2]
	}
}

#endif	// OBLIVION

// Patch compiler check on end of line when calling commands from within other commands
// TODO: implement for run-time compiler
// Hook differs for 1.0/1.2 CS versions as the patched code differs
#if CS_VERSION == CS_VERSION_1_0

void PatchEndOfLineCheck(bool bDisableCheck)
{
	if (bDisableCheck)
		WriteRelJump(kEndOfLineCheckPatchAddr, kEndOfLineCheckJumpAddr);
	else
	{
		SafeWrite8(kEndOfLineCheckPatchAddr, 0x0F);
		SafeWrite8(kEndOfLineCheckPatchAddr + 1, 0x83);
		SafeWrite32(kEndOfLineCheckPatchAddr + 2, kEndOfLineCheckJumpDelta);
	}
}

#elif CS_VERSION == CS_VERSION_1_2

void PatchEndOfLineCheck(bool bDisableCheck)
{

	if (bDisableCheck)
		SafeWrite8(kEndOfLineCheckPatchAddr, 0xEB);		// unconditional (short) jump
	else
		SafeWrite8(kEndOfLineCheckPatchAddr, 0x73);		// conditional jnb (short)
}

#else

void PatchEndOfLineCheck(bool bDisableCheck)
{
	// ###TODO: implement for run-time
}

#endif

static bool s_bParsingExpression = false;

bool ParsingExpression()
{
	return s_bParsingExpression;
}

bool ParseNestedFunction(CommandInfo* cmd, ScriptLineBuffer* lineBuf, ScriptBuffer* scriptBuf)
{
	// disable check for end of line
	PatchEndOfLineCheck(true);

	s_bParsingExpression = true;
	bool bParsed = cmd->parse(cmd->numParams, cmd->params, lineBuf, scriptBuf);
	s_bParsingExpression = false;

	// re-enable EOL check
	PatchEndOfLineCheck(false);

	return bParsed;
}

static std::stack<UInt8*> s_loopStartOffsets;

void RegisterLoopStart(UInt8* offsetPtr)
{
	s_loopStartOffsets.push(offsetPtr);
}

bool HandleLoopEnd(UInt32 offsetToEnd)
{
	if (!s_loopStartOffsets.size())
		return false;

	UInt8* startPtr = s_loopStartOffsets.top();
	s_loopStartOffsets.pop();

	*((UInt32*)startPtr) = offsetToEnd;
	return true;
}

#if !OBLIVION

bool __stdcall HandleBeginCompile(ScriptBuffer* buf)
{
	// empty out the loop stack
	while (s_loopStartOffsets.size())
		s_loopStartOffsets.pop();

	// Preprocess the script:
	//  - check for usage of array variables in Set statements (disallowed)
	//  - check loop structure integrity
	//  - check for use of ResetAllVariables on scripts containing string/array vars

	bool bResult = PrecompileScript(buf);
	if (bResult) {
		PluginManager::Dispatch_Message(0, OBSEMessagingInterface::kMessage_Precompile, buf, sizeof(buf), NULL);
	}

	return bResult;
}

static __declspec(naked) void CompileScriptHook(void)
{
	static bool precompileResult;

	__asm	
	{
		mov		eax,	[esp+4]					// grab the second arg (ScriptBuffer*)
		pushad
		push	eax
		call	HandleBeginCompile				// Precompile
		mov		[precompileResult],	al			// save result
		popad
		call	[kBeginScriptCompileCallAddr]	// let the compiler take over
		test	al, al
		jz		EndHook							// return false if CompileScript() returned false
		mov		al,	[precompileResult]			// else return result of Precompile
	EndHook:
		jmp		[kBeginScriptCompileRetnAddr]
	}
}

// replace special characters ("%q" -> '"', "%r" -> '\n')
UInt32 __stdcall CopyStringArg(char* dest, const char* src, UInt32 len, ScriptLineBuffer* lineBuf)
{
	if (!src || !len || !dest)
		return len;

	std::string str(src);
	UInt32 pos = 0;

	while ((pos = str.find('%', pos)) != -1 && pos < str.length() - 1)
	{
		char toInsert = 0;
		switch (str[pos + 1])
		{
		case '%':
			pos += 2;
			continue;
		case 'r':
		case 'R':
			toInsert = '\n';
			break;
		case 'q':
		case 'Q':
			toInsert = '"';
			break;
		default:
			pos += 1;
			continue;
		}
		
		str.insert(pos, 1, toInsert);	// insert char at current pos
		str.erase(pos + 1, 2);			// erase format specifier
		pos += 1;
	}

	// copy the string to script data
	memcpy(dest, str.c_str(), str.length());

	// write length of string
	lineBuf->dataOffset -= 2;
	lineBuf->Write16(str.length());

	return str.length();
}

// major code differences between CS versions here so hooks differ significantly
static __declspec(naked) void __cdecl CopyStringArgHook(void)
{
#if CS_VERSION == CS_VERSION_1_2	
	// overwrite call to memcpy()

	// On entry:
	//	eax: dest buffer
	//	edx: src string
	//	edi: string len
	//  esi: ScriptLineBuffer* (must be preserved)
	// edi must be updated to reflect length of modified string (added to dataOffset on return)

	__asm
	{
		push	esi				

		push	esi
		push	edi
		push	edx
		push	eax
		call	CopyStringArg

		mov		edi, eax
		pop		esi

		retn
	}
#elif CS_VERSION == CS_VERSION_1_0	
	// string copy is done inline by CS 1.0 using rep movsd
	// need to write length of string to scriptdata before writing string

	// ###TODO

#else
#error unsupported CS version
#endif
}

void Hook_Compiler_Init()
{
	// hook beginning of compilation process
	WriteRelJump(kBeginScriptCompilePatchAddr, (UInt32)&CompileScriptHook);

	// hook copying of string argument to compiled data
	// lets us modify the string before its copied
#if CS_VERSION == CS_VERSION_1_2		// 1.0 hook not yet implemented
	WriteRelCall(kCopyStringArgHookAddr, (UInt32)&CopyStringArgHook);
#endif

	// hook code in the vanilla expression parser's subroutine to fix the buffer overflow
	WriteRelJump(kExpressionParserBufferOverflowHookAddr_1, (UInt32)&ExpressionParserBufferOverflowHook_1);
	WriteRelJump(kExpressionParserBufferOverflowHookAddr_2, (UInt32)&ExpressionParserBufferOverflowHook_2);
}

#endif