#include <spslr.h>
#include <stddef.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, uint32_t target);
static void spslr_selfpatch_patch_ipins(void);

static void reorder_object(void* dst, const void* src, uint32_t target);
static int64_t spslr_calculate_ipin_value(uint32_t start);

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");

	spslr_selfpatch_patch_dpins();
	spslr_selfpatch_patch_ipins();

	spslr_randomizer_clear();
}

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

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

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

	for (uint32_t 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);
	}
}

void spslr_selfpatch_patch_dpin(void* addr, uint32_t 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];

	void* reorder_buffer = spslr_env_malloc(t->size);
	spslr_env_memset(reorder_buffer, 0, t->size);
	reorder_object(reorder_buffer, addr, target);

	spslr_env_poke_data(addr, reorder_buffer, t->size);

	spslr_env_free(reorder_buffer);
}

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

		int64_t value = spslr_calculate_ipin_value(ip->program);

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

int64_t spslr_calculate_ipin_value(uint32_t start) {
	int64_t res = 0;

	uint32_t 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;
}