// Copyright (C) 2012 Zeex
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
// of the Software, and to permit persons to whom the Software is furnished to do
// so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#if defined DISASM_INC
	#endinput
#endif
#define DISASM_INC

#include <string>

#include "amx_base"
#include "amx_header"
#include "amx_memory"
#include "opcode"

#define MAX_INSN_NAME 15

enum DisasmInsnInfo {
	DisasmInsnInfo_name[MAX_INSN_NAME],
	DisasmInsnInfo_num_opers,
	bool:DisasmInsnInfo_needs_reloc
}

enum DisasmContext {
	DisasmContext_start_ip,
	DisasmContext_end_ip,
	DisasmContext_nip,
	DisasmContext_cip,
	Opcode:DisasmContext_opcode,
	DisasmContext_insn[DisasmInsnInfo]
}

static stock const insn_table[][DisasmInsnInfo] = {
	{ "none",        0, false },
	{ "load.pri",    1, false },
	{ "load.alt",    1, false },
	{ "load.s.pri",  1, false },
	{ "load.s.alt",  1, false },
	{ "lref.pri",    1, false },
	{ "lref.alt",    1, false },
	{ "lref.s.pri",  1, false },
	{ "lref.s.alt",  1, false },
	{ "load.i",      0, false },
	{ "lodb.i",      1, false },
	{ "const.pri",   1, false },
	{ "const.alt",   1, false },
	{ "addr.pri",    1, false },
	{ "addr.alt",    1, false },
	{ "stor.pri",    1, false },
	{ "stor.alt",    1, false },
	{ "stor.s.pri",  1, false },
	{ "stor.s.alt",  1, false },
	{ "sref.pri",    1, false },
	{ "sref.alt",    1, false },
	{ "sref.s.pri",  1, false },
	{ "sref.s.alt",  1, false },
	{ "stor.i",      0, false },
	{ "strb.i",      1, false },
	{ "lidx",        0, false },
	{ "lidx.b",      1, false },
	{ "idxaddr",     0, false },
	{ "idxaddr.b",   1, false },
	{ "align.pri",   1, false },
	{ "align.alt",   1, false },
	{ "lctrl",       1, false },
	{ "sctrl",       1, false },
	{ "move.pri",    0, false },
	{ "move.alt",    0, false },
	{ "xchg",        0, false },
	{ "push.pri",    0, false },
	{ "push.alt",    0, false },
	{ "push.r",      1, false },
	{ "push.c",      1, false },
	{ "push",        1, false },
	{ "push.s",      1, false },
	{ "pop.pri",     0, false },
	{ "pop.alt",     0, false },
	{ "stack",       1, false },
	{ "heap",        1, false },
	{ "proc",        0, false },
	{ "ret",         0, false },
	{ "retn",        0, false },
	{ "call",        1, true  },
	{ "call.pri",    0, false },
	{ "jump",        1, true  },
	{ "jrel",        1, true  },
	{ "jzer",        1, true  },
	{ "jnz",         1, true  },
	{ "jeq",         1, true  },
	{ "jneq",        1, true  },
	{ "jless",       1, true  },
	{ "jleq",        1, true  },
	{ "jgrtr",       1, true  },
	{ "jgeq",        1, true  },
	{ "jsless",      1, true  },
	{ "jsleq",       1, true  },
	{ "jsgrtr",      1, true  },
	{ "jsgeq",       1, true  },
	{ "shl",         0, false },
	{ "shr",         0, false },
	{ "sshr",        0, false },
	{ "shl.c.pri",   1, false },
	{ "shl.c.alt",   1, false },
	{ "shr.c.pri",   1, false },
	{ "shr.c.alt",   1, false },
	{ "smul",        0, false },
	{ "sdiv",        0, false },
	{ "sdiv.alt",    0, false },
	{ "umul",        0, false },
	{ "udiv",        0, false },
	{ "udiv.alt",    0, false },
	{ "add",         0, false },
	{ "sub",         0, false },
	{ "sub.alt",     0, false },
	{ "and",         0, false },
	{ "or",          0, false },
	{ "xort",        0, false },
	{ "not",         0, false },
	{ "neg",         0, false },
	{ "invert",      0, false },
	{ "add.c",       1, false },
	{ "smul.c",      1, false },
	{ "zero.pri",    0, false },
	{ "zero.alt",    0, false },
	{ "zero",        1, false },
	{ "zero.s",      1, false },
	{ "sign.pri",    0, false },
	{ "sign.alt",    0, false },
	{ "eq",          0, false },
	{ "neq",         0, false },
	{ "less",        0, false },
	{ "leq",         0, false },
	{ "grtr",        0, false },
	{ "geq",         0, false },
	{ "sless",       0, false },
	{ "sleq",        0, false },
	{ "sgrtr",       0, false },
	{ "sgeq",        0, false },
	{ "eq.c.pri",    1, false },
	{ "eq.c.alt",    1, false },
	{ "inc.pri",     0, false },
	{ "inc.alt",     0, false },
	{ "inc",         1, false },
	{ "inc.s",       1, false },
	{ "inc.i",       0, false },
	{ "dec.pri",     0, false },
	{ "dec.alt",     0, false },
	{ "dec",         1, false },
	{ "dec.s",       1, false },
	{ "dec.i",       0, false },
	{ "movs",        1, false },
	{ "cmps",        1, false },
	{ "fill",        1, false },
	{ "halt",        1, false },
	{ "bounds",      1, false },
	{ "sysreq.pri",  1, false },
	{ "sysreq.c",    1, false },
	{ "file",       -1, false }, // obsolete
	{ "line",       -1, false }, // obsolete
	{ "symbol",     -1, false }, // obsolete
	{ "srange",     -1, false }, // obsolete
	{ "jump.pri",    0, false },
	{ "switch",      1, false },
	{ "casetbl",    -1, true },
	{ "swap.pri",    0, false },
	{ "swap.alt",    0, false },
	{ "push.adr",    1, false },
	{ "nop",         0, false },
	{ "sysreq.d",    1, false },
	{ "symtag",      1, false }, // obsolete
	{ "break",       0, false }
};

