#include "config.h"

#include <check.h>
#include <errno.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#include "asmsh_check.h"
#include "asm_env.h"
#include "compile.h"

#define REAL_CHILD "../child"
#define STAMP_CHILD "samples/asm_env_stamps"

static const long int stamp1 = 0x4242; // %r15
static const long int stamp2 = -1; // %r14

START_TEST(test_env)
{
	const char childpath[] = REAL_CHILD;
	asmsh_env_t *env = asmsh_env(childpath);
	
	ck_assert_ptr_nonnull(env);
	ck_assert_str_eq(childpath, env->childpath);

	asmsh_env_free(env);
}
END_TEST

START_TEST(test_asm_stamps)
{
	const char childpath[] = STAMP_CHILD;
	asmsh_env_t *env = asmsh_env(childpath);
	
	ck_assert_ptr_nonnull(env);
	ck_assert_str_eq(env->childpath, childpath);

	ck_assert_mem_eq(&stamp1, &(env->regs.r15), 8);
	ck_assert_mem_eq(&stamp2, &(env->regs.r14), 8);

	asmsh_env_free(env);
}
END_TEST

START_TEST(test_env_update)
{
	const char childpath[] = STAMP_CHILD;
	asmsh_env_t *env = asmsh_env(childpath);

	ck_assert_ptr_nonnull(env);
	ck_assert_str_eq(env->childpath, childpath);

	asmsh_env_update_maps(env);

	asmsh_env_free(env);
}
END_TEST

START_TEST(test_env_exit42)
{
	const char childpath[] = REAL_CHILD;

	asmsh_asmc_ctx_t *cctx;
	asmsh_bytecode_t bcode;
	int ret, status;

	asmsh_env_t *env = asmsh_env(childpath);
	ck_assert_ptr_nonnull(env);

	cctx = asmsh_asmc_ctx_default();
	ck_assert_ptr_nonnull(cctx);

	ck_assert_int_ge(asmsh_asmc_compile(cctx, "mov $60, %rax", &bcode), 0);
	ck_assert_int_ge(asmsh_env_write_code(env, &bcode), 0);
	ck_assert_int_ge(asmsh_asmc_compile(cctx, "mov $42, %rdi", &bcode), 0);
	ck_assert_int_ge(asmsh_env_write_code(env, &bcode), 0);
	ck_assert_int_ge(asmsh_asmc_compile(cctx, "syscall", &bcode), 0);
	ck_assert_int_ge(asmsh_env_write_code(env, &bcode), 0);

	asmsh_asmc_ctx_free(cctx);

	for(int i=0; i<2; i++)
	{
		ret = asmsh_env_step(env, &status);
		ck_assert_int_eq(ret, 0);
		ck_assert_msg(WCOREDUMP(status) ==  0, "SEGFAULT");
		ck_assert_int_ne(WIFSTOPPED(status), 0);
		ck_assert_int_eq(status>>8, 5);
	}
	ret = asmsh_env_step(env, &status);
	ck_assert_int_eq(ret, 1);
	ck_assert(WIFEXITED(status));
	ck_assert_int_eq(WEXITSTATUS(status), 42);

	asmsh_env_free(env);

}
END_TEST

START_TEST(test_env_steps)
{
	const char childpath[] = REAL_CHILD;
	const char *instr[] = {
		"mov $60, %rax",
		"mov $0x42,%rdi",
		"syscall",
	};

	asmsh_asmc_ctx_t *cctx;
	asmsh_bytecode_t bcode;
	int ret, status, i;

	asmsh_env_t *env = asmsh_env(childpath);
	ck_assert_ptr_nonnull(env);

	cctx = asmsh_asmc_ctx_default();
	ck_assert_ptr_nonnull(cctx);


	for(i=0; i<(sizeof(instr)/sizeof(*instr))-1; i++)
	{
		ck_assert_int_ge(asmsh_asmc_compile(cctx, instr[i], &bcode), 0);
		ck_assert_int_ge(asmsh_env_write_code(env, &bcode), 0);
		ret = asmsh_env_step(env, &status);
		ck_assert_int_eq(ret, 0);
		ck_assert_msg(WCOREDUMP(status) ==  0, "SEGFAULT");
		ck_assert_int_ne(WIFSTOPPED(status), 0);
		ck_assert_int_eq(status>>8, 5);
	}
	ck_assert_int_ge(asmsh_asmc_compile(cctx, instr[i], &bcode), 0);
	ck_assert_int_ge(asmsh_env_write_code(env, &bcode), 0);
	ret = asmsh_env_step(env, &status);
	ck_assert_int_eq(ret, 1);
	ck_assert(WIFEXITED(status));
	ck_assert_int_eq(WEXITSTATUS(status), 0x42);

	asmsh_asmc_ctx_free(cctx);
	asmsh_env_free(env);

}
END_TEST

