// todo: on weird jump: stop the disassembling there or treat as normal instruction?
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <stdbool.h>
#include <stddef.h>

#ifdef WINDOWS
#include <windows.h>
#else
#include <sys/mman.h>
#include <string.h>

#endif

#include "udis86.h"
#include "list.h"
#include "hook.h"
#include "misc.h"

#define MINIMUM_REQUIRED_FUNCTION_LENGTH_SHORT_HOOK 5
#define MINIMUM_REQUIRED_FUNCTION_LENGTH_LONG_HOOK 16
#define MAXIMUM_DISTANCE_FOR_SHORT_HOOK 0xFFFFFFFF
#define PAGE_BOUNDARY 0x1000

/**
	\brief Estimates the function size by checking opcodes until 
		* an instruction that signifies the end of the function or
		* MINIMUM_REQUIRED_FUNCTION_LENGTH_LONG_HOOK bytes were found
	\return 0 on error
*/
static size_t estimate_function_length(void* function);

/**
	\brief Either builds the trampoline
	\param functionSize How much of the function shall be copied
	\param trampolineSize out parameter
*/
static void* build_trampoline(void* function, size_t functionSize, size_t* trampolineSize);

/**
\brief jmp +X == +X + sizeof(jump type) + eip
*/
static size_t get_jump_offset(ud_t* ud, size_t offsetOfInstr);

/**
\todo add loop & jecx (and possibly more) - though both are probably only used to jump back
and currently is_jump is only called to from functions that don't care about that
*/
static bool is_jump(enum ud_mneomic_code op);

/**
\brief Writes
jmp [$+0] // $+0 == addr
addr: xxx
*/
static size_t write_x64_jump(unsigned char* where, void* toWhere);

static bool is_end_of_function(enum ud_mneomic_code opc, ud_t* ud, size_t furthestJump, size_t instrOffset);
static bool loops_into_overwritten_code(void* function);
static size_t write_jcc_jump(const uint8_t* instruction, void* whereToWrite, void* originalAddr);
static void fix_jcc_jumps(struct ListHead* instructions, struct ListHead* jccs);
static int32_t get_rip_delta(ud_t* ud);
static size_t write_instr_with_rip_delta(ud_t* ud, uint64_t absoluteAddr, const uint8_t* opcBuf, unsigned char* where);

int hook(void* function, size_t functionLength, void* replacement, void** trampoline)
{
	ptrdiff_t d = 0; // function minus replacement
	size_t needed = 0,
		trampolineSizeNeeded = 0;

	assert(function);
	assert(replacement);
	assert(trampoline);

	if(!functionLength)
		functionLength = estimate_function_length(function);

	/* is it actually possible to hook, space wise? */
	d = (ptrdiff_t)function - (ptrdiff_t)replacement;
	needed = d > MAXIMUM_DISTANCE_FOR_SHORT_HOOK ? 
		MINIMUM_REQUIRED_FUNCTION_LENGTH_LONG_HOOK : MINIMUM_REQUIRED_FUNCTION_LENGTH_SHORT_HOOK;
	if(functionLength < needed)
		return NOT_ENOUGH_SPACE;

	if(loops_into_overwritten_code(function))
		return LOOPS_INTO_OVERWRITTEN_CODE;

	//! fixme: Remove in real code
	needed = MINIMUM_REQUIRED_FUNCTION_LENGTH_LONG_HOOK; 

	printf("---\nNow build the trampoline\n");
	if((*trampoline = build_trampoline(function, needed, &trampolineSizeNeeded)) < 0)
		return (int)*trampoline;
	printf("Needed for trampoline %p: %d (Found %d)\n", *trampoline, trampolineSizeNeeded, functionLength);

#ifdef WINDOWS
	FlushInstructionCache(GetCurrentProcess(), function, needed);
#endif
	return SUCCESS;
}

