#include "ms.h"
#include <list>
#include <map>
#include <cstring>
#include <sstream>
#include <cstdlib>

/*
 * Stack frame layout:
 * 0 |  oldpc
 * 1 |  oldasm
 * 2 |  oldenv
 * 3 |  oldbp // for debug 
 * 4 V  ...
 */

namespace ms {

static Value *new_lambda(Assembly *assm, Value *parent_env) {
	Value *v = unshield().new_object(Value::LAMBDA);
	v->p1 = assm;
	v->p2 = parent_env;
	return v;
}

static Value *new_environment(Value *parent, size_t vsize) {
	std::vector<Value *> *vec = new std::vector<Value *>(vsize);
	std::fill(vec->begin(), vec->end(), (Value *)0);

	Value *v = get_memory()->alloc(Value::ENVIRON);
	v->p1 = parent;
	v->p2 = vec;
	return v;
}

void environ_mark(Memory &mem, const Value *v) {
	if (v->p1)
		mem.mark_object(static_cast<Value *>(v->p1));
	const std::vector<Value *> *vec = static_cast<const std::vector<Value *> *>(v->p2);
	for (std::vector<Value *>::const_iterator i = vec->begin(); i != vec->end(); i++) {
		if (*i)
			mem.mark_object(*i);
	}
}

void environ_free(Value *v) {
	delete static_cast<std::vector<Value *> *>(v->p2);
}

void asm_mark(Memory &mem, const Assembly *v) {
	for (std::vector<Instruction>::const_iterator i = v->seq()->begin();
	     i != v->seq()->end(); i++) {
		switch (i->type) {
		case Instruction::PUSH:
			mem.mark_object(reinterpret_cast<const Value *>(i->arg));
			break;
		case Instruction::LAMBDA:
			mem.mark_object(reinterpret_cast<const Assembly *>(i->arg));
			break;
		default:
			break;
		}
	}
}

void asm_free(Assembly *v) {
	delete v->seq();
}

void Continuation::mark(Memory &mem) const {
	size_t stop = sp();
	while (stop--)
		mem.mark_object(copied_stack()[stop]);
}

void Continuation::free() {
	std::free(p1); // copied_stack
}

static Continuation *new_continuation(Value **copied_stack, size_t sp) {
	Value *v = unshield().new_object(Value::CONT);
	v->p1 = copied_stack;
	v->p2 = reinterpret_cast<void *>(sp);
	return static_cast<Continuation *>(v);
}

static bool check_argc(int expected, int given) {
	return expected == given || expected < 0 && given >= -expected - 1;
}

static Value *argument_error(int expected, int given) {
	std::ostringstream sstream;
	sstream << "wrong number of arguments: expected "
		<< expected << ", given " << given;
	return unshield().new_error(sstream.str());
}


class VM {
private:
#define stack_push(v) do { \
	Value *vv = (v); \
	if (sp == stack_depth) \
		return unshield().new_pair(unshield().new_nil(), unshield().new_error("stack overflow")); \
	stack[sp++] = vv; \
} while (0)
#define stack_pop() \
	stack[--sp]
#define stack_popn(n) \
	(sp -= n)

	static inline std::vector<Value *> *unwrap_env(Value *v) {
		return static_cast<std::vector<Value *> *>(v->p2);
	}

	static inline Value **resolve_var(Value *env, size_t idx) {
		while (idx >= unwrap_env(env)->size()) {
			idx -= unwrap_env(env)->size();
			env = static_cast<Value *>(env->p1);
		}
		return &unwrap_env(env)->at(idx);
	}

public:
	VM(const Assembly *assm, Value *parent_env, int stack_depth) :
		stack_depth(stack_depth),
		toplevel_assm(assm), env(parent_env), sp(0) {
		stack = (Value **)std::malloc(sizeof(Value *) * stack_depth);
		assert(stack);
	}

	~VM() {
		std::free(stack);
	}

