/* -*- mode:c -*- */

#include "compiler.h"
#include "gc.h"
#include "lib/std.h"
#include "lisp.h"
#include "plat/plat.h"

#include <dasm_proto.h>
#include <dasm_x86.h>

#include <libgen.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define value_size sizeof(value_t)

|.arch x86;

|.macro setup, nvars;
|1:;
| push ebp;
| mov ebp, esp;
| sub esp, (value_size * nvars);
|.endmacro;

|.macro cleanup;
| mov esp, ebp;
| pop ebp;
| ret;
|.endmacro;

|.macro call_extern, address;
| call &address;
|.endmacro;

dasm_State *d;
unsigned int npc = 8;

|.macro run_gc;
| mov eax, esp;
| push ebp;
| push eax;
| mov eax, _do_gc;
| call eax;
|.endmacro;

struct function *find_function(struct environment *env, char *name)
{
	struct function *f;

	for (f = env->first; f && strcmp(f->name, name); f = f->prev)
	{
	}

	return f;
}

unsigned int local_alloc(struct local *local)
{
	for (int i = 0; i < local->num_stack_slots; i++)
	{
		if (local->stack_slots[i] == false)
		{
			local->stack_slots[i] = true;

			if (i >= local->num_stack_entries)
				local->num_stack_entries++;

			return i;
		}
	}

	int old_size = local->num_stack_slots;
	local->num_stack_slots += 4;
	local->stack_slots =
	    realloc(local->stack_slots, local->num_stack_slots * sizeof(bool));
	// unreadable: set the remaining slots to unused
	memset(local->stack_slots + old_size, 0, local->num_stack_slots - old_size);
	local->stack_slots[old_size] = true;

	return old_size;
}

void local_free(struct local *local, unsigned int slot)
{
	local->stack_slots[slot] = false;
}

void del_local(struct local *local)
{
	free(local->stack_slots);

	for (struct variable *next, *f = local->first; f; f = next)
	{
		next = f->prev;
		free(f);
	}
}

void del_env(struct environment *env)
{
	for (struct function *next, *f = env->first; f; f = next)
	{
		next = f->prev;
		// We're not gonna bother munmap()ing the function
		free(f);
	}

	for (struct loaded_file *next, *l = env->first_loaded; l; l = next)
	{
		next = l->previous;
		free(l->resolved_path);
		free(l);
	}

	free(env);
}

void add_load(struct environment *env, char *path)
{
	static char buffer[512];
	long size = readlink(path, buffer, 512);
	buffer[size] = '\0';
	char *resolved = strdup(buffer);

	struct loaded_file *f = malloc(sizeof(struct loaded_file));
	f->resolved_path = resolved;
	f->previous = env->first_loaded;
	env->first_loaded = f;
}

struct dasm_State *compile_function(value_t args, enum namespace namespace,
                                    struct environment *env,
                                    struct local *local_out,
                                    struct local *local_parent,
                                    struct args **args_out, char *name,
                                    char *path)
{
	dasm_State *d;
	dasm_State **Dst = &d;

	|.section code, imports;
	dasm_init(&d, DASM_MAXSECTION);

	|.globals lbl_;
	void *labels[lbl__MAX];
	dasm_setupglobal(&d, labels, lbl__MAX);

	|.actionlist lisp_actions;
	dasm_setup(&d, lisp_actions);

	struct local local;
	local.parent = NULL;
	local.first = NULL;
	local.num_vars = 0;
	local.npc = 8;
	local.nextpc = 0;
	local.stack_slots = malloc(sizeof(bool) * 4);
	memset(local.stack_slots, 0, sizeof(bool) * 4);
	local.num_stack_slots = 4;
	local.num_stack_entries = 0;
	local.num_closure_slots = 0;
	local.parent = local_parent;
	local.current_function_name = name;
	local.current_file_path = path;

	dasm_growpc(&d, local.npc);

	value_t arglist = car(args);
	value_t body = cdr(args);

	// This will add the arguments to local too.
	struct args *ar = list_to_args(env, arglist, &local);
	local.args = ar;

	if (!ar)
	{
		err("Malformed args list");
	}

	for (value_t body_ = body; !nilp(body_); body_ = cdr(body_))
	{
		walk_and_alloc(&local, car(body_));
	}

	| setup (local.num_stack_entries);

	memset(local.stack_slots, 0, local.num_stack_slots * sizeof(bool));
	local.num_stack_entries = 0;

	for (; !nilp(body); body = cdr(body))
	{
		bool tail = nilp(cdr(body));
		compile_expression(env, &local, car(body), tail, Dst);
	}

	| cleanup;

	if (local_out)
		*local_out = local;

	if (args_out)
		*args_out = ar;

	return d;
}

