#include "CommandTable.h"
#include "ParamInfos.h"
#include "Commands_General.h"
#include "ScriptUtils.h"
#include "Hooks_Script.h"

#if OBLIVION

#if OBLIVION_VERSION == OBLIVION_VERSION_1_1
	static const UInt32 kDataDeltaStackOffset = 480;
#elif OBLIVION_VERSION == OBLIVION_VERSION_1_2
	static const UInt32 kDataDeltaStackOffset = 482;
#elif OBLIVION_VERSION == OBLIVION_VERSION_1_2_416
	static const UInt32 kDataDeltaStackOffset = 482;
#else
#error unsupported oblivion version
#endif

#include <stack>
#include "GameAPI.h"

#define GET_EXECUTION_STATE(STATE)			\
{											\
	UInt32	_esi;							\
	__asm { mov _esi, esi }					\
	STATE = (ScriptExecutionState*)_esi;		\
}

class Loop
{
public:
	Loop() { }
	virtual ~Loop() { }

	virtual bool Update(COMMAND_ARGS) = 0;
};

class WhileLoop : public Loop
{
	UInt32		m_exprOffset;		// offset of test expression in script data
public:
	WhileLoop(UInt32 exprOffset) : m_exprOffset(exprOffset) { }
	virtual ~WhileLoop() { }

	virtual bool Update(COMMAND_ARGS);
};

class ForEachLoop : public Loop
{
public:
	virtual bool Update(COMMAND_ARGS) = 0;
	virtual bool IsEmpty() = 0;
};

class ArrayIterLoop : public ForEachLoop
{
	ArrayID		m_srcID;
	ArrayID		m_iterID;
	ArrayKey	m_curKey;

	void UpdateIterator(const ArrayElement* elem);
public:
	ArrayIterLoop(ForEachContext* context);
	virtual ~ArrayIterLoop();

	virtual bool Update(COMMAND_ARGS);
	bool IsEmpty() { return (g_ArrayMap.SizeOf(m_srcID) == -1 || g_ArrayMap.SizeOf(m_srcID) == 0);	}
};

class StringIterLoop : public ForEachLoop
{
	std::string		m_src;
	UInt32			m_curIndex;
	UInt32			m_iterID;

public:
	StringIterLoop(ForEachContext* context);
	virtual ~StringIterLoop() { }

	virtual bool Update(COMMAND_ARGS);
	bool IsEmpty() { return m_src.length() == 0; }
};

class LoopManager
{
	struct LoopInfo 
	{
		Loop*		loop;
		SavedIPInfo	ipInfo;		// stack depth, ip of loop start
		UInt32		endIP;		// ip of instruction following loop end
	};

	static LoopManager* s_singleton;
	std::stack<LoopInfo> m_loops;

	LoopManager()	{ delete s_singleton;	}

	void RestoreStack(ScriptExecutionState* state, SavedIPInfo* info);
public:
	static LoopManager* GetSingleton() { if (!s_singleton) s_singleton = new LoopManager; return s_singleton; }

	void Add(Loop* loop, ScriptExecutionState* state, UInt32 startOffset, UInt32 endOffset, COMMAND_ARGS);
	bool Break(ScriptExecutionState* state, COMMAND_ARGS);
	bool Continue(ScriptExecutionState* state, COMMAND_ARGS);
};

LoopManager* LoopManager::s_singleton = NULL;

// ***************************
// Loop class member functions
// ***************************

ArrayIterLoop::ArrayIterLoop(ForEachContext* context)
{
	m_srcID = context->sourceID;
	m_iterID = 0;
	g_ArrayMap.AddReference(&m_iterID, context->iteratorID, 0xFF);

	ArrayElement elem;
	ArrayKey key;

	if (g_ArrayMap.GetFirstElement(m_srcID, &elem, &key))
	{
		m_curKey = key;
		UpdateIterator(&elem);		// initialize iterator to first element in array
	}
}

