#define i_implement
#include"vm.h"
#undef i_implement

#include"str.h"

#include<math.h>
#include<malloc.h>

size_t lvm_call(LVM *L, LFunc *func, size_t arg_count, set_LValueU *heap, LRegSet *regset) {
	if(func->is_native) {
		return func->native_func(L, func->ud, arg_count, heap, regset);
	}
	
	static const void *dispatch_table[] = {
		[L_GETGLOBAL] = &&do_getglobal,
		[L_SETGLOBAL] = &&do_setglobal,
		[L_SETINT16] = &&do_setint16,
		[L_SETINT32] = &&do_setint32,
		[L_SETFLOAT] = &&do_setfloat,
		[L_SETSTR] = &&do_setstr,
		[L_SETTABLE] = &&do_settable,
		[L_SETBOOL] = &&do_setbool,
		[L_SETNIL] = &&do_setnil,
		[L_SETFUNC] = &&do_setfunc,
		[L_ADD] = &&do_add,
		[L_SUB] = &&do_sub,
		[L_MUL] = &&do_mul,
		[L_DIV] = &&do_div,
		[L_MOD] = &&do_mod,
		[L_RET] = &&do_ret,
		[L_JNOTCOND] = &&do_jnotcond,
		[L_MOVE] = &&do_move,
		[L_CALL] = &&do_call,
		[L_JUMP] = &&do_jump,
		[L_ADVANCETEST] = &&do_advancetest,
		[L_COND_EQ] = &&do_cond_eq,
		[L_COND_NEQ] = &&do_cond_neq,
		[L_SETFIELD] = &&do_setfield,
		[L_GETFIELD] = &&do_getfield,
		[L_LEN] = &&do_len,
	};
	
	LUnit *unit = func->unit;
	
	LInst *inst = func->lua_instrs;
	#define DISPATCH() goto *dispatch_table[(++inst)->opcode]
	
	LThreadPrivates privates = {.regset = regset, .heap = heap};
	
	inst--;
	DISPATCH();
	
do_getglobal:;
	{
		uint8_t *area = unit->abyss + inst->bc;
		size_t len = *(uint16_t*) area;
		area += 2;
		
		LString *str = realloc(NULL, sizeof(*str) + len);
		str->length = len;
		memcpy(str->data, area, len);
		
		regset->regs[inst->a] = ltable_get(func->env, lvalue_from_string(str));
		
		set_LValueU_insert(heap, lvalue_from_string(str).u);
		
		lvm_gc_alert(L, &privates, sizeof(*str) + len);
	}
	DISPATCH();
	
do_setglobal:;
	{
		uint8_t *area = unit->abyss + inst->bc;
		size_t len = *(uint16_t*) area;
		area += 2;
		
		LString *str = realloc(NULL, sizeof(*str) + len);
		str->length = len;
		memcpy(str->data, area, len);
		
		ltable_set(func->env, lvalue_from_string(str), regset->regs[inst->a]);
		
		set_LValueU_insert(heap, lvalue_from_string(str).u);
		
		lvm_gc_alert(L, &privates, sizeof(*str) + len);
	}
	DISPATCH();
	
do_setint16:;
	regset->regs[inst->a] = lvalue_from_int32((int16_t) inst->bc);
	DISPATCH();
	
do_setint32:;
	regset->regs[inst->a] = lvalue_from_int32(*(int32_t*) &unit->abyss[inst->bc]);
	DISPATCH();
	
do_setfloat:;
	DISPATCH();
	
do_setstr:;
	{
		uint8_t *area = unit->abyss + inst->bc;
		size_t len = *(uint16_t*) area;
		area += 2;
		
		LString *str = realloc(NULL, sizeof(*str) + len);
		str->length = len;
		memcpy(str->data, area, len);
		
		regset->regs[inst->a] = lvalue_from_string(str);
		
		set_LValueU_insert(heap, lvalue_from_string(str).u);
		
		lvm_gc_alert(L, &privates, sizeof(*str) + len);
	}
	DISPATCH();
	
do_settable:;
	{
		LTable *tbl = ltable_new(inst->bc);
		regset->regs[inst->a] = lvalue_from_table(tbl);
		set_LValueU_insert(heap, lvalue_from_table(tbl).u);
	}
	DISPATCH();
	
do_setbool:;
	regset->regs[inst->a] = lvalue_from_bool(inst->b);
	DISPATCH();
	
do_setnil:;
	regset->regs[inst->a] = lvalue_from_nil();
	DISPATCH();
	
do_setfunc:;
	regset->regs[inst->a] = lvalue_from_func(&func->unit->funcs[inst->bc]);
	DISPATCH();
	
do_add:;
	{
		LValue x = regset->regs[inst->b];
		LValue y = regset->regs[inst->c];
		if(lvalue_tag(x) == LTAG_I32 && lvalue_tag(y) == LTAG_FLOAT) {
			regset->regs[inst->a] = lvalue_from_double(lvalue_to_int32(x) + y.f);
		} else if(lvalue_tag(x) == LTAG_FLOAT && lvalue_tag(y) == LTAG_I32) {
			regset->regs[inst->a] = lvalue_from_double(x.f + lvalue_to_int32(y));
		} else if(lvalue_tag(x) == LTAG_I32 && lvalue_tag(y) == LTAG_I32) {
			regset->regs[inst->a] = lvalue_from_int32(lvalue_to_int32(x) + lvalue_to_int32(y));
		} else goto err;
	}
	DISPATCH();
	
do_sub:;
	{
		LValue x = regset->regs[inst->b];
		LValue y = regset->regs[inst->c];
		if(lvalue_tag(x) == LTAG_I32 && lvalue_tag(y) == LTAG_FLOAT) {
			regset->regs[inst->a] = lvalue_from_double(lvalue_to_int32(x) - y.f);
		} else if(lvalue_tag(x) == LTAG_FLOAT && lvalue_tag(y) == LTAG_I32) {
			regset->regs[inst->a] = lvalue_from_double(x.f - lvalue_to_int32(y));
		} else if(lvalue_tag(x) == LTAG_I32 && lvalue_tag(y) == LTAG_I32) {
			regset->regs[inst->a] = lvalue_from_int32(lvalue_to_int32(x) - lvalue_to_int32(y));
		} else goto err;
	}
	DISPATCH();
	
do_mul:;
	{
		LValue x = regset->regs[inst->b];
		LValue y = regset->regs[inst->c];
		if(lvalue_tag(x) == LTAG_I32 && lvalue_tag(y) == LTAG_FLOAT) {
			regset->regs[inst->a] = lvalue_from_double(lvalue_to_int32(x) * y.f);
		} else if(lvalue_tag(x) == LTAG_FLOAT && lvalue_tag(y) == LTAG_I32) {
			regset->regs[inst->a] = lvalue_from_double(x.f * lvalue_to_int32(y));
		} else if(lvalue_tag(x) == LTAG_I32 && lvalue_tag(y) == LTAG_I32) {
			regset->regs[inst->a] = lvalue_from_int32(lvalue_to_int32(x) * lvalue_to_int32(y));
		} else goto err;
	}
	DISPATCH();
	
do_div:;
	{
		LValue x = regset->regs[inst->b];
		LValue y = regset->regs[inst->c];
		if(lvalue_tag(x) == LTAG_I32 && lvalue_tag(y) == LTAG_FLOAT) {
			regset->regs[inst->a] = lvalue_from_double(lvalue_to_int32(x) / y.f);
		} else if(lvalue_tag(x) == LTAG_FLOAT && lvalue_tag(y) == LTAG_I32) {
			regset->regs[inst->a] = lvalue_from_double(x.f / lvalue_to_int32(y));
		} else if(lvalue_tag(x) == LTAG_I32 && lvalue_tag(y) == LTAG_I32) {
			int32_t yv = lvalue_to_int32(y);
			if(yv == 0) {
				regset->regs[inst->a] = lvalue_from_nil();
			} else {
				regset->regs[inst->a] = lvalue_from_int32(lvalue_to_int32(x) / yv);
			}
		} else goto err;
	}
	DISPATCH();
	
do_mod:;
	{
		LValue x = regset->regs[inst->b];
		LValue y = regset->regs[inst->c];
		if(lvalue_tag(x) == LTAG_I32 && lvalue_tag(y) == LTAG_FLOAT) {
			regset->regs[inst->a] = lvalue_from_double(fmod(fmod(lvalue_to_int32(x), y.f) + y.f, y.f));
		} else if(lvalue_tag(x) == LTAG_FLOAT && lvalue_tag(y) == LTAG_I32) {
			int32_t yv = lvalue_to_int32(y);
			regset->regs[inst->a] = lvalue_from_double(fmod(fmod(x.f, yv) + yv, yv));
		} else if(lvalue_tag(x) == LTAG_I32 && lvalue_tag(y) == LTAG_I32) {
			int32_t yv = lvalue_to_int32(y);
			if(yv == 0) {
				goto err;
			} else {
				regset->regs[inst->a] = lvalue_from_int32((lvalue_to_int32(x) % yv + yv) % yv);
			}
		} else goto err;
	}
	DISPATCH();

do_jump:;
	inst += (int16_t) inst->bc;
	
	L->safepoint_func(L, heap, regset);
	DISPATCH();

do_jnotcond:;
	{
		LValue v = regset->regs[inst->a];
		if(v.u == LTAG_NIL || v.u == LTAG_FALSE) {
			inst += (int16_t) inst->bc;
		}
	}
	DISPATCH();

do_call:;
	{
		if(lvalue_tag(regset->regs[inst->a]) != LTAG_FUNCTION) {
			goto err;
		}
		
		uint8_t *abyss_data = unit->abyss + inst->bc;
		
		uint8_t ret_vreg = abyss_data[0];
		uint8_t arg_count = abyss_data[1];
		uint8_t *args = &abyss_data[2];

		LRegSet regset2 = {.parent = regset};
		lvm_reset_regs(&regset2);
		for(int i = 0; i < arg_count; i++) {
			regset2.regs[i] = regset->regs[args[i]];
		}
		size_t returned_count = lvm_call(L, (LFunc*) (regset->regs[inst->a].u & ~LTAG_MASK), arg_count, heap, &regset2);
		
		if(returned_count) {
			// TODO: more than 1 return
			regset->regs[ret_vreg] = regset2.regs[0];
		}
	}
	DISPATCH();

do_move:;
	regset->regs[inst->a] = regset->regs[inst->b];
	DISPATCH();

do_advancetest:;
	{
		int64_t a = lvalue_to_int32(regset->regs[inst->a]);
		int64_t b = lvalue_to_int32(regset->regs[inst->b]);
		int64_t c = lvalue_to_int32(regset->regs[inst->c]);
		if(!((c >= 0 && a > b) || (c < 0 && a < b))) {
			inst++;
		}
	}
	DISPATCH();
	
do_cond_eq:;
	regset->regs[inst->a] = lvalue_from_bool(lvalue_eq(regset->regs[inst->b], regset->regs[inst->c]));
	DISPATCH();
	
do_cond_neq:;
	regset->regs[inst->a] = lvalue_from_bool(!lvalue_eq(regset->regs[inst->b], regset->regs[inst->c]));
	DISPATCH();

do_setfield:;
	{
		if(lvalue_tag(regset->regs[inst->a]) != LTAG_TABLE) {
			goto err;
		}

		if(regset->regs[inst->b].u == LTAG_NIL) {
			goto err;
		}
		
		LTable *tbl = (void*) (regset->regs[inst->a].u & ~LTAG_MASK);

		ltable_set(tbl, regset->regs[inst->b], regset->regs[inst->c]);
	}
	DISPATCH();

do_getfield:;
	{
		if(lvalue_tag(regset->regs[inst->a]) != LTAG_TABLE) {
			goto err;
		}

		LTable *tbl = (void*) (regset->regs[inst->b].u & ~LTAG_MASK);

		regset->regs[inst->a] = ltable_get(tbl, regset->regs[inst->c]);
	}
	DISPATCH();

do_len:
	if(lvalue_tag(regset->regs[inst->b]) == LTAG_STRING) {
		regset->regs[inst->a] = lvalue_from_int32(((LString*) (regset->regs[inst->b].u & ~LTAG_MASK))->length);
	} else if(lvalue_tag(regset->regs[inst->b]) == LTAG_TABLE) {
		regset->regs[inst->a] = lvalue_from_int32(ltable_len((LTable*) (regset->regs[inst->b].u & ~LTAG_MASK)));
	} else goto err;
	DISPATCH();
	
err:;
	puts("Error");
do_ret:;
	return 0;
	
}