START_TEST(test_env_regs)
{
	const char childpath[] = REAL_CHILD;
	asmsh_env_t *env = asmsh_env(childpath);
	ck_assert_ptr_nonnull(env);

	asmsh_asmc_ctx_t *cctx;
	asmsh_bytecode_t bcode;
	int ret, status, i;
	struct _instr_tests {
		char *instr;
		long long unsigned int *reg;
		u_int64_t val;
	} tinstr[] = {
		{
			.instr = "mov $42, %r15",
			.reg = &env->regs.r15,
			.val = 42,
		},
		{
			.instr = "mov %r15, %rcx",
			.reg = &env->regs.rcx,
			.val = 42,
		},
		{
			.instr = "mov $24, %rbx",
			.reg = &env->regs.rbx,
			.val = 24,
		},
		{
			.instr = "add %rbx, %rcx",
			.reg = &env->regs.rcx,
			.val = 0x42,
		},
		{
			.instr = "mov $60, %rax",
			.reg = &env->regs.rax,
			.val = 60,
		},
		{
			.instr = "xor %rdi, %rdi",
			.reg = &env->regs.rdi,
			.val = 0,
		},
		{
			.instr = "syscall",
			.reg = NULL,
			.val = 0,
		},
	};

	cctx = asmsh_asmc_ctx_default();
	ck_assert_ptr_nonnull(cctx);

	for(i=0; i<(sizeof(tinstr)/sizeof(*tinstr))-1;i++)
	{
		ck_assert_int_ge(asmsh_asmc_compile(cctx, tinstr[i].instr, &bcode), 0);
		ck_assert_int_ge(asmsh_env_write_code(env, &bcode), 0);
		ret = asmsh_env_step(env, &status);
		ck_assert_int_eq(ret, 0);
		ck_assert_msg(WCOREDUMP(status) ==  0, "SEGFAULT");
		ck_assert_int_ne(WIFSTOPPED(status), 0);
		ck_assert_int_eq(status>>8, 5);
		if(!tinstr[i].reg) { continue; }
		asmsh_env_update_regs(env);
		ck_assert_int_eq(*(tinstr[i].reg), tinstr[i].val);
	}
	ck_assert_int_ge(asmsh_asmc_compile(cctx, tinstr[i].instr, &bcode), 0);
	ck_assert_int_ge(asmsh_env_write_code(env, &bcode), 0);
	ret = asmsh_env_step(env, &status);
	ck_assert_int_eq(ret, 1);
	ck_assert(WIFEXITED(status));
	ck_assert_int_eq(WEXITSTATUS(status), 0);

	asmsh_asmc_ctx_free(cctx);
	asmsh_env_free(env);

}
END_TEST


START_TEST(test_embed_env_exit42)
{
	const char childpath[] = REAL_CHILD;

	asmsh_asmc_ctx_t *cctx;
	asmsh_bytecode_t bcode;
	int ret, status;

	asmsh_env_t *env = asmsh_env(NULL);
	ck_assert_ptr_nonnull(env);

	cctx = asmsh_asmc_ctx_default();
	ck_assert_ptr_nonnull(cctx);

	ck_assert_int_ge(asmsh_asmc_compile(cctx, "mov $60, %rax", &bcode), 0);
	ck_assert_int_ge(asmsh_env_write_code(env, &bcode), 0);
	ck_assert_int_ge(asmsh_asmc_compile(cctx, "mov $42, %rdi", &bcode), 0);
	ck_assert_int_ge(asmsh_env_write_code(env, &bcode), 0);
	ck_assert_int_ge(asmsh_asmc_compile(cctx, "syscall", &bcode), 0);
	ck_assert_int_ge(asmsh_env_write_code(env, &bcode), 0);

	asmsh_asmc_ctx_free(cctx);

	for(int i=0; i<2; i++)
	{
		ret = asmsh_env_step(env, &status);
		ck_assert_int_eq(ret, 0);
		ck_assert_msg(WCOREDUMP(status) ==  0, "SEGFAULT");
		ck_assert_int_ne(WIFSTOPPED(status), 0);
		ck_assert_int_eq(status>>8, 5);
	}
	ret = asmsh_env_step(env, &status);
	ck_assert_int_eq(ret, 1);
	ck_assert(WIFEXITED(status));
	ck_assert_int_eq(WEXITSTATUS(status), 42);

	asmsh_env_free(env);

}
END_TEST


ASMSH_CHECK_START("Testing asm env", "testing asm environment")
	ASMSH_ADD_TEST(test_env);
	ASMSH_ADD_TEST(test_asm_stamps);
	ASMSH_ADD_TEST(test_env_update);
	ASMSH_ADD_TEST(test_env_exit42);
	ASMSH_ADD_TEST(test_env_steps);
	ASMSH_ADD_TEST(test_env_regs);
	ASMSH_ADD_TEST(test_embed_env_exit42);
ASMSH_CHECK_END