void compile_tl(value_t val, struct environment *env, char *fname)
{
	if (!listp(val))
		err("Top level must be a list");

	value_t form = car(val);
	value_t args = cdr(val);

	if (symstreq(form, "defun") || symstreq(form, "defmacro"))
	{
		enum namespace namespace = NS_FUNCTION;

		if (symstreq(form, "defmacro"))
		namespace = NS_MACRO;

		struct local local;
		struct args *a;
		char *name = (char *)(car(args) ^ SYMBOL_TAG);

		dasm_State *d = compile_function(cdr(args), namespace, env, &local,
		                                 NULL, &a, name, fname);

		add_function(env, name, link_program(&d), a, namespace);

		dasm_free(&d);
		del_local(&local);
	}
	else if (symstreq(form, "progn"))
	{
		for (value_t val = args; !nilp(val); val = cdr(val))
		{
			compile_tl(car(val), env, fname);
		}
	}
	else if (symstreq(form, "load"))
	{
		if (length(args) != 1)
		{
			err_at(val, "load expects exactly 1 argument, %d given",
				   length(args));
		}
		load_relative(env, fname, car(args));
	}
}

void walk_and_alloc(struct local *local, value_t body)
{
	// TODO: handle macros
	if (!listp(body))
		return;

	value_t args = cdr(body);

	if (symstreq(car(body), "let1"))
	{
		int slot = local_alloc(local);

		value_t expr = cdr(args);
		for (; !nilp(expr); expr = cdr(expr))
		{
			walk_and_alloc(local, car(expr));
		}

		local_free(local, slot);
	}
	else if (symstreq(car(body), "lambda"))
	{
		// We don't want to walk the lambda because it's another function. When
		// the lambda is compiled it will be walked.
		return;
	}
	else
	{
		for (; !nilp(args); args = cdr(args))
		{
			walk_and_alloc(local, car(args));
		}
	}
}

bool load(struct environment *env, char *path)
{
	if (!file_exists(path))
		return false;

	add_load(env, path);

	unsigned char pool = make_pool();
	unsigned char pop = push_pool(pool);

	struct istream *is = new_fistream(path, false);
	if (!is)
		return false;

	value_t val;

	while (read1(is, &val))
	{
		compile_tl(val, env, path);
	}

	del_fistream(is);
	pop_pool(pop);

	return true;
}

value_t load_relative(struct environment *env, char *to, value_t name)
{
	if (!stringp(name))
		return nil;

	char *new_path = (char *)(name ^ STRING_TAG);
	char *relative_to = strdup(to);
	char full_path[512];

	snprintf(full_path, 512, "%s/%s", dirname(relative_to), new_path);

	if (load(env, full_path))
		return t;
	else
		return nil;
}

struct environment *compile_file(char *filename, bool *ok)
{
	value_t val;
	struct environment *env = malloc(sizeof(struct environment));
	env->first = NULL;
	env->first_loaded = NULL;

	add_load(env, filename);
	load_std(env);

	bool ok_ = load(env, filename);

	if (ok)
		*ok = ok_;

	return env;
}

int nextpc(struct local *local, dasm_State **Dst)
{
	int n = local->nextpc++;
	if (n > local->npc)
	{
		local->npc += 16;
		dasm_growpc(Dst, local->npc);
	}
	return n;
}