void ArrayIterLoop::UpdateIterator(const ArrayElement* elem)
{
	std::string val("value");
	std::string key("key");

	// iter["value"] = element data
	switch (elem->DataType())
	{
	case kDataType_String:
		g_ArrayMap.SetElementString(m_iterID, val, elem->m_data.str);
		break;
	case kDataType_Numeric:
		g_ArrayMap.SetElementNumber(m_iterID, val, elem->m_data.num);
		break;
	case kDataType_Form:
		{
			g_ArrayMap.SetElementFormID(m_iterID, val, elem->m_data.formID);
			break;
		}
	case kDataType_Array:
		{
			ArrayID arrID = elem->m_data.num;
			g_ArrayMap.SetElementArray(m_iterID, val, arrID);
			break;
		}
	default:
		DEBUG_PRINT("ArrayIterLoop::UpdateIterator(): unknown datatype %d found for element value", elem->DataType());
	}

	// iter["key"] = element key
	switch (m_curKey.KeyType())
	{
	case kDataType_String:
		g_ArrayMap.SetElementString(m_iterID, key, m_curKey.Key().str);
		break;
	default:
		g_ArrayMap.SetElementNumber(m_iterID, key, m_curKey.Key().num);
	}
}

bool ArrayIterLoop::Update(COMMAND_ARGS)
{
	ArrayElement elem;
	ArrayKey key;

	if (g_ArrayMap.GetNextElement(m_srcID, &m_curKey, &elem, &key))
	{
		m_curKey = key;
		UpdateIterator(&elem);	
		return true;
	}

	return false;
}

ArrayIterLoop::~ArrayIterLoop()
{
	g_ArrayMap.RemoveReference(&m_iterID, 0xFF);
}

StringIterLoop::StringIterLoop(ForEachContext* context)
{
	StringVar* srcVar = g_StringMap.Get(context->sourceID);
	StringVar* iterVar = g_StringMap.Get(context->iteratorID);
	if (srcVar && iterVar)
	{
		m_src = srcVar->String();
		m_curIndex = 0;
		m_iterID = context->iteratorID;
		if (m_src.length())
			iterVar->Set(m_src.substr(0, 1).c_str());
	}
}

bool StringIterLoop::Update(COMMAND_ARGS)
{
	StringVar* iterVar = g_StringMap.Get(m_iterID);
	if (iterVar)
	{
		m_curIndex++;
		if (m_curIndex < m_src.length())
		{
			iterVar->Set(m_src.substr(m_curIndex, 1).c_str());
			return true;
		}
	}

	return false;
}

void LoopManager::Add(Loop* loop, ScriptExecutionState* state, UInt32 startOffset, UInt32 endOffset, COMMAND_ARGS)
{
	// save the instruction offsets
	LoopInfo loopInfo;
	loopInfo.loop = loop;
	loopInfo.endIP = endOffset;

	// save the stack
	SavedIPInfo* savedInfo = &loopInfo.ipInfo;
	savedInfo->ip = startOffset;
	savedInfo->stackDepth = state->stackDepth;
	memcpy(savedInfo->stack, state->stack, (savedInfo->stackDepth + 1) * sizeof(UInt32));

	// add the loop
	m_loops.push(loopInfo);
}

void LoopManager::RestoreStack(ScriptExecutionState* state, SavedIPInfo* info)
{
	state->stackDepth = info->stackDepth;
	memcpy(state->stack, info->stack, (info->stackDepth + 1) * sizeof(UInt32));
}

bool LoopManager::Break(ScriptExecutionState* state, COMMAND_ARGS)
{
	if (!m_loops.size())
		return false;

	LoopInfo* loopInfo = &m_loops.top();

	RestoreStack(state, &loopInfo->ipInfo);
	opcodeOffsetPtr[kDataDeltaStackOffset] += loopInfo->endIP - (*opcodeOffsetPtr);

	delete loopInfo->loop;
	m_loops.pop();

	return true;
}