size_t lvm_run(LVM *L, LFunc *func, size_t arg_count, LRegSet *regset) {
	set_LValueU heap = {};
	
	atomic_fetch_add(&L->active_thread_count, 1);
	
	size_t ret = lvm_call(L, func, arg_count, &heap, regset);
	
	mtx_lock(&L->dead_heap_mut);
		for(c_each(i, set_LValueU, heap)) {
			set_LValueU_insert(&L->dead_heap, *i.ref);
		}
		atomic_fetch_sub(&L->active_thread_count, 1);
	mtx_unlock(&L->dead_heap_mut);

	set_LValueU_drop(&heap);
	
	return ret;
}

LFunc *lvm_func_from_native(LFuncCallback cb, void *ud) {
	LFunc *f = calloc(1, sizeof(*f));
	f->is_native = true;
	f->ud = ud;
	f->native_func = cb;
	return f;
}

bool lvalue_eq(LValue a, LValue b) {
	if(a.u == b.u) {
		return true;
	}
	
	if(lvalue_tag(a) == LTAG_I32 && lvalue_tag(b) == LTAG_FLOAT) {
		return (a.u & ~LTAG_MASK) == b.f;
	} else if(lvalue_tag(a) == LTAG_FLOAT && lvalue_tag(b) == LTAG_I32) {
		return (b.u & ~LTAG_MASK) == a.f;
	} else if(lvalue_tag(a) == LTAG_STRING && lvalue_tag(b) == LTAG_STRING) {
		LString *sa = (LString*) (a.u & ~LTAG_MASK);
		LString *sb = (LString*) (b.u & ~LTAG_MASK);
		
		if(sa->length != sb->length) {
			return false;
		}
		
		return !memcmp(sa->data, sb->data, sa->length);
	}
	
	return false;
}