void compile_backquote(struct environment *env, struct local *local,
                       value_t val, dasm_State **Dst)
{
	if (!listp(val))
	{
		| mov eax, (val);
	}
	else
	{
		value_t fsym = car(val), args = cdr(val);
		int nargs = length(args),
			n = length(val);

		if (symstreq(fsym, "unquote"))
		{
			if (nargs != 1)
			{
				err_at(val, "unquote (or ,) takes exactly 1 argument");
			}

			compile_expression(env, local, car(args), false, Dst);
		}
		else
		{
			| push nil;

			for (int i = n - 1; i >= 0; i--)
			{
				compile_backquote(env, local, elt(val, i), Dst);
				| push eax;
				| call_extern cons;
				| add esp, 8;

				// Remove unnecessary pop
				| push eax;
			}
				| pop eax;
		}
	}
}

value_t eval(struct environment *env, value_t form)
{
	// Eval!
	value_t function = cons(nil, cons(form, nil));

	struct local local;
	struct args *args;

	dasm_State *d = compile_function(function, NS_ANONYMOUS, env, &local, NULL,
	                                 &args, NULL, "/");

	del_local(&local);

	value_t (*f)() = link_program(&d);
	return f();
}

void compile_variable(struct variable *v, dasm_State *Dst)
{
	switch (v->type)
	{
	case V_ARGUMENT:
		| mov eax, dword[ebp + (value_size * (v->number + 2))];
		break;
	case V_BOUND:
		| mov eax, dword[ebp - ((v->number + 1) * value_size)];
		break;
	case V_FREE:
		// edi is the closure context pointer
		| mov eax, dword[edi + (v->number * value_size)];
		break;
	default:
		err("Sorry, can only access V_ARGUMENT, V_FREE and V_BOUND variables "
		    "for now :(");
	}
}