bool LoopManager::Continue(ScriptExecutionState* state, COMMAND_ARGS)
{
	if (!m_loops.size())
		return false;

	LoopInfo* loopInfo = &m_loops.top();

	if (!loopInfo->loop->Update(PASS_COMMAND_ARGS))
	{
		Break(state, PASS_COMMAND_ARGS);
		return true;
	}

	RestoreStack(state, &loopInfo->ipInfo);
	opcodeOffsetPtr[kDataDeltaStackOffset] += loopInfo->ipInfo.ip - (*opcodeOffsetPtr);

	return true;
}


bool WhileLoop::Update(COMMAND_ARGS)
{
	// save *opcodeOffsetPtr so we can calc IP to branch to after evaluating loop condition
	UInt32 originalOffset = *opcodeOffsetPtr;

	// update offset to point to loop condition, evaluate
	*opcodeOffsetPtr = m_exprOffset;
	ExpressionEvaluator eval(PASS_COMMAND_ARGS);
	bool bResult = eval.ExtractArgs();

	*opcodeOffsetPtr = originalOffset;

	if (bResult && eval.Arg(0))
	{
		ScriptToken boolToken = eval.Arg(0)->ConvertTo(kTokenType_Boolean, &eval);
		if (boolToken.IsGood())
			bResult = boolToken.value.num ? true : false;
	}

	return bResult;
}

// *********** commands

static bool Cmd_Let_Execute(COMMAND_ARGS)
{
	ExpressionEvaluator evaluator(PASS_COMMAND_ARGS);
	evaluator.ExtractArgs();

	return true;
}

// used to evaluate OBSE expressions within 'if' statements
// i.e. if eval (array[idx] == someThing)
static bool Cmd_eval_Execute(COMMAND_ARGS)
{
	*result = 0;
	ExpressionEvaluator eval(PASS_COMMAND_ARGS);
	
	if (eval.ExtractArgs() && eval.Arg(0))
	{
		ScriptToken boolToken = eval.Arg(0)->ConvertTo(kTokenType_Boolean, &eval);
		if (boolToken.IsGood())
			*result = boolToken.value.num;
	}

	return true;
}

// attempts to evaluate an expression. Returns false if error occurs, true otherwise. Suppresses error messages
static bool Cmd_testexpr_Execute(COMMAND_ARGS)
{
	*result = 0;
	ExpressionEvaluator eval(PASS_COMMAND_ARGS);
	ExpressionEvaluator::ToggleErrorSuppression(true);

	if (eval.ExtractArgs() && eval.Arg(0) && eval.Arg(0)->IsGood())
	{
		if (eval.Arg(0)->type == kTokenType_ArrayElement)		// is it an array elem with valid index?
			*result = eval.Arg(0)->ConvertTo(kTokenType_Boolean, &eval).IsGood() ? 1 : 0;
		else
			*result = 1;
	}

	ExpressionEvaluator::ToggleErrorSuppression(false);
	return true;
}

static bool Cmd_While_Execute(COMMAND_ARGS)
{
	ScriptExecutionState* state;
	GET_EXECUTION_STATE(state);

	// read offset to end of loop
	UInt8* data = (UInt8*)arg1 + *opcodeOffsetPtr;
	UInt32 offsetToEnd = *(UInt32*)data;

	// calc offset of first instruction following this While expression
	data += 5;		// UInt32 offset + UInt8 numArgs
	UInt16 exprLen = *((UInt16*)data);
	UInt32 startOffset = *opcodeOffsetPtr + 5 + exprLen;

	// create the loop and add it to loop manager
	WhileLoop* loop = new WhileLoop(*opcodeOffsetPtr + 4);
	LoopManager* mgr = LoopManager::GetSingleton();
	mgr->Add(loop, state, startOffset, offsetToEnd, PASS_COMMAND_ARGS);

	// test condition, break immediately if false
	*opcodeOffsetPtr = startOffset;		// need to update to point to _next_ instruction before calling loop->Update()
	if (!loop->Update(PASS_COMMAND_ARGS))
		mgr->Break(state, PASS_COMMAND_ARGS);

	return true;
}

