#include <assert.h>
#include <sys/mman.h>

#include <ejit/ejit.h>

#include "common.h"

static void emit_insn_i(struct ejit_func *f, enum ejit_opcode op, size_t r0,
                        size_t r1, long o)
{
	struct ejit_insn i = {.op = op, .r0 = r0, .r1 = r1, .o = o};
	vec_append(&f->insns, &i);
}

static void emit_insn_r(struct ejit_func *f, enum ejit_opcode op, size_t r0,
                        size_t r1, size_t r2)
{
	struct ejit_insn i = {.op = op, .r0 = r0, .r1 = r1, .r2 = r2};
	vec_append(&f->insns, &i);
}

static void emit_insn_p(struct ejit_func *f, enum ejit_opcode op, size_t r0,
                        size_t r1, void *p)
{
	struct ejit_insn i = {.op = op, .r0 = r0, .r1 = r1, .p = p};
	vec_append(&f->insns, &i);
}

static void emit_insn_f(struct ejit_func *f, enum ejit_opcode op, size_t r0,
                        size_t r1, double d)
{
	struct ejit_insn i = {.op = op, .r0 = r0, .r1 = r1, .d = d};
	vec_append(&f->insns, &i);
}

struct ejit_func *ejit_create_func(enum ejit_type rtype, size_t argc,
                                   const struct ejit_operand args[argc])
{
	struct ejit_func *f = malloc(sizeof(struct ejit_func));
	assert(f);

	f->rtype = rtype;

	f->insns = vec_create(sizeof(struct ejit_insn));
	f->labels = vec_create(sizeof(size_t));
	f->arena = NULL;
	f->size = 0;

	for (size_t i = 0; i < argc; ++i) {
		switch (args[i].kind) {
		case EJIT_OPERAND_GPR: emit_insn_r(f, PARAM, i, args[i].type,
			                           args[i].r); break;
		case EJIT_OPERAND_FPR: emit_insn_r(f, PARAM_F, i, args[i].type,
			                           args[i].r); break;
		default: abort();
		}
	}

	emit_insn_i(f, START, 0, 0, 0);
	return f;
}

void ejit_compile_func(struct ejit_func *f, size_t gpr, size_t fpr, bool use_64)
{
	ejit_select_compile_func(f, gpr, fpr, use_64, true);
}

void ejit_select_compile_func(struct ejit_func *f, size_t gpr, size_t fpr,
                              bool use_64, bool try_jit)
{
	/* emit a final end instruction in case user didn't do a return */
	emit_insn_i(f, END, 0, 0, 0);

	f->gpr = gpr;
	f->fpr = fpr;

	/* try to jit compile if possible */
	if (try_jit && ejit_compile(f, use_64))
		return;

	/* otherwise, convert opcodes to address labels */

	void **labels;
	/* just get labels, don't actually run anything yet */
	ejit_interp(f, 0, NULL, NULL, false, &labels);
	foreach_vec(ii, f->insns) {
		struct ejit_insn i = vect_at(struct ejit_insn, f->insns, ii);
		void *addr = labels[i.op];
		assert(addr);
		i.addr = addr;
		vect_at(struct ejit_insn, f->insns, ii) = i;
	}
}

void ejit_destroy_func(struct ejit_func *f)
{
	if (f->arena)
		munmap(f->arena, f->size);

	vec_destroy(&f->insns);
	vec_destroy(&f->labels);
	free(f);
}

struct ejit_label ejit_label(struct ejit_func *f)
{
	size_t addr = vec_len(&f->insns);
	vec_append(&f->labels, &addr);
	return (struct ejit_label){.addr = addr};
}

void ejit_patch(struct ejit_func *f, struct ejit_reloc r, struct ejit_label l)
{
	struct ejit_insn i = vect_at(struct ejit_insn, f->insns, r.insn);
	/** @todo some assert that checks the opcode? */
	i.r0 = l.addr;
	vect_at(struct ejit_insn, f->insns, r.insn) = i;
}

void ejit_calli(struct ejit_func *s, struct ejit_func *f, size_t argc,
                const struct ejit_operand args[argc])
{
	for (size_t i = 0; i < argc; ++i) {
		switch (args[i].kind) {
		case EJIT_OPERAND_GPR: emit_insn_r(s, ARG, i, args[i].type,
			                           args[i].r); break;
		case EJIT_OPERAND_FPR: emit_insn_r(s, ARG_F, i, args[i].type,
			                           args[i].r); break;
		case EJIT_OPERAND_IMM: emit_insn_i(s, ARG_I, i, args[i].type,
			                           args[i].r); break;
		case EJIT_OPERAND_FLT: emit_insn_f(s, ARG_FI, i, args[i].type,
			                           args[i].d); break;
		default: abort();
		}
	}

	emit_insn_p(s, CALLI, 0, 0, f);
}