void compile_expression(struct environment *env, struct local *local,
                        value_t val, bool tail, dasm_State **Dst)
{
	if (symstreq(val, "nil") || nilp(val))
	{
		| mov eax, (nil);
	}
	else if (symstreq(val, "t"))
	{
		| mov eax, (t);
	}
	else if (integerp(val) || stringp(val))
	{
		| mov eax, val;
	}
	else if (listp(val))
	{
		value_t fsym = car(val);
		value_t args = cdr(val);
		int nargs = length(args);

		if (!symbolp(fsym))
		{
			printval(val, 2);
			err_at(val, "function name must be a symbol");
		}

		if (symstreq(fsym, "if"))
		{
			if (nargs < 2 || nargs > 3)
				err("Must give at least 2 arguments to if");

			compile_expression(env, local, car(args), false, Dst);
			int false_label = nextpc(local, Dst),
			    after_label = nextpc(local, Dst);

			// result is in eax
			| cmp eax, (nil);
			| je =>false_label;

			compile_expression(env, local, elt(args, 1), tail, Dst);
			| jmp =>after_label;
			|=>false_label:;
			if (nargs == 3)
				compile_expression(env, local, elt(args, 2), tail, Dst);
			|=>after_label:;
		}
		else if (symstreq(fsym, "progn"))
		{
			for (value_t val = args; !nilp(val); val = cdr(val))
			{
				bool t = tail && nilp(cdr(val));
				compile_expression(env, local, car(val), t, Dst);
			}
		}
		else if (symstreq(fsym, "let1"))
		{
			if (nargs < 2)
			{
				err("Must give at least 2 arguments to let1");
			}
			value_t binding = car(args);
			value_t rest = cdr(args);

			if (length(binding) != 2)
			{
				err("Binding list in let1 must contain exactly two entries");
			}

			value_t name = car(binding);
			value_t value = car(cdr(binding));

			compile_expression(env, local, value, false, Dst);

			int i = local_alloc(local);

			add_variable(local, V_BOUND, (char *)(name ^ SYMBOL_TAG), i);

			| mov dword[ebp - ((i + 1) * value_size)], eax;

			for (; !nilp(rest); rest = cdr(rest))
			{
				bool t = tail && nilp(cdr(rest));
				compile_expression(env, local, car(rest), t, Dst);
			}

			local_free(local, i);
		}
		else if (symstreq(fsym, "gc"))
		{
			if (nargs)
			{
				err_at(val, "gc takes no arguments");
			}

			| run_gc;
		}
		else if (symstreq(fsym, "quote"))
		{
			if (nargs != 1)
				err("quote should take exactly 1 argument");

			// Simple!
			| mov eax, (car(args));
		}
		else if (symstreq(fsym, "backquote"))
		{
			if (nargs != 1)
				err("backquote should take exactly 1 argument");

			compile_backquote(env, local, car(args), Dst);
		}
		else if (symstreq(fsym, "function"))
		{
			if (nargs != 1)
			{
				err("function should take exactly 1 argument");
			}

			if (!symbolp(car(args)))
			{
				err("argument to function should be a symbol resolvable at "
				    "compile time");
			}

			struct function *f =
			    find_function(env, (char *)(car(args) ^ SYMBOL_TAG));
			value_t closure = create_closure(f->code_ptr, f->args, 0);

			| mov eax, (closure);
		}
		else if (symstreq(fsym, "list"))
		{
			| push (nil);

			for (int i = nargs - 1; i >= 0; i--)
			{
				compile_expression(env, local, elt(args, i), false, Dst);

				| push eax;
				| call_extern cons;
				| add esp, (2 * value_size);
				| push eax;
			}
				| pop eax;
		}
		else if (symstreq(fsym, "lambda"))
		{
			// Compile the function with this as the parent scope
			struct local new_local;
			int nargs_out;
			dasm_State *d = compile_function(
			    args, NS_ANONYMOUS, env, &new_local, local, &nargs_out,
			    "recurse", local->current_file_path);

			// Link the function
			void *func_ptr = link_program(&d);

			// Create a closure object with the correct number of captures at
			// runtime
			| push (new_local.num_closure_slots);
			| push (nargs_out);
			| push (func_ptr);
			| call_extern create_closure;
			| add esp, 12;

			// Walk the generated local scope for V_FREE variables, since each
			// of these exists in our scope (or higher), evaluate it and set it
			// as a member of the lambda capture.

			for (struct variable *var = new_local.first; var; var = var->prev)
			{
				if (var->type == V_FREE)
				{
					// Closure in eax
					| push eax;
					// Variable now in eax
					compile_variable(find_variable(local, var->name), Dst);
					| push eax;

					// The capture offset
					| push (var->number);
					| call_extern set_closure_capture_variable;
					// Skip the value and index
					| add esp, 8;
					// Pop the closure back in to eax
					| pop eax;
				}
			}

			// Closure is still in eax

			dasm_free(&d);
			del_local(&new_local);
		}
		else if (symstreq(fsym, "eval"))
		{
			if (nargs != 1)
			{
				err("eval takes exactly 1 argument");
			}

			compile_expression(env, local, car(args), false, Dst);
			| push eax;
			| push (env);
			| call_extern eval;
		}
		else if (symstreq(fsym, "load"))
		{
			if (nargs != 1)
			{
				err_at(val, "load takes exactly 1 argument, %d given", nargs);
			}

			compile_expression(env, local, car(args), false, Dst);
			| push eax;
			| push (local->current_file_path);
			| push (env);
			| call_extern load_relative;
		}
		else
		{
			char *name = (char *)(fsym ^ SYMBOL_TAG);
			struct function *func = find_function(env, name);

			bool is_recursive = false;
			struct args *nargs_needed = NULL;

			// The number of arguments actually passed on the stack,
			// i.e. all varargs are 1.
			int real_nargs;

			if (local->current_function_name &&
			    symstreq(fsym, local->current_function_name))
			{
				is_recursive = true;
				nargs_needed = local->args;
			}
			else
			{
				if (func == NULL)
				{
					err_at(val, "Function %s undefined", name);
				}

				nargs_needed = func->args;
			}

			if (!are_args_acceptable(nargs_needed, nargs))
			{
				err_at(val,
				       "wrong number of args in function call: %s at %s:%d, "
				       "want %d args but given %d\n",
				       name, cons_file(val), cons_line(val),
				       nargs_needed->num_required, nargs);
			}

			int total_taken = nargs_needed->num_optional +
				nargs_needed->num_required;

			real_nargs = total_taken + (nargs_needed->variadic ? 1 : 0);

			if (is_recursive || func->namespace == NS_FUNCTION)
			{
				int nargs = length(args);

				int line = cons_line(val);
				char *file = cons_file(val);

				if (nargs_needed->variadic)
				{
					| push (nil);
				}

				if (nargs > total_taken && nargs_needed->variadic)
				{
					// We are passing varargs, which means we need to make a list

					for (int i = nargs - 1; i >= total_taken; i--)
					{
						compile_expression(env, local, elt(args, i), false, Dst);
						| push eax;
						| call_extern cons;
						| add esp, 8;
						| push eax;
					}
				}

				for (int i = nargs_needed->num_optional - 1;
				     i >= nargs - nargs_needed->num_required; i--)
				{
					// Push the default optional values
					| push (nargs_needed->optional_arguments[i].value);
				}

				int min = MIN(nargs, total_taken);

				for (int i = min - 1; i >= 0; i--)
				{
					compile_expression(env, local, elt(args, i), false, Dst);
					| push eax;
				}

				if (is_recursive)
				{
					if (tail)
					{
						// Move all the arguments pushed to the stack
						// back up to the argument bit of the stack.

						for (int i = 0; i < real_nargs; i++)
						{
							| pop eax;
							| mov dword[ebp + (value_size * (i + 2))], eax;
						}

						// Jmp back to start
						| mov esp, ebp;
						| pop ebp;
						| jmp <1;
					}
					else
					{
						| call <1;
					}
				}
				else
				{
					// | mov ebx, (func->code_addr);
					| call_extern func->code_addr;
				}
				| add esp, (real_nargs * value_size);
				// result in eax
			}
			else if (func->namespace == NS_MACRO)
			{
				// Make sure that the stuff allocated by the macro isn't in a
				// pool
				unsigned char pool = push_pool(0);

				value_t expanded_to = call_list(func, args);

				pop_pool(pool);

				compile_expression(env, local, expanded_to, false, Dst);
			}
		}
	}
	else if (symbolp(val))
	{
		if (symstreq(val, "+current-file+"))
		{
			value_t file_name_val = strval(local->current_file_path);

			| mov eax, (file_name_val);
		}
		else
		{
			struct variable *v =
			    find_variable(local, (char *)(val ^ SYMBOL_TAG));

			if (!v)
			{
				fprintf(stderr, "var: %s\n", (char *)(val ^ SYMBOL_TAG));
				err("Variable unbound");
			}

			compile_variable(v, Dst);
		}
	}
}