stock DisasmInit(ctx[DisasmContext], start = 0, end = 0) {
	new hdr[AMX_HDR];
	GetAmxHeader(hdr);

	new dat = hdr[AMX_HDR_DAT];
	new cod = hdr[AMX_HDR_COD];

	new code_base = cod - dat;

	start += code_base;

	ctx[DisasmContext_nip] = start;
	ctx[DisasmContext_cip] = start;

	ctx[DisasmContext_start_ip] = start;
	if (end != 0) {
		ctx[DisasmContext_end_ip] = code_base + end;
	} else {
		ctx[DisasmContext_end_ip] = code_base + (dat - cod);
	}
}

stock bool:DisasmDecodeInsn(ctx[DisasmContext]) {
	new ip = ctx[DisasmContext_nip];
	if (ip >= 0) {
		return false;
	}

	new Opcode:opcode = UnrelocateOpcode(Opcode:ReadAmxMemory(ctx[DisasmContext_nip]));
	if (opcode <= OP_NONE || _:opcode >= NUM_OPCODES) {
		return false;
	}

	ctx[DisasmContext_cip] = ctx[DisasmContext_nip];
	ctx[DisasmContext_opcode] = opcode;
	ctx[DisasmContext_insn] = insn_table[_:opcode];

	ip += 4;

	if (opcode == OP_CASETBL) {
		new n = ReadAmxMemory(ip);
		ip += 4;
		ip += (2 * n + 1) * 4;
		ctx[DisasmContext_insn][DisasmInsnInfo_num_opers] = n;
	} else {
		ip += 4 * insn_table[_:opcode][DisasmInsnInfo_num_opers];
	}

	ctx[DisasmContext_nip] = ip;
	return true;
}

stock Opcode:DisasmGetOpcode(ctx[DisasmContext]) {
	return ctx[DisasmContext_opcode];
}

stock DisasmGetOperand(ctx[DisasmContext], index = 0) {
	return ReadAmxMemory(ctx[DisasmContext_cip] + (index + 1) * 4);
}

stock DisasmGetNumOperands(ctx[DisasmContext]) {
	return ctx[DisasmContext_insn][DisasmInsnInfo_num_opers];
}

stock bool:DisasmNeedReloc(ctx[DisasmContext]) {
	return ctx[DisasmContext_insn][DisasmInsnInfo_needs_reloc];
}

stock DisasmGetNextIp(ctx[DisasmContext]) {
	return ctx[DisasmContext_nip];
}

stock DisasmGetCurIp(ctx[DisasmContext]) {
	return ctx[DisasmContext_cip];
}

stock DisasmGetInsnName(ctx[DisasmContext], name[], size = sizeof(name)) {
	name[0] = '\0';
	strcat(name, ctx[DisasmContext_insn][DisasmInsnInfo_name], size);
}

static stock ToHexStr(x) {
	new s[11];
	new i = 0;
	new j = 0;

	while (i < sizeof(s) && j < 8) {
		new n = x >> (7 - j) * 4 & 0xF;
		switch (n) {
			case 0x0..0x9:
				s[i] = n + '0';
			case 0xA..0xF:
				s[i] = n + 'a' - 0xA;
		}
		i++;
		j++;
	}

	return s;
}

static stock bool:IsPrintableAscii(c) {
	return 32 <= c <= 126;
}

static stock ToPrintableAscii(c) {
	if (IsPrintableAscii(c)) {
		return c;
	}
	return ' ';
}