static bool Cmd_ForEach_Execute(COMMAND_ARGS)
{
	ScriptExecutionState* state;
	GET_EXECUTION_STATE(state);

	// get offset to end of loop
	UInt8* data = (UInt8*)arg1 + *opcodeOffsetPtr;
	UInt32 offsetToEnd = *(UInt32*)data;

	// calc offset to first instruction within loop
	data += 5;
	UInt16 exprLen = *((UInt16*)data);
	UInt32 startOffset = *opcodeOffsetPtr + 5 + exprLen;

	// evaluate the expression to get the context
	*opcodeOffsetPtr += 4;				// set to start of expression
	ExpressionEvaluator eval(PASS_COMMAND_ARGS);
	bool bExtracted = eval.ExtractArgs();
	*opcodeOffsetPtr -= 4;				// restore

	if (!bExtracted || !eval.Arg(0))
	{
		ShowRuntimeError(scriptObj, "ForEach expression failed to return a value");
		return false;
	}

	ForEachLoop* loop = NULL;
	ForEachContext* context = eval.Arg(0)->GetForEachContext();
	if (!context)
		ShowRuntimeError(scriptObj, "Cmd_ForEach: Expression does not evaluate to a ForEach context");
	else		// construct the loop 
	{
		if (context->variableType == Script::eVarType_Array)
		{
			ArrayIterLoop* arrayLoop = new ArrayIterLoop(context);
			loop = arrayLoop;
		}
		else if (context->variableType == Script::eVarType_String)
		{
			StringIterLoop* stringLoop = new StringIterLoop(context);
			loop = stringLoop;
		}
	}

	if (loop)
	{
		LoopManager* mgr = LoopManager::GetSingleton();
		mgr->Add(loop, state, startOffset, offsetToEnd, PASS_COMMAND_ARGS);
		if (loop->IsEmpty())
		{
			*opcodeOffsetPtr = startOffset;
			mgr->Break(state, PASS_COMMAND_ARGS);
		}
	}

	return true;
}

static bool Cmd_Break_Execute(COMMAND_ARGS)
{
	ScriptExecutionState* state;
	GET_EXECUTION_STATE(state);

	LoopManager* mgr = LoopManager::GetSingleton();
	if (mgr->Break(state, PASS_COMMAND_ARGS))
		return true;

	ShowRuntimeError(scriptObj, "Break called outside of a valid loop context.");
	return false;
}

static bool Cmd_Continue_Execute(COMMAND_ARGS)
{
	ScriptExecutionState* state;
	GET_EXECUTION_STATE(state);

	LoopManager* mgr = LoopManager::GetSingleton();
	if (mgr->Continue(state, PASS_COMMAND_ARGS))
	{
		*opcodeOffsetPtr += 4;
		return true;
	}

	ShowRuntimeError(scriptObj, "Continue called outside of a valid loop context.");
	return false;
}

static bool Cmd_Loop_Execute(COMMAND_ARGS)
{
	ScriptExecutionState* state;
	GET_EXECUTION_STATE(state);

	LoopManager* mgr = LoopManager::GetSingleton();
	if (mgr->Continue(state, PASS_COMMAND_ARGS))
		return true;

	ShowRuntimeError(scriptObj, "Loop called outside of a valid loop context.");
	return false;
}

static bool Cmd_ToString_Execute(COMMAND_ARGS)
{
	*result = 0;
	std::string tokenAsString = "NULL";

	ExpressionEvaluator eval(PASS_COMMAND_ARGS);
	if (eval.ExtractArgs() && eval.Arg(0))
	{
		ScriptToken token;


		token = eval.Arg(0)->ConvertTo(kTokenType_String, &eval);
		if (token.IsGood())
			tokenAsString = token.value.str;
		else
		{
			token = eval.Arg(0)->ConvertTo(kTokenType_Number, &eval);
			if (token.IsGood())
			{
				char buf[0x20];
				sprintf_s(buf, sizeof(buf), "%g", token.value.num);
				tokenAsString = buf;
			}
			else
			{
				token = eval.Arg(0)->ConvertTo(kTokenType_Form, &eval);
				if (token.IsGood())
					tokenAsString = GetFullName(token.GetForm());
			}
		}
	}

	AssignToStringVar(PASS_COMMAND_ARGS, tokenAsString.c_str());
	return true;
}