struct variable *add_variable(struct local *local, enum var_type type,
                              char *name, int number)
{
	struct variable *var = malloc(sizeof(struct variable));
	var->prev = local->first;
	var->type = type;
	var->name = name;
	var->number = number;

	local->first = var;

	return var;
}

void destroy_local(struct local *local)
{
	for (struct variable *v = local->first; v;)
	{
		struct variable *t = v;
		v = v->prev;
		free(t);
	}
}

struct variable *find_variable(struct local *local, char *name)
{
	struct variable *v = local->first;

	for (; v && strcmp(v->name, name) != 0; v = v->prev)
	{
	}

	if (!v)
	{
		if (local->parent)
		{
			v = find_variable(local->parent, name);

			if (v)
			{
				// We found this in a parent scope, add it as a V_FREE variable
				// to skip the search.
				v = add_variable(local, V_FREE, name,
				                 local->num_closure_slots++);
			}
		}
	}
	return v;
}

extern value_t _call_list(void *addr, value_t list, value_t *edi);

value_t call_list_args(void *code_ptr, struct args *args, value_t list,
                       void *data)
{
	list = deep_copy(list);

	int nargs = length(list);

	value_t *val = &list;

	for (value_t i = list; !nilp(i); i = cdr(i))
	{
		val = cdrref(i);
	}

	int total_required = args->num_required + args->num_optional;

	if (nargs > total_required)
	{
		// Take the remainder of the list and put it as the last item in the
		// list.
		value_t trailing = cxdr(list, total_required);
		value_t last_item = cons(trailing, nil);

		*cxdrref(&list, total_required) = last_item;
	}
	else if (nargs < total_required)
	{
		for (int i = nargs - args->num_required; i < args->num_optional; i++)
		{
			// Append the i-th defualt argument
			value_t appended = cons(args->optional_arguments[i].value, nil);
			*val = appended;
			val = cdrref(appended);
		}
	}

	// We want to call this if we pass the correct # of arguments or less, just
	// not if we have already passed varargs. Appends a nil argument.
	if (nargs <= total_required)
	{
		// Enough real arguments but no variadic arguments. Pass a nil list.
		*val = cons(nil, nil);
	}

	return _call_list(code_ptr, list, data);
}