void ejit_escapei(struct ejit_func *s, ejit_escape_t f, size_t argc,
                  const struct ejit_operand args[argc])
{
	for (size_t i = 0; i < argc; ++i) {
		switch (args[i].kind) {
		case EJIT_OPERAND_GPR: emit_insn_r(s, ARG, i, args[i].type,
			                           args[i].r); break;
		case EJIT_OPERAND_FPR: emit_insn_r(s, ARG_F, i, args[i].type,
			                           args[i].r); break;
		case EJIT_OPERAND_IMM: emit_insn_r(s, ARG_I, i, args[i].type,
			                           args[i].r); break;
		case EJIT_OPERAND_FLT: emit_insn_f(s, ARG_FI, i, args[i].type,
			                           args[i].d); break;
		default: abort();
		}
	}

	emit_insn_p(s, ESCAPEI, 0, 0, f);
}

void ejit_escapei_f(struct ejit_func *s, ejit_escape_f_t f, size_t argc,
                    const struct ejit_operand args[argc])
{
	for (size_t i = 0; i < argc; ++i) {
		switch (args[i].kind) {
		case EJIT_OPERAND_GPR: emit_insn_r(s, ARG, i, args[i].type,
			                           args[i].r); break;
		case EJIT_OPERAND_FPR: emit_insn_r(s, ARG_F, i, args[i].type,
			                           args[i].r); break;
		case EJIT_OPERAND_IMM: emit_insn_r(s, ARG_I, i, args[i].type,
			                           args[i].r); break;
		case EJIT_OPERAND_FLT: emit_insn_f(s, ARG_FI, i, args[i].type,
			                           args[i].d); break;
		default: abort();
		}
	}

	emit_insn_p(s, ESCAPEI_F, 0, 0, f);
}

void ejit_retval(struct ejit_func *s, struct ejit_gpr r0)
{
	emit_insn_i(s, RETVAL, r0.r, 0, 0);
}

void ejit_stxi_64(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1,
                  long o)
{
	emit_insn_i(s, STXI64, r0.r, r1.r, o);
}

void ejit_ldxi_u64(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1,
                   long o)
{
	emit_insn_i(s, LDXIU64, r0.r, r1.r, o);
}

void ejit_retr(struct ejit_func *s, struct ejit_gpr r0)
{
	emit_insn_r(s, RETR, r0.r, 0, 0);
}

void ejit_retr_f(struct ejit_func *s, struct ejit_fpr r0)
{
	emit_insn_r(s, RETR_F, r0.f, 0, 0);
}

void ejit_reti(struct ejit_func *s, long i)
{
	emit_insn_i(s, RETI, 0, 0, i);
}

void ejit_reti_f(struct ejit_func *s, double f)
{
	emit_insn_f(s, RETI_F, 0, 0, f);
}

void ejit_addr(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1,
               struct ejit_gpr r2)
{
	emit_insn_r(s, ADDR, r0.r, r1.r, r2.r);
}

void ejit_addr_f(struct ejit_func *s, struct ejit_fpr r0, struct ejit_fpr r1,
                 struct ejit_fpr r2)
{
	emit_insn_r(s, ADDR_F, r0.f, r1.f, r2.f);
}

void ejit_addi(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1,
               long o)
{
	emit_insn_i(s, ADDI, r0.r, r1.r, o);
}

void ejit_absr_f(struct ejit_func *s, struct ejit_fpr r0, struct ejit_fpr r1)
{
	emit_insn_i(s, ABSR_F, r0.f, r1.f, 0);
}

void ejit_subr(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1,
               struct ejit_gpr r2)
{
	emit_insn_r(s, SUBR, r0.r, r1.r, r2.r);
}

void ejit_subr_f(struct ejit_func *s, struct ejit_fpr r0, struct ejit_fpr r1,
                 struct ejit_fpr r2)
{
	emit_insn_r(s, SUBR_F, r0.f, r1.f, r2.f);
}

void ejit_subi(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1,
               long o)
{
	emit_insn_i(s, SUBI, r0.r, r1.r, o);
}

