#include <spslr.h>

#include "randomizer.h"
#include "env.h"
#include "spslr_list_link.h"

#define SPSLR_SANITY_CHECK

static void spslr_selfpatch_load_targets(void);
static void spslr_selfpatch_randomize_targets(void);
static void spslr_selfpatch_patch_dpins(void);
static void spslr_selfpatch_patch_dpin(void* addr, spslr_u32 target);
static void spslr_selfpatch_patch_ipins(void);

static void reorder_object(void* dst, const void* src, spslr_u32 target);
static spslr_s64 spslr_calculate_ipin_value(spslr_u32 start);

static void* reorder_buffer = SPSLR_NULL;

static void allocate_reorder_buffer();
static void release_reorder_buffer();

void spslr_selfpatch(void) {
	if (spslr_randomizer_init() < 0)
		spslr_env_panic("failed to initialize randomizer");

	if (spslr_randomize() < 0)
		spslr_env_panic("failed to randomize targets");

	allocate_reorder_buffer();
	spslr_selfpatch_patch_dpins();
	release_reorder_buffer();

	spslr_selfpatch_patch_ipins();

	spslr_randomizer_clear();
}

static void allocate_reorder_buffer() {
	if (reorder_buffer)
		return;

	spslr_u32 max_target_size = 0;
	for (spslr_u32 i = 0; i < spslr_target_cnt; i++) {
		if (spslr_targets[i].size > max_target_size)
			max_target_size = spslr_targets[i].size;
	}

	reorder_buffer = spslr_env_malloc(max_target_size);
	if (!reorder_buffer)
		spslr_env_panic("failed to allocate reorder buffer");
}

static void release_reorder_buffer() {
	if (!reorder_buffer)
		return;

	spslr_env_free(reorder_buffer);
	reorder_buffer = SPSLR_NULL;
}

static void spslr_selfpatch_patch_dpins(void) {
	for (spslr_u32 dpidx = 0; dpidx < spslr_dpin_cnt; dpidx++) {
		const struct spslr_dpin* dp = &spslr_dpins[dpidx];
		spslr_selfpatch_patch_dpin((void*)dp->addr, dp->target);
	}
}

static void reorder_object(void* dst, const void* src, spslr_u32 target) {
	spslr_u32 field_count;
	if (spslr_randomizer_get_target(target, SPSLR_NULL, &field_count))
		spslr_env_panic("failed to get target field count");

	const spslr_u8* src_countable = (const spslr_u8*)src;
	spslr_u8* dst_countable = (spslr_u8*)dst;

	for (spslr_u32 i = 0; i < field_count; i++) {
		struct spslr_randomizer_field_info finfo;
		if (spslr_randomizer_get_field(target, i, SPSLR_RANDOMIZER_FIELD_IDX_MODE_FINAL, &finfo))
			spslr_env_panic("failed to get ordered field descriptor");

		spslr_env_memcpy(dst_countable + finfo.offset, src_countable + finfo.initial_offset, finfo.size);
	}
}

static void spslr_selfpatch_patch_dpin(void* addr, spslr_u32 target) {
#ifdef SPSLR_SANITY_CHECK
	if (target >= spslr_target_cnt)
		spslr_env_panic("dpin refers to invalid target");
#endif

	const struct spslr_target* t = &spslr_targets[target];

	spslr_env_memset(reorder_buffer, 0, t->size);
	reorder_object(reorder_buffer, addr, target);
	spslr_env_poke_data(addr, reorder_buffer, t->size);
}

static void spslr_selfpatch_patch_ipins(void) {
	for (spslr_u32 ipidx = 0; ipidx < spslr_ipin_cnt; ipidx++) {
		const struct spslr_ipin* ip = &spslr_ipins[ipidx];

		spslr_s64 value = spslr_calculate_ipin_value(ip->program);

		switch (ip->size) {
			case 1:
				spslr_env_poke_text_8((void*)ip->addr, (spslr_u8)value);
				break;
			case 2:
				spslr_env_poke_text_16((void*)ip->addr, (spslr_u16)value);
				break;
			case 4:
				spslr_env_poke_text_32((void*)ip->addr, (spslr_u32)value);
				break;
			case 8:
				spslr_env_poke_text_64((void*)ip->addr, (spslr_u64)value);
				break;
#ifdef SPSLR_SANITY_CHECK
			default:
				spslr_env_panic("invalid ipin size");
#endif
		}
	}
}

static spslr_s64 spslr_calculate_ipin_value(spslr_u32 start) {
	spslr_s64 res = 0;

	spslr_u32 pc = start;
	while (true) {
#ifdef SPSLR_SANITY_CHECK
		if (pc >= spslr_ipin_op_cnt)
			spslr_env_panic("ipin op out of bounds");
#endif

		int end_flag = 0;

		const struct spslr_ipin_op* op = &spslr_ipin_ops[pc++];
		switch (op->code) {
			case SPSLR_IPIN_OP_PATCH:
				end_flag = 1;
				break;
			case SPSLR_IPIN_OP_ADD_INITIAL_OFFSET:
				{
					struct spslr_randomizer_field_info finfo;
					if (spslr_randomizer_get_field(op->op0.add_initial_offset_target,
							op->op1.add_initial_offset_field, SPSLR_RANDOMIZER_FIELD_IDX_MODE_ORIGINAL, &finfo))
						spslr_env_panic("failed to get initial field offset");

					res += finfo.initial_offset;
				}
				break;
			case SPSLR_IPIN_OP_ADD_OFFSET:
				{
					struct spslr_randomizer_field_info finfo;
					if (spslr_randomizer_get_field(op->op0.add_initial_offset_target,
							op->op1.add_offset_field, SPSLR_RANDOMIZER_FIELD_IDX_MODE_ORIGINAL, &finfo))
						spslr_env_panic("failed to get initial field offset");

					res += finfo.offset;
				}
				break;
			case SPSLR_IPIN_OP_SUB_INITIAL_OFFSET:
				{
					struct spslr_randomizer_field_info finfo;
					if (spslr_randomizer_get_field(op->op0.add_initial_offset_target,
							op->op1.sub_initial_offset_field, SPSLR_RANDOMIZER_FIELD_IDX_MODE_ORIGINAL, &finfo))
						spslr_env_panic("failed to get initial field offset");

					res -= finfo.initial_offset;
				}
				break;
			case SPSLR_IPIN_OP_SUB_OFFSET:
				{
					struct spslr_randomizer_field_info finfo;
					if (spslr_randomizer_get_field(op->op0.add_initial_offset_target,
							op->op1.sub_offset_field, SPSLR_RANDOMIZER_FIELD_IDX_MODE_ORIGINAL, &finfo))
						spslr_env_panic("failed to get initial field offset");

					res -= finfo.offset;
				}
				break;
			case SPSLR_IPIN_OP_ADD_CONST:
				res += op->op0.add_const_value;
				break;
#ifdef SPSLR_SANITY_CHECK
			default:
				spslr_env_panic("invalid ipin op");
#endif
		}

		if (end_flag)
			break;
	}

	return res;
}