	Value *execute() {
		size_t pc = 0, bp = 0;
		const Assembly *assm = toplevel_assm;
		env = new_environment(env, assm->num_variables());

#ifdef THREADED_CODE
# error FIXME
#else
# define CASE(n) case Instruction::n:
# define NEXT break
#endif

		// FIXME: threaded
		while (pc < assm->seq()->size()) {
			const Instruction *i = &assm->seq()->at(pc++);
			//std::cout << "ins:" << pc-1 << ",itype:" << i->type << "\n";
			switch (i->type) {
			CASE(NOP)
				NEXT;
			CASE(PUSH)
				stack_push((Value *)i->arg);
				NEXT;
			CASE(POP)
				(void)stack_pop();
				NEXT;
			CASE(DUP) {
				Value *v = stack_pop();
				stack_push(v);
				stack_push(v);
				NEXT;
			}
			CASE(GETVAR)
				stack_push(*resolve_var(env, i->arg));
				NEXT;
			CASE(SETVAR)
				*resolve_var(env, i->arg) = stack_pop();
				NEXT;
			CASE(LAMBDA)
				stack_push(new_lambda(reinterpret_cast<Assembly *>(i->arg), env));
				NEXT;
			CASE(APPLY)
			CASE(APPLY_TAILCALL)
apply_again: {
				Memory::Shield shield(get_memory());
				int argc = i->arg;
				Value **args = &stack[sp - argc];
				Value *callable = stack[sp - argc - 1];
				// std::cout << "CALLING:";
				// dump(callable);
				// std::cout << "\nARGS:" << i->arg << "\n";
				// for (int ik = 0; ik < i->arg; ik++)
				// 	dump(args[ik]);
				// std::cout << "\nEND\n";
				// std::cout << "bp="<<bp<<",sp="<<sp<<"\n";
				if (callable->type == Value::LAMBDA) {
					shield.add(callable);

					const Assembly *casm = static_cast<const Assembly *>(callable->p1);
					if (!check_argc(casm->num_args(), argc))
						return shield.new_pair(shield.new_nil(),
								       argument_error(casm->num_args(), argc));

					if (casm->num_args() < 0) {
						args[-casm->num_args() - 1] =
							shield.new_list(argc - (-casm->num_args() - 1),
								 &args[-casm->num_args() - 1]);
						argc = -casm->num_args();
					}
					Value *new_env = shield.add(new_environment(static_cast<Value *>(callable->p2),
										    casm->num_variables()));
					if (argc > 0)
						std::memcpy(&unwrap_env(new_env)->at(0),
							    args, argc * sizeof(Value *));
					stack_popn(i->arg + 1);

					if (bp == 0 || i->type != Instruction::APPLY_TAILCALL) {
						// Non-tail call && non-toplevel call
						stack_push(shield.new_number(pc));
						stack_push(const_cast<Assembly *>(assm));
						stack_push(env);
						stack_push(shield.new_number(bp));
					}
					pc = 0;
					assm = casm;
					env = new_env;
					bp = sp;
				} else if (callable->type == Value::CFUNC) {
					int expected_argc = (intptr_t)callable->p2;
					if (!check_argc(expected_argc, argc))
						return shield.new_pair(shield.new_nil(), argument_error(expected_argc, argc));

					Value *(*func)(int, Value **) =
						reinterpret_cast<Value *(*)(int, Value **)>(callable->p1);
					Value *ret = func(i->arg, args);
					shield.add(ret);
					if (ret->is_a(Value::ERROR))
						return shield.new_pair(shield.new_nil(), ret);
					stack_popn(argc + 1);
					stack_push(ret);
				} else if (callable->type == Value::CALLCC) {
					if (!check_argc(1, argc))
						return shield.new_pair(shield.new_nil(), argument_error(1, argc));

					Value *called = shield.add(stack_pop()); // arg for call/cc
					(void)shield.add(stack_pop());  // call/cc

					Value **copied_stack = (Value **)std::malloc(sizeof(Value *) * stack_depth);
					assert(copied_stack);
					std::memcpy(copied_stack, stack, stack_depth * sizeof(Value *));
					copied_stack[sp + 0] = shield.new_number(pc);
					copied_stack[sp + 1] = const_cast<Assembly *>(assm);
					copied_stack[sp + 2] = env;
					copied_stack[sp + 3] = shield.new_number(bp);
					// std::cout << "cont saved! bp=" << bp << ",sp=" << sp << ": \n";
					Value *cont = shield.add(new_continuation(copied_stack, sp + 4));

					stack_push(called);
					stack_push(cont);
					goto apply_again;
				} else if (callable->type == Value::CONT) {
					if (!check_argc(1, argc))
						return shield.new_pair(shield.new_nil(), argument_error(1, argc));
					// std::cout << "cont! bp=" << bp << ",sp=" << sp << ": \n";
					Value *arg = shield.add(stack_pop()); // arg for cont

					const Continuation *cont = static_cast<Continuation *>(callable);
					std::memcpy(stack, cont->copied_stack(), stack_depth * sizeof(Value *));
					sp = cont->sp();

					bp = reinterpret_cast<size_t>(stack_pop()->p1);
					env = stack_pop();
					assm = static_cast<const Assembly *>(stack_pop());
					pc = reinterpret_cast<size_t>(stack_pop()->p1);

					// std::cout << "cont restored! bp=" << bp << ",sp=" << sp << ": \n";
					stack_push(arg);
				} else
					return shield.new_pair(shield.new_nil(), shield.new_error("invalid application"));
				NEXT;
			}
			CASE(RET) {
				// std::cout << "ret! bp=" << bp << ",sp=" << sp << ": ";
				Value *ret = stack_pop();
				// dump(ret);
				// std::cout << "\n";
				assert(bp == sp);
				if (sp != 0) {
					bp = reinterpret_cast<size_t>(stack_pop()->p1);
					env = stack_pop();
					assm = static_cast<const Assembly *>(stack_pop());
					pc = reinterpret_cast<size_t>(stack_pop()->p1);
				}

				stack_push(ret);
				NEXT;
			}
			CASE(JUMP)
				pc += i->arg;
				NEXT;
			CASE(JUMPUNLESS)
				if (!stack_pop()->test())
					pc += i->arg;
				NEXT;
			CASE(JUMPIF)
				if (stack_pop()->test())
					pc += i->arg;
				NEXT;
			}
		}
		// Toplevel
		assert(bp == 0 && sp == 1);
		{
			Memory::Shield shield(get_memory());
			return shield.new_pair(shield.add(stack_pop()), shield.new_nil());
		}
	}

	size_t stack_depth;
	Value **stack;
	const Assembly *toplevel_assm;
	Value *env;
	size_t sp;
};

void vm_mark_objects(Memory &mem, const VM *vm) {
	if (!vm)
		return;
	size_t stop = vm->sp;
	while (stop--)
		mem.mark_object(vm->stack[stop]);
	if (vm->env)
		mem.mark_object(vm->env);
	mem.mark_object(vm->toplevel_assm);
}

Value *eval(const Assembly *assm, Value *env) {
	Memory::Shield shield(get_memory());
	Value *wrapper = shield.new_object(Value::VM);
	wrapper->p1 = wrapper->p2 = nullptr;

	VM vm(assm, env, 1024);
	wrapper->p1 = &vm;

	Value *ret = vm.execute();
	wrapper->p1 = nullptr;
	return ret;
}

}