value_t call_list(struct function *fun, value_t list)
{
	return call_list_args(fun->code_ptr, fun->args, list, NULL);
}

value_t call_list_closure(struct closure *c, value_t list)
{
	return call_list_args(c->function, c->args, list, c->data);
}

struct args *new_args()
{
	struct args *a = malloc(sizeof(struct args));
	a->num_optional = 0;
	a->num_required = 0;
	a->variadic = false;

	return a;
}

struct args *add_optional_arg(struct args *args, value_t name, value_t value)
{
	int i = args->num_optional++;
	args =
	    realloc(args, sizeof(struct args) + sizeof(struct optional_argument) *
	                                            args->num_optional);

	args->optional_arguments[i] = (struct optional_argument){
	    .value = value,
	    .name = name,
	};

	return args;
}

bool are_args_acceptable(struct args *args, int number)
{
	if (args->variadic)
	{
		return number >= args->num_required;
	}
	else
	{
		return number >= args->num_required &&
		       number <= args->num_required + args->num_optional;
	}
}

struct args *list_to_args(struct environment *env, value_t list,
                          struct local *local)
{
	struct args *args = new_args();

	bool in_optional = false;

	for (value_t i = list; !nilp(i); i = cdr(i))
	{
		value_t val = car(i);
		if (symbolp(val))
		{
			if (!args->variadic && symstreq(val, "&"))
			{
				i = cdr(i);
				value_t name = car(i);

				if (!symbolp(name))
				{
					err("You must provide a symbol after & in an argument list "
					    "to bind the\n"
					    "variadic arguments to.");
				}

				args->variadic = true;

				add_variable(local, V_ARGUMENT, (char *)(name ^ SYMBOL_TAG),
				             args->num_optional + args->num_required);

				continue;
			}

			if (!in_optional)
			{
				add_variable(local, V_ARGUMENT, (char *)(val ^ SYMBOL_TAG),
				             args->num_required++);
			}
			else
			{
				char *name = (char *)(val ^ SYMBOL_TAG);
				if (name[0] == '&')
				{
					err("Non-optional argument following optional arguments "
					    "starts with a &\n"
					    "did you mean to declare a variadic argument? If so "
					    "leave a space\n"
					    "between the & and name.");
				}
				else
				{
					err("Cannot define a non-optional argument after an "
					    "optional one.");
				}
			}
		}
		else if (listp(val))
		{
			in_optional = true;
			int len = length(val);

			if (len != 2)
			{
				err("A list defining an optional value must be structured like "
				    "(name expr)\n"
				    "with exactly two arguments.");
			}

			value_t name = car(val);
			value_t expr = car(cdr(val));

			value_t function = cons(nil, cons(expr, nil));

			dasm_State *d =
			    compile_function(function, NS_ANONYMOUS, env, NULL, NULL, NULL,
			                     NULL, local->current_file_path);

			// TODO: GC stack top!
			value_t (*compiled)() = link_program(&d);

			value_t value = compiled();
			args = add_optional_arg(args, name, value);

			add_variable(local, V_ARGUMENT, (char *)(name ^ SYMBOL_TAG),
			             args->num_required + args->num_optional - 1);
		}
	}

	return args;
}

void display_args(struct args *args)
{
	printf("Args object taking %d require arguments and %d optionals:\n",
	       args->num_required, args->num_optional);

	for (int i = 0; i < args->num_optional; i++)
	{
		printf("   %d\t%s\n", i,
		       (char *)(args->optional_arguments[i].name ^ SYMBOL_TAG));
		printval(args->optional_arguments[i].value, 2);
	}
}