static void* build_trampoline(void* function, size_t functionSize, size_t* trampolineSize)
{
	ud_t ud;
	void* trampoline;

	assert(trampolineSize);

	ud_init(&ud);
	ud_set_input_buffer(&ud, function, (size_t)function | (PAGE_BOUNDARY - 1)); // Can't read further than page boundary - there be dragons
	ud_set_pc(&ud, (uint64_t)function);
	ud_set_mode(&ud, 64);
	ud_set_syntax(&ud, UD_SYN_INTEL);

	*trampolineSize = 0;

	if(!(trampoline = alloc_rwx(PAGE_BOUNDARY)))
		return (void*)CANT_ALLOC;
	printf("trampoline @ %p\n", trampoline);

	size_t sz = 0;
	for(; sz <= functionSize;)
	{
		size_t instrLen = 0,
			offsetOfInstr = sz;
		bool neededFixup = false;
		uint8_t* addrForInstr = (uint8_t*)trampoline + *trampolineSize;
		const uint8_t* buf = NULL;

		assert(*trampolineSize < PAGE_BOUNDARY);

		if(!(instrLen = ud_disassemble(&ud)) || ud.error)
			return 0;
		buf = ud_insn_ptr(&ud);

		sz += instrLen;
		printf("%p %s", ud_insn_off(&ud), ud_insn_asm(&ud));

		
		// Check for XXX [rip + ?]
		const struct ud_operand* op = NULL;
		for(int i = 0; ; i++, op = ud_insn_opr(&ud, i))
		{
			if((op->type & UD_OP_REG) == UD_OP_REG)
			{
				if(op->base != UD_R_RIP)
					continue;
				neededFixup = true;
				
				
				int32_t ripDelta = get_rip_delta(&ud);
				void* absoluteAddr = ripDelta + ud_insn_off(&ud) + instrLen;
				printf("RIP + %x == %p\n", ripDelta, absoluteAddr);
				
				*trampolineSize += write_instr_with_rip_delta(&ud, absoluteAddr, buf, addrForInstr);
				
				break;
			}
		}

		// Check for jmp/jcc that would jump into the void after the trampoline
		size_t originalOff;
		if((originalOff = get_jump_offset(&ud, offsetOfInstr)) >= functionSize)
		{
			neededFixup = true;

			void* absoluteAddr = originalOff + (uint8_t*)function;

#if 0
			printf("\nWould jump behind trampoline (to %p -- offs: %d) into the void\n", absoluteAddr, originalOff);
			printf("Instr in Trampoline @ %d jumps %d forward\n", *trampolineSize, originalOff);
			printf("%p - %p = %p\n", trampoline, ((ptrdiff_t)function + originalOff), (ptrdiff_t)trampoline - ((ptrdiff_t)function + originalOff));
#endif

			// Normal jump
			if(buf[0] == 0xEB || buf[0] == 0xE9)
			{
				*trampolineSize += write_x64_jump(addrForInstr, absoluteAddr);
			}
			// jcc todo: jecxz, loop(?)
			else
			{
				*trampolineSize += write_jcc_jump(buf, addrForInstr, absoluteAddr);
			}
		}

		// todo: Fix relative calls
		

		printf("\n");

		// If it was a normal instruction, just copy it
		if(!neededFixup)
		{
			memcpy(addrForInstr, buf, instrLen);

			*trampolineSize += instrLen;
		}
	}



	// Build jmp back to original function
	*trampolineSize += write_x64_jump((uint8_t*)trampoline + *trampolineSize, (uint8_t*)function + sz);


	return trampoline;
}


static size_t write_instr_with_rip_delta(ud_t* ud, uint64_t value, const uint8_t* opcBuf, unsigned char* where)
{
	// 5, d, 15, 1d, 25
	printf(value);
	

	/*XXX yyy, [addr]
	jmp behind
	addr:
		value
	behind:
	*/

	return 0;
}

static int32_t get_rip_delta(ud_t* ud)
{
	const struct ud_operand* op = NULL;
	int32_t ret = 0;

	for(int i = 0; ; i++, op = ud_insn_opr(ud, i))
	{
		// Only the RIP offset is interesting
		if((op->type & UD_OP_REG) != UD_OP_REG || op->base != UD_R_RIP)
			continue;

		return op->lval.sdword;
		
	}

	return ret;
}

