#include <assert.h>

#include <string.h>

#include <stdlib.h>

#include "decompiler.h"

#include "code_analyzer.h"

#include "platform_util.h"

struct Block;

#define PCM_BRANCH_TARGET		(1<<0)

#define PCM_OPCODE				(1<<1)

#define PCM_HAS_BLOCK			(1<<2)

#define PCM_NUM_BITS			3

#define NEXT_OP(cip, op)		cip += (dc->opdef[op].params + 1) * sizeof(cell_t)

struct Edge

{

	Edge *next;

	Block *target;

	uint32_t ip;

};

struct Block

{

	cell_t ip_start;

	cell_t ip_end;

	Edge *edge_list;

	Edge *final_edge;

	Block *next_in_table;

};

struct ControlFlowGraph

{

	int err;

	unsigned int max_blocks;

	unsigned int max_edges;

	unsigned int used_blocks;

	unsigned int used_edges;

	Block *blocks;

	Edge *edges;

	Block *entry;

	FunctionInfo *func;

};

DefineNode::DefineNode(BaseNode *o, const char *fmt, ...) : BaseNode(_OP_DEFINE), other(o)

{

	va_list ap;

	va_start(ap, fmt);

	Sp_FormatArgs(name, sizeof(name), fmt, ap);

	va_end(ap);

}

static FunctionInfo *sp_InferFuncPrototype(sp_decomp_t *dc, uint32_t code_addr, bool is_public)

{

	sp_tag_t *pTag;

	sp_file_t *plugin;

	FunctionInfo *info;

	sp_fdbg_symbol_t *sym;

	plugin = dc->plugin;

	info = new FunctionInfo;

	memset(info, 0, sizeof(FunctionInfo));

	info->code_addr = code_addr;

	if ((sym = Sp_FindFunctionSym(plugin, code_addr)) != NULL)

	{

		Sp_Format(info->name, sizeof(info->name), "%s", plugin->debug.stringbase + sym->name);

		if (sym->tagid != 0)

		{

			info->tag = Sp_FindTag(plugin, sym->tagid);

		}

	}

	else

	{

		Sp_Format(info->name, sizeof(info->name), "function_%x", code_addr);

		for (uint32_t i = 0; i < plugin->num_publics; i++)

		{

			if (plugin->publics[i].code_offs == code_addr)

			{

				Sp_Format(info->name,

					sizeof(info->name),

					"%s",

					plugin->publics[i].name);

				break;

			}

		}

	}

	info->sym = sym;

	info->is_public = is_public;

	/* Search for all locals */

	{

		uint32_t i;

		uint8_t *cursor;

		cursor = (uint8_t *)plugin->debug.symbols;

		for (i = 0; i < plugin->debug.syms_num; i++)

		{

			sym = (sp_fdbg_symbol_t *)cursor;

			if (sym->ident == SP_SYM_VARIABLE

				|| sym->ident == SP_SYM_REFARRAY

				|| sym->ident == SP_SYM_REFERENCE

				|| sym->ident == SP_SYM_ARRAY)

			{

				if (code_addr >= sym->codestart

					&& code_addr < sym->codeend

					&& sym->addr >= 0xC

					&& (sym->addr & 0x3) == 0

					&& ((sym->addr - 0xC) >> 2) < SP_MAX_EXEC_PARAMS

					&& sym->vclass == 1)

				{

					unsigned int num;

					num = (sym->addr - 0xC) >> 2;

					info->known_args[num].sym = sym;

					if (num + 1 > info->num_known_args)

					{

						info->num_known_args = num + 1;

					}

				}

			}

			cursor += sizeof(sp_fdbg_symbol_t);

			if (sym->dimcount > 0)

			{

				cursor += sizeof(sp_fdbg_arraydim_t) * sym->dimcount;

			}

		}

	}

	/* Map all arguments into place. */

	for (unsigned int i = 0; i < info->num_known_args; i++)

	{

		if ((sym = info->known_args[i].sym) == NULL)

		{

			continue;

		}

		if (sym->tagid != 0 && (pTag = Sp_FindTag(plugin, sym->tagid)) != NULL)

		{

			info->known_args[i].tag = pTag;

		}

		info->known_args[i].name = plugin->debug.stringbase + sym->name;

		info->known_args[i].dimcount = sym->dimcount;

		if (sym->dimcount > 0)

		{

			info->known_args[i].dims = 

				(sp_fdbg_arraydim_t *)((char *)sym + sizeof(sp_fdbg_symbol_t));

		}

	}

	return info;

}

static FunctionInfo *sp_GetFuncPrototype(sp_decomp_t *dc,

										 uint32_t code_addr,

										 bool known,

										 bool is_public = false);

static FunctionInfo *sp_GetFuncPrototype(sp_decomp_t *dc,

										 uint32_t code_addr,

										 bool known,

										 bool is_public)