static stock DumpDataRowChar(File:file, start, num, max) {
	new cur = start;
	new end = start + num*4;

	while (cur < max) {
		new p[4 char + 1];
		p[0] = ReadAmxMemory(cur);

		new u[4 + 1];
		u[0] = ToPrintableAscii(p{0});
		u[1] = ToPrintableAscii(p{1});
		u[2] = ToPrintableAscii(p{2});
		u[3] = ToPrintableAscii(p{3});
		u[4] = '\0';

		if (cur < end) {
			fwrite(file, u);
		} else {
			fwrite(file, " ");
		}
		cur += 4;
	}
}

static stock DumpDataRowHex(File:file, start, num, max) {
	new cur = start;
	new end = start + num*4;

	while (cur < max) {
		if (cur < end) {
			fwrite(file, ToHexStr(ReadAmxMemory(cur)));
		} else {
			fwrite(file, "        ");
		}
		fwrite(file, " ");
		cur += 4;
	}
}

stock DisasmDump(const out_filename[]) {
	new ctx[DisasmContext];
	DisasmInit(ctx);

	new File:file = fopen(out_filename, io_write);
	if (!file) {
		return;
	}

	new base = GetAmxBaseAddress();

	new hdr[AMX_HDR];
	GetAmxHeader(hdr);

	new dat = hdr[AMX_HDR_DAT];
	new cod = hdr[AMX_HDR_COD];
	new hea = hdr[AMX_HDR_HEA];

	// DISASSAMBLE CODE
	fwrite(file, "; CODE\n\n");

	while (DisasmDecodeInsn(ctx)) {
		new Opcode:opcode = DisasmGetOpcode(ctx);

		if (opcode == OP_PROC) {
			fwrite(file, "\n");
		}

		new cip = DisasmGetCurIp(ctx);

		new insn_name[15];
		DisasmGetInsnName(ctx, insn_name);
		fwrite(file, ToHexStr(cip + dat - cod));
		fwrite(file, " ");
		fwrite(file, insn_name);
		fwrite(file, " ");

		if (opcode == OP_PROC) {
			new name[32];
			new address = cip + dat - cod;
			if (address == hdr[AMX_HDR_CIP]) {
				strcat(name, "main");
			} else {
				new index = GetPublicIndexFromAddress(address);
				if (index >= 0) {
					GetPublicNameFromIndex(index, name);
				}
			}
			if (strlen(name) != 0) {
				fwrite(file, "; ");
				fwrite(file, name);
			}
		}

		if (opcode == OP_CASETBL) {
			new num = DisasmGetNumOperands(ctx);
			for (new i = 0; i < num; i++) {
				fwrite(file, "\n         ");
				new addr = ReadAmxMemory(cip + 4 * i * 2);
				new rel_addr = addr - (base + cod) + (cip - 4 + dat - cod);
				fwrite(file, ToHexStr(rel_addr));
				fwrite(file, " ");
				new val = ReadAmxMemory(cip + 4 * i * 2 + 4);
				fwrite(file, ToHexStr(val));
			}
			goto next;
		}

		switch (opcode) {
			case OP_CALL: {
				new name[32];
				new address = ReadAmxMemory(cip) - base - cod;
				if (address == hdr[AMX_HDR_CIP]) {
					strcat(name, "main");
				} else {
					new index = GetPublicIndexFromAddress(address);
					if (index >= 0) {
						GetPublicNameFromIndex(index, name);
					}
				}
				if (strlen(name) > 0) {
					fwrite(file, name);
					goto next;
				}
			}
			case OP_SYSREQ_C, OP_SYSREQ_D: {
				new name[100];
				if (opcode == OP_SYSREQ_C) {
					new index = DisasmGetOperand(ctx);
					GetNativeNameFromIndex(index, name);
				} else {
					new address = DisasmGetOperand(ctx);
					new index = GetNativeIndexFromAddress(address);
					if (index >= 0) {
						GetNativeNameFromIndex(index, name);
					}
				}
				if (strlen(name) > 0) {
					fwrite(file, name);
					goto next;
				}
			}
		}

		new n = DisasmGetNumOperands(ctx);
		for (new i = 0; i < n; i++) {
			new operand = DisasmGetOperand(ctx, i);
			if (DisasmNeedReloc(ctx)) {
				operand -= base + cod;
			}
			fwrite(file, ToHexStr(operand));
		}

next:
		fwrite(file, "\n");
	}

	// DUMP DATA
	fwrite(file, "\n\n; DATA\n");

	new data_end = hea - dat;

	for (new i = 0; i < data_end; i += 0x10) {
		fwrite(file, ToHexStr(i));
		fwrite(file, "  ");
		DumpDataRowHex(file, i, 4, min(i + 0x10, data_end));
		fwrite(file, " ");
		DumpDataRowChar(file, i, 4, min(i + 0x10, data_end));
		fwrite(file, "\n");
	}

	fclose(file);
}