static size_t write_jcc_jump(const uint8_t* instruction, void* whereToWrite, void* originalAddr)
{
	/*
	jX $+?

	is transformed into:

	jnX behind
	jmp [originalAddr]
	originalAddr:
		?
	behind:
	*/

	// Stolen from https://github.com/TsudaKageyu/minhook/blob/master/src/trampoline.c
	uint8_t cond = ((instruction[0] != 0x0F ? instruction[0] : instruction[1]) & 0x0F);
	uint8_t newConditionOpc = 0x71 ^ cond; // Invert the condition in x64 mode to simplify the conditional jump logic.

	// jnx behind
	char jccTemplate[] = {newConditionOpc, 0x0F};
	memcpy((unsigned char*)whereToWrite, jccTemplate, sizeof(jccTemplate));
	size_t written = sizeof(jccTemplate);

	// jmp [originalAddr]
	written += write_x64_jump((unsigned char*)whereToWrite + written, originalAddr);

	return written;
}

static size_t get_jump_offset(ud_t* ud, size_t offsetOfInstr)
{
	const struct ud_operand* op = NULL;
	for(int i = 0; ; i++, op = ud_insn_opr(ud, i))
	{
		if(op->type == UD_OP_JIMM)
		{
			return op->lval.sdword + offsetOfInstr + ud_insn_len(ud);
		}
		// all other types are weird jumps, like jumps to a register or to []
	}

	return 0;
}

static bool is_jump(enum ud_mneomic_code op)
{
	return op >= UD_Ijo && op <= UD_Ijmp;
}

static bool is_end_of_function(enum ud_mneomic_code opc, ud_t* ud, size_t furthestJump, size_t instrOffset)
{
	if(opc == UD_Iret && furthestJump <= instrOffset)
		return true;
	else if(opc == UD_Ijmp)
	{
		const struct ud_operand* op = ud_insn_opr(ud, 0);
		if(op->type == UD_OP_MEM || op->type == UD_OP_REG)
			return true;
	}
	// todo: 0xCC ?

	return false;
}

static size_t estimate_function_length(void* function)
{
	size_t furthestJump = 0; //! Functions can have multiple return points. The library assumes that if a RET is hit and a jmp/jcc went over it, this RET isn't the last
	unsigned char* p = function;
	size_t funcLen = 0;
	ud_t ud;

	ud_init(&ud);
	ud_set_input_buffer(&ud, function, (size_t)function | (PAGE_BOUNDARY - 1)); // Can't read further than page boundary - there be dragons
	ud_set_pc(&ud, (uint64_t)function);
	ud_set_mode(&ud, 64);
	ud_set_syntax(&ud, UD_SYN_INTEL);

	while(funcLen < MINIMUM_REQUIRED_FUNCTION_LENGTH_LONG_HOOK && !ud_input_end(&ud))
	{
		size_t instrLen = 0,
			offsetOfInstr = funcLen;
		if(!(instrLen = ud_disassemble(&ud)) || ud.error)
			return 0;

		p += instrLen;
		funcLen += instrLen;

		printf("%p %s", ud_insn_off(&ud), ud_insn_asm(&ud));

		enum ud_mneomic_code op = ud_insn_mnemonic(&ud);
		if(is_jump(op))
		{
			size_t off = get_jump_offset(&ud, offsetOfInstr);
			if(off > furthestJump)
				furthestJump = off;
		}

		if(is_end_of_function(op, &ud, furthestJump, offsetOfInstr))
		{
			printf("end of function\n");
			break;
		}

		printf("\n");
	}

	return funcLen;
}

static size_t write_x64_jump(unsigned char* where, void* toWhere)
{
	const char jmpMemTemplate[] = {0x48, 0xFF, 0x25, 0x00, 0x00, 0x00, 0x00};

	memcpy(where, jmpMemTemplate, sizeof(jmpMemTemplate));
	where += sizeof(jmpMemTemplate);

	memcpy(where, &toWhere, sizeof(toWhere));

	return sizeof(jmpMemTemplate)+sizeof(toWhere);
}