static bool Cmd_Print_Execute(COMMAND_ARGS)
{
	*result = 0;
	ExpressionEvaluator eval(PASS_COMMAND_ARGS);
	if (eval.ExtractArgs() && eval.Arg(0))
	{
		ScriptToken strResult = eval.Arg(0)->ConvertTo(kTokenType_String, &eval);
		if (strResult.IsGood())
			Console_Print(strResult.value.str.c_str());
	}

	return true;
}

static bool Cmd_TypeOf_Execute(COMMAND_ARGS)
{
	std::string typeStr = "NULL";

	ExpressionEvaluator eval(PASS_COMMAND_ARGS);
	if (eval.ExtractArgs() && eval.Arg(0))
	{
		if (eval.Arg(0)->ConvertTo(kTokenType_Number, &eval).IsGood())
			typeStr = "Number";
		else if (eval.Arg(0)->ConvertTo(kTokenType_String, &eval).IsGood())
			typeStr = "String";
		else if (eval.Arg(0)->ConvertTo(kTokenType_Form, &eval).IsGood())
			typeStr = "Form";
		else
		{
			ScriptToken arrToken = eval.Arg(0)->ConvertTo(kTokenType_Array, &eval);
			if (arrToken.IsGood())
				typeStr = g_ArrayMap.GetTypeString(arrToken.value.arrID);
		}
	}

	AssignToStringVar(PASS_COMMAND_ARGS, typeStr.c_str());
	return true;
}

#endif

static bool Cmd_Let_Parse(UInt32 numParams, ParamInfo* paramInfo, ScriptLineBuffer* lineBuf, ScriptBuffer* scriptBuf)
{
#if OBLIVION
	Console_Print("Let cannot be called from the console.");
	return false;
#endif

	ExpressionParser parser(scriptBuf, lineBuf);
	if (!parser.ParseArgs(paramInfo, numParams))
		return false;

	// verify that assignment operator is last data recorded
	UInt8 lastData = lineBuf->dataBuf[lineBuf->dataOffset - 1];
	if (lastData != kOpType_Assignment)
	{
		ShowCompilerError(lineBuf, "Expected assignment in Let statement.\n\nCompiled script not saved.");
		return false;
	}

	return true;
}

static bool Cmd_BeginLoop_Parse(UInt32 numParams, ParamInfo* paramInfo, ScriptLineBuffer* lineBuf, ScriptBuffer* scriptBuf)
{
	// reserve space for a UInt32 storing offset past end of loop
	RegisterLoopStart(scriptBuf->scriptData + scriptBuf->dataOffset + lineBuf->dataOffset + 4);	
	lineBuf->dataOffset += sizeof(UInt32);

	// parse the loop condition
	ExpressionParser parser(scriptBuf, lineBuf);
	return parser.ParseArgs(paramInfo, numParams);
}

static bool Cmd_Loop_Parse(UInt32 numParams, ParamInfo* paramInfo, ScriptLineBuffer* lineBuf, ScriptBuffer* scriptBuf)
{
	UInt8* endPtr = scriptBuf->scriptData + scriptBuf->dataOffset + lineBuf->dataOffset + 4;
	UInt32 offset = endPtr - scriptBuf->scriptData;		// num bytes between beginning of script and instruction following Loop

	if (!HandleLoopEnd(offset))
	{
		ShowCompilerError(lineBuf, "'Loop' encountered without matching 'While' or 'ForEach'.\n\nCompiled script not saved.");
		return false;
	}

	return true;	
}

static bool Cmd_Null_Parse(UInt32 numParams, ParamInfo* paramInfo, ScriptLineBuffer* lineBuf, ScriptBuffer* scriptBuf)
{
	// takes no args, writes no data
	return true;
}

static ParamInfo kParams_OneBasicType[] =
{
	{	"expression",	kOBSEParamType_BasicType,	0	},
};