{

	FunctionInfo *func;

	func = dc->funcs;

	while (func != NULL)

	{

		if (func->code_addr == code_addr)

		{

			return func;

		}

		func = func->next;

	}

	/* See if this is public. */

	if (!known)

	{

		is_public = false;

		for (unsigned int i = 0; i < dc->plugin->num_publics; i++)

		{

			if (dc->plugin->publics[i].code_offs == code_addr)

			{

				is_public = true;

				break;

			}

		}

	}

	func = sp_InferFuncPrototype(dc, code_addr, is_public);

	func->next = dc->funcs;

	dc->funcs = func;

	return func;

}

static void sp_CacheLocals(sp_decomp_t *dc, FunctionInfo *info)

{

	sp_file_t *plugin;

	sp_fdbg_symbol_t *sym;

	plugin = dc->plugin;

	/* Search for all locals */

	{

		uint32_t i;

		uint8_t *cursor;

		cursor = (uint8_t *)plugin->debug.symbols;

		for (i = 0; i < plugin->debug.syms_num; i++)

		{

			sym = (sp_fdbg_symbol_t *)cursor;

			if ((sym->codestart >= info->code_addr && sym->codestart < info->code_end)

				&& (sym->codeend <= info->code_end && sym->codeend >= info->code_addr)

				&& sym->vclass == 1)

			{

				info->num_known_vars++;

			}

			cursor += sizeof(sp_fdbg_symbol_t);

			if (sym->dimcount > 0)

			{

				cursor += sizeof(sp_fdbg_arraydim_t) * sym->dimcount;

			}

		}

	}

	/* Map all locals into place. */

	if (info->num_known_vars)

	{

		uint32_t i, num;

		uint8_t *cursor;

		num = 0;

		info->known_vars = new FuncVar[info->num_known_vars];

		memset(info->known_vars, 0, sizeof(FuncVar) * info->num_known_vars);

		cursor = (uint8_t *)plugin->debug.symbols;

		for (i = 0; i < plugin->debug.syms_num; i++)

		{

			sym = (sp_fdbg_symbol_t *)cursor;

			if (sym->ident == SP_SYM_VARIABLE

				|| sym->ident == SP_SYM_REFARRAY

				|| sym->ident == SP_SYM_REFERENCE

				|| sym->ident == SP_SYM_ARRAY)

			{

				/* Detect all function variables. */

				if ((sym->codestart >= info->code_addr && sym->codestart < info->code_end)

					&& (sym->codeend <= info->code_end && sym->codeend >= info->code_addr)

					&& sym->vclass == 1)

				{

					assert(num < info->num_known_vars);

					info->known_vars[num].sym = sym;

					info->known_vars[num].name = plugin->debug.stringbase + sym->name;

					info->known_vars[num].tag = 

						(sym->tagid == 0) ? NULL : Sp_FindTag(plugin, sym->tagid);

					info->known_vars[num].dimcount = sym->dimcount;

					info->known_vars[num].dims = 

						(sym->dimcount == 0)

						? NULL

						: (sp_fdbg_arraydim_t *)(cursor + sizeof(sp_fdbg_symbol_t));

					num++;

				}

			}

			cursor += sizeof(sp_fdbg_symbol_t);

			if (sym->dimcount > 0)

			{

				cursor += sizeof(sp_fdbg_arraydim_t) * sym->dimcount;

			}

		}

		assert(num == info->num_known_vars);

	}

}

static void sp_PrintFunctionPrototype(PrintBuffer *printer, sp_decomp_t *dc, FunctionInfo *func)

{

	FuncVar *var;

	if (func->is_public)

	{

		printer->Append("public ");

	}

	if (func->tag != NULL)

	{

		printer->Append("%s:", func->tag->name);

	}

	printer->Append("%s(", func->name);

	for (unsigned int i = 0; i < func->num_known_args; i++)

	{

		var = &func->known_args[i];

		if (var->sym == NULL)

		{

			printer->Append("unknown_arg_%d%s, ", i, i == func->num_known_args - 1 ? "" : ",");

			continue;

		}

		if (var->sym->ident == SP_SYM_REFERENCE)

		{

			printer->Append("&");

		}

		if (var->tag != NULL)

		{

			printer->Append("%s:", var->tag->name);

		}

		printer->Append("%s", var->name);

		if (var->sym->dimcount > 0)

		{

			sp_fdbg_arraydim_t *dims = 

				(sp_fdbg_arraydim_t *)((char *)var->sym + sizeof(sp_fdbg_symbol_t));

			for (uint16_t j = 0; j < var->sym->dimcount; j++)

			{

				if (dims[j].size == 0)

				{

					printer->Append("[]");

				}

				else

				{

					printer->Append("[%d]", dims[j].size);

				}

			}

		}

		if (i != func->num_known_args - 1)

		{

			printer->Append(", ");

		}

	}

	printer->Append(")");

}

void sp_DebugExprTree(PrintBuffer *buffer, BaseNode *node, bool is_stmt=false)