static void gc_unmark_heap(set_LValueU *heap) {
	for(c_each(i, set_LValueU, *heap)) {
		LValue v = (LValue) {.u = *i.ref};

		assert(lvalue_tag(v) == LTAG_TABLE || lvalue_tag(v) == LTAG_STRING);

		void *gco = (void*) (v.u & ~LTAG_MASK);

		if(lvalue_tag(v) == LTAG_TABLE) {
			LTable *tbl = gco;
			tbl->ref = false;
		} else if(lvalue_tag(v) == LTAG_STRING) {
			LString *str = gco;
			str->ref = false;
		}
	}
}
static void gc_unmark_all(LVM *L, size_t thread_count) {
	for(size_t thrd = 0; thrd < thread_count; thrd++) {
		LThreadPrivates *privates = &L->privates[thrd];
		gc_unmark_heap(privates->heap);
	}
	gc_unmark_heap(&L->dead_heap);
}
static void gc_mark(LValue v) {
	if(lvalue_tag(v) != LTAG_TABLE && lvalue_tag(v) != LTAG_STRING) {
		return;
	}

	void *gco = (void*) (v.u & ~LTAG_MASK);
	
	if(lvalue_tag(v) == LTAG_TABLE) {
		LTable *tbl = gco;
		tbl->ref = true;

		for(size_t i = 0; i < tbl->buckets->capacity; i++) {
			LTableEntry e = tbl->buckets->data[i];
			gc_mark(e.key);
			gc_mark(e.val);
		}
	} else if(lvalue_tag(v) == LTAG_STRING) {
		LString *str = gco;
		str->ref = true;
	}
}
static void gc_mark_units(LVM *L) {
	for(size_t u = 0; u < L->unit_count; u++) {
		LUnit *unit = &L->units[u];
		for(size_t f = 0; f < unit->func_count; f++) {
			LFunc *func = &unit->funcs[f];
			gc_mark(lvalue_from_table(func->env));
			for(size_t upv = 0; upv < func->upvalue_count; upv++) {
				gc_mark(func->upvalues[upv]);
			}
		}
	}
}
static void safepoint_active(LVM *L, set_LValueU *heap, LRegSet *regset) {
	size_t my_privates_index = atomic_fetch_add(&L->privates_index, 1);
	
	L->privates[my_privates_index].heap = heap;
	L->privates[my_privates_index].regset = regset;
	
	atomic_fetch_add(&L->privates_ready, 1);
	
	// Wait until GC finishes
	while(atomic_load(&L->safepoint_func) == safepoint_active);
	
	atomic_fetch_sub(&L->privates_ready, 1);
}
static void gc_mark_stacks(LVM *L, size_t thread_count) {
	for(size_t thrd = 0; thrd < thread_count; thrd++) {
		LThreadPrivates *privates = &L->privates[thrd];
		
		LRegSet *rset = privates->regset;
		while(rset) {
			for(size_t r = 0; r < 256; r++) {
				gc_mark(rset->regs[r]);
			}
			rset = rset->parent;
		}
	}
}
static void safepoint_inactive(LVM *L, set_LValueU *heap, LRegSet *regset) {
}
static void gc_delete_unmarked_in_heap(LVM *L, set_LValueU *heap) {
	for(set_LValueU_iter i = set_LValueU_begin(heap); i.ref;) {
		LValue v = (LValue) {.u = *i.ref};

		void *gco = (void*) (v.u & ~LTAG_MASK);

		if(lvalue_tag(v) == LTAG_TABLE) {
			LTable *tbl = gco;
			
			if(tbl->ref == false) {
				free(tbl->buckets);
				free(tbl);
				
				i = set_LValueU_erase_at(heap, i);
				continue;
			}
		} else if(lvalue_tag(v) == LTAG_STRING) {
			LString *str = gco;

			if(str->ref == false) {
				lvm_gc_alert(L, NULL, -sizeof(*str) - str->length);
				
				free(str);
				
				i = set_LValueU_erase_at(heap, i);
				continue;
			}
		}

		set_LValueU_next(&i);
	}
}
static void gc_delete_unmarked(LVM *L, size_t thread_count) {
	for(size_t thrd = 0; thrd < thread_count; thrd++) {
		LThreadPrivates *privates = &L->privates[thrd];
		gc_delete_unmarked_in_heap(L, privates->heap);
	}
	gc_delete_unmarked_in_heap(L, &L->dead_heap);
}
static void lvm_gc_force(LVM *L, LThreadPrivates *callerPrivates) {
	// At most one thread can force GC, while others must behave as usual and enter a safepoint instead
	if(atomic_compare_exchange_strong(&L->gcInProgress, &(bool) {false}, true)) {
		
		//static size_t gcidx = 0;
		//fprintf(stderr, "GC %i (%lu bytes)\n", atomic_fetch_add(&gcidx, 1), L->memUsage);
		
		if(callerPrivates) {
			// Called from within VM
			atomic_store(&L->privates_index, 1);
			atomic_store(&L->privates_ready, 1);
			
			L->privates[0] = *callerPrivates;
		} else {
			// Called outside of VM, probably by lvm_destroy
			atomic_store(&L->privates_index, 0);
			atomic_store(&L->privates_ready, 0);
		}
		
		L->safepoint_func = safepoint_active;
		
		// Wait until other threads have entered GC stage
		while(atomic_load(&L->privates_ready) < atomic_load(&L->active_thread_count));
		
		size_t thread_count = atomic_load(&L->privates_ready);
		
		mtx_lock(&L->dead_heap_mut);
			gc_unmark_all(L, thread_count);
			gc_mark_stacks(L, thread_count);
			gc_mark_units(L);
			gc_delete_unmarked(L, thread_count);
		mtx_unlock(&L->dead_heap_mut);
		
		while(L->memUsage > L->nextGCThreshold) {
			L->nextGCThreshold <<= 1;
		}
		
		L->safepoint_func = safepoint_inactive;
		
		if(callerPrivates) {
			// Called from within VM
			atomic_fetch_sub(&L->privates_ready, 1);
		}
		
		// Wait until other threads have left GC stage
		while(atomic_load(&L->privates_ready) > 0);
		
		atomic_store(&L->gcInProgress, false);
		
	}
}

// `privates` can be NULL but ONLY IF diff < 0
void lvm_gc_alert(LVM *L, LThreadPrivates *privates, intmax_t diff) {
	L->memUsage += diff;
	
	assert(L->memUsage >= 0);
	
	if(L->memUsage > L->nextGCThreshold) {
		lvm_gc_force(L, privates);
	}
}

void lvm_init(LVM *L) {
	memset(L, 0, sizeof(*L));
	
	L->safepoint_func = safepoint_inactive;
	L->nextGCThreshold = 16384;
	mtx_init(&L->dead_heap_mut, mtx_plain);
}

void lvm_destroy(LVM *L) {
	mtx_destroy(&L->dead_heap_mut);
	
	lvm_gc_force(L, NULL);
	
	set_LValueU_drop(&L->dead_heap);
}