CommandInfo kCommandInfo_Let =
{
	"Let",
	"",
	0,
	"assigns the value of an expression to a variable",
	0,
	1,
	kParams_OneBasicType,
	HANDLER(Cmd_Let_Execute),
	Cmd_Let_Parse,
	NULL,
	0
};

static ParamInfo kParams_OneBoolean[] =
{
	{	"boolean expression",	kOBSEParamType_Boolean,	0	},
};

CommandInfo kCommandInfo_eval =
{
	"eval",
	"",
	0,
	"evaluates an expression and returns a boolean result.",
	0,
	1,
	kParams_OneBoolean,
	HANDLER(Cmd_eval_Execute),
	Cmd_Expression_Parse,
	NULL,
	0
};

static ParamInfo kParams_NoTypeChecking[] =
{
	{	"expression",	kOBSEParamType_NoTypeCheck,	0	},
};

CommandInfo kCommandInfo_testexpr =
{
	"testexpr",
	"",
	0,
	"returns false if errors occur while evaluating expression, true otherwise",
	0,
	1,
	kParams_NoTypeChecking,
	HANDLER(Cmd_testexpr_Execute),
	Cmd_Expression_Parse,
	NULL,
	0
};

CommandInfo kCommandInfo_While =
{
	"while",
	"",
	0,
	"loops until the given condition evaluates to false",
	0,
	1,
	kParams_OneBoolean,
	HANDLER(Cmd_While_Execute),
	Cmd_BeginLoop_Parse,
	NULL,
	0
};

CommandInfo kCommandInfo_Loop =
{
	"loop",
	"",
	0,
	"marks the end of a While or ForEach loop",
	0,
	1,
	kParams_OneOptionalInt,				// unused, but need at least one param for Parse() to be called
	HANDLER(Cmd_Loop_Execute),
	Cmd_Loop_Parse,
	NULL,
	0
};

CommandInfo kCommandInfo_Continue =
{
	"continue",
	"",
	0,
	"returns control to the top of a loop",
	0,
	1,
	kParams_OneOptionalInt,				// unused, but need at least one param for Parse() to be called
	HANDLER(Cmd_Continue_Execute),
	Cmd_Null_Parse,
	NULL,
	0
};

CommandInfo kCommandInfo_Break =
{
	"break",
	"",
	0,
	"exits a loop",
	0,
	1,
	kParams_OneOptionalInt,				// unused, but need at least one param for Parse() to be called
	HANDLER(Cmd_Break_Execute),
	Cmd_Null_Parse,
	NULL,
	0
};

static ParamInfo kParams_ForEach[] =
{
	{	"ForEach expression",	kOBSEParamType_ForEachContext,	0	},
};

CommandInfo kCommandInfo_ForEach =
{
	"ForEach",
	"",
	0,
	"iterates over the elements of an array",
	0,
	1,
	kParams_ForEach,
	HANDLER(Cmd_ForEach_Execute),
	Cmd_BeginLoop_Parse,
	NULL,
	0
};

static ParamInfo kParams_OneOBSEString[] =
{
	{	"string",	kOBSEParamType_String,	0	},
};

CommandInfo kCommandInfo_ToString = 
{
	"ToString",
	"",
	0,
	"attempts to convert an expression to a string",
	0,
	1,
	kParams_OneBasicType,
	HANDLER(Cmd_ToString_Execute),
	Cmd_Expression_Parse,
	NULL,
	0
};

CommandInfo kCommandInfo_TypeOf =
{
	"TypeOf",
	"",
	0,
	"returns a string representing the type of the expression",
	0,
	1,
	kParams_OneBasicType,
	HANDLER(Cmd_TypeOf_Execute),
	Cmd_Expression_Parse,
	NULL,
	0
};

CommandInfo kCommandInfo_Print =
{
	"Print",
	"",
	0,
	"prints a string expression to the console",
	0,
	1,
	kParams_OneOBSEString,
	HANDLER(Cmd_Print_Execute),
	Cmd_Expression_Parse,
	NULL,
	0
};