{

	switch (node->op)

	{

	case _OP_DEFINE:

		{

			DefineNode *def = (DefineNode *)node;

			buffer->Append("new %s = ", def->name);

			sp_DebugExprTree(buffer, def->other, is_stmt);

			break;

		}

	case OP_CALL:

		{

			CallNode *call = (CallNode *)node;

			buffer->Append("%s(", call->func->name);

			for (unsigned int i = 0; i < call->argc; i++)

			{

				sp_DebugExprTree(buffer, call->args[i], false);

			}

			buffer->Append(")");

			break;

		}

	case _OP_USE:

		{

			UseNode *unode = (UseNode *)node;

			buffer->Append("%s", unode->def->name);

			break;

		}

	case _OP_VAR:

		{

			VarNode *var_node = (VarNode *)node;

			buffer->Append("%s", var_node->var->name);

			break;

		}

	case _OP_NEW:

		{

			DeclNode *d_node = (DeclNode *)node;

			if (d_node->expr == NULL)

			{

				buffer->Append("decl ");

				sp_DebugExprTree(buffer, d_node->var, true);

			}

			else

			{

				buffer->Append("new ");

				sp_DebugExprTree(buffer, d_node->var, true);

				buffer->Append(" = ");

				sp_DebugExprTree(buffer, d_node->expr, true);

			}

			break;

		}

	case _OP_STMT:

	case OP_RETN:

		{

			StmtNode *u_node = (StmtNode *)node;

			while (u_node != NULL)

			{

				buffer->Append("\t");

				if (u_node->op == OP_RETN)

				{

					buffer->Append("return ");

				}

				sp_DebugExprTree(buffer, u_node->node, true);

				buffer->Append(";\n");

				u_node = u_node->next;

			}

			break;

		}

	case _OP_CONST:

		{

			ConstNode *c_node = (ConstNode *)node;

			buffer->Append("%d", c_node->val);

			break;

		}

	case OP_ADD:

	case OP_ADD_C:

	case OP_SMUL:

	case OP_SMUL_C:

	case OP_SDIV:

	case OP_SUB_ALT:

	case _OP_STORE:

	case _OP_MODULO:

		{

			BinOpNode *bin_node = (BinOpNode *)node;

			if (!is_stmt)

			{

				buffer->Append("(");

			}

			sp_DebugExprTree(buffer, bin_node->left);

			switch (node->op)

			{

			case OP_ADD:

			case OP_ADD_C:

				{

					buffer->Append(" + ");

					break;

				}

			case OP_SMUL:

			case OP_SMUL_C:

				{

					buffer->Append(" * ");

					break;

				}

			case _OP_STORE:

				{

					buffer->Append(" = ");

					break;

				}

			case OP_SUB_ALT:

				{

					buffer->Append(" - ");

					break;

				}

			case OP_SDIV:

				{

					buffer->Append(" / ");

					break;

				}

			case _OP_MODULO:

				{

					buffer->Append(" % ");

					break;

				}

			}

			sp_DebugExprTree(buffer, bin_node->right, node->op == _OP_STORE);

			if (!is_stmt)

			{

				buffer->Append(")");

			}

			break;

		}

	}

}

void sp_DebugPrint(FunctionInfo *info, sp_decomp_t *dc, BaseNode *node)

{

	PrintBuffer buffer;

	sp_PrintFunctionPrototype(&buffer, dc, info);

	buffer.Append("\n{\n");

	sp_DebugExprTree(&buffer, node, true);

	buffer.Append("}\n");

	buffer.Dump(stdout);

}

#define PCODE_VAL(pcode, offs)			*(cell_t *)((char *)(pcode)+(offs))

#define PCODE_PARAM(pcode, offs, num)	*(cell_t *)((char *)(pcode)+(offs)+((num)*sizeof(cell_t)))

FuncVar *sp_FindGraphVar(FunctionInfo *func, ucell_t cip, cell_t sp, cell_t stack)

{

	if (stack >= 0xC)

	{

		if ((stack & 0x3) != 0)

		{

			return NULL;

		}

		stack -= 0xC;

		stack >>= 2;

		if (stack >= SP_MAX_EXEC_PARAMS)

		{

			return NULL;

		}

		if (func->known_args[stack].sym == NULL)

		{

			return NULL;

		}

		return &func->known_args[stack];

	}

	else if (stack < 0)

	{

		FuncVar *var;

		if (stack < sp)

		{

			return NULL;

		}

		for (unsigned int i = 0; i < func->num_known_vars; i++)

		{

			var = &func->known_vars[i];

			if (cip >= var->sym->codestart

				&& cip <= var->sym->codeend

				&& stack == var->sym->addr)

			{

				return var;

			}

		}

		return NULL;

	}

	else

	{

		return NULL;

	}

}

#define CHECK_VALUE(v)							\

	assert(v != NULL);							\

	if (v == NULL) {							\

		err = SP_ERROR_INVALID_INSTRUCTION;		\

		goto return_error;						\

	}