void ejit_mulr(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1,
               struct ejit_gpr r2)
{
	emit_insn_r(s, MULR, r0.r, r1.r, r2.r);
}

void ejit_divr(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1,
               struct ejit_gpr r2)
{
	emit_insn_r(s, DIVR, r0.r, r1.r, r2.r);
}

void ejit_negr(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1)
{
	emit_insn_i(s, NEGR, r0.r, r1.r, 0);
}

void ejit_movi(struct ejit_func *s, struct ejit_gpr r0, long o)
{
	emit_insn_i(s, MOVI, r0.r, 0, o);
}

void ejit_movr(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1)
{
	if (r0.r == r1.r)
		return;

	emit_insn_i(s, MOVR, r0.r, r1.r, 0);
}

void ejit_eqr(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1,
              struct ejit_gpr r2)
{
	emit_insn_r(s, EQR, r0.r, r1.r, r2.r);
}

void ejit_ltr(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1,
              struct ejit_gpr r2)
{
	emit_insn_r(s, LTR, r0.r, r1.r, r2.r);
}

struct ejit_reloc ejit_bltr(struct ejit_func *s, struct ejit_gpr r0,
                            struct ejit_gpr r1)
{
	size_t addr = vec_len(&s->insns);
	emit_insn_r(s, BLTR, 0, r0.r, r1.r);
	return (struct ejit_reloc){.insn = addr};
}

struct ejit_reloc ejit_bnei(struct ejit_func *s, struct ejit_gpr r0, long o)
{
	size_t addr = vec_len(&s->insns);
	emit_insn_i(s, BNEI, 0, r0.r, o);
	return (struct ejit_reloc){.insn = addr};
}

struct ejit_reloc ejit_beqi(struct ejit_func *s, struct ejit_gpr r0, long o)
{
	size_t addr = vec_len(&s->insns);
	emit_insn_i(s, BEQI, 0, r0.r, o);
	return (struct ejit_reloc){.insn = addr};
}

struct ejit_reloc ejit_bgti(struct ejit_func *s, struct ejit_gpr r0, long o)
{
	size_t addr = vec_len(&s->insns);
	emit_insn_i(s, BGTI, 0, r0.r, o);
	return (struct ejit_reloc){.insn = addr};
}

struct ejit_reloc ejit_jmp(struct ejit_func *s)
{
	size_t addr = vec_len(&s->insns);
	emit_insn_i(s, JMP, 0, 0, 0);
	return (struct ejit_reloc){.insn = addr};
}

static struct interp_state create_interp_state()
{
	struct interp_state state;
	state.gprs = vec_create(sizeof(long));
	state.fprs = vec_create(sizeof(double));
	state.args = vec_create(sizeof(struct ejit_arg));
	return state;
}

static void destroy_interp_state(struct interp_state state)
{
	vec_destroy(&state.gprs);
	vec_destroy(&state.fprs);
	vec_destroy(&state.args);
}

long ejit_run_interp(struct ejit_func *f, size_t argc,
                     struct ejit_arg args[argc], struct interp_state *state)
{
	assert(f->gpr && "trying to run a function that hasn't been compiled");
	assert(f->rtype == EJIT_VOID || ejit_int_type(f->rtype));
	if (f->arena)
		return ((ejit_escape_t)f->arena)(argc, args);

	return ejit_interp(f, argc, args, state, true, NULL).r;
}

long ejit_run_func(struct ejit_func *f, size_t argc, struct ejit_arg args[argc])
{
	assert(f->gpr && "trying to run a function that hasn't been compiled");
	assert(f->rtype == EJIT_VOID || ejit_int_type(f->rtype));
	if (f->arena)
		return ((ejit_escape_t)f->arena)(argc, args);

	struct interp_state state = create_interp_state();
	long r = ejit_interp(f, argc, args, &state, true, NULL).r;
	destroy_interp_state(state);
	return r;
}

double ejit_run_func_f(struct ejit_func *f, size_t argc,
                       struct ejit_arg args[argc])
{
	assert(f->fpr && "trying to run a function that hasn't been compiled");
	assert(ejit_float_type(f->rtype));
	if (f->arena)
		return ((ejit_escape_f_t)f->arena)(argc, args);

	struct interp_state state = create_interp_state();
	double r = ejit_interp(f, argc, args, &state, true, NULL).d;
	destroy_interp_state(state);
	return r;
}