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

#include "vm.h"
#include "ast.h"
#include "read.h"
#include "state.h"

static size_t placeholder(State *S, Chunk *ch) {
	size_t old_ix = ch->bc.len;
	chunk_wbc(S, ch, 0x00);
	chunk_wbc(S, ch, 0x00);
	return old_ix;
}

static void patch(State *S, Chunk *ch, size_t addr, uint16_t val) {
	ch->bc.d[addr] = val & 0xff;
	ch->bc.d[addr+1] = (val & 0xff00) >> 8;
}



static void compile_node(State *S, Chunk *ch, AstNode a);

typedef void (*form_compiler)(State *S, Chunk *ch, AstVec l, Op op);
typedef struct {
	char *name;
	int min_params;
	bool ellipsis;
	form_compiler action;
	Op op;
} BuiltinForm;


void single_form(State *S, Chunk *ch, AstVec l, Op op) {
	compile_node(S, ch, l.vals[1]);
	chunk_wbc(S, ch, op);
}

void set_form(State *S, Chunk *ch, AstVec l, Op op) {
	AstNode ident = l.vals[1];
	if (ident.ty != AST_IDENT) {
		fprintf(stderr, "set's first argument must be identifier");
		exit(1);
	}
	ObjString *o = objstring_copy_cstr(S, ident.as.str);
	compile_node(S, ch, l.vals[2]);
	chunk_wbc(S, ch, OP_SETGLOBAL);
	chunk_wbc(S, ch, chunk_wconst(S, ch, VAL_OBJ(o)));
}

void do_form(State *S, Chunk *ch, AstVec l, Op op) {
	for (int i = 1; i < l.len - 1; i++) {
		compile_node(S, ch, l.vals[i]);
		chunk_wbc(S, ch, OP_DROP);
	}
	compile_node(S, ch, l.vals[l.len - 1]);
}

void if_form(State *S, Chunk *ch, AstVec l, Op op) {
	// (if cond if-true if-false)
	// cond
	// 0branch ->A
	// if-true
	// skip ->B
	// A: if-false
	// B:
	compile_node(S, ch, l.vals[1]);
	chunk_wbc(S, ch, OP_0BRANCH);
	size_t ph_a = placeholder(S, ch);
	compile_node(S, ch, l.vals[2]);
	chunk_wbc(S, ch, OP_SKIP);
	size_t ph_b = placeholder(S, ch);
	size_t dest_a = ch->bc.len;
	compile_node(S, ch, l.vals[3]);
	size_t dest_b = ch->bc.len;

	patch(S, ch, ph_a, dest_a - ph_a - 2);
	patch(S, ch, ph_b, dest_b - ph_b - 2);
}

void while_form(State *S, Chunk *ch, AstVec l, Op op) {
	// (while cond body ...)
	// A:
	// cond
	// 0branch ->B
	// body ....
	// redo ->A
	// B:
	// nil   (while loop always returns nil)
	size_t dest_a = ch->bc.len;
	compile_node(S, ch, l.vals[1]);
	chunk_wbc(S, ch, OP_0BRANCH);
	size_t ph_b = placeholder(S, ch);
	for (int i = 2; i < l.len; i++)  {
		compile_node(S, ch, l.vals[i]);
		chunk_wbc(S, ch, OP_DROP);
	}
	chunk_wbc(S, ch, OP_REDO);
	size_t ph_a = placeholder(S, ch);
	size_t dest_b = ch->bc.len;
	chunk_wbc(S, ch, OP_NIL);

	patch(S, ch, ph_a, ph_a - dest_a + 2);
	patch(S, ch, ph_b, dest_b - ph_b - 2);
}

void arith_form(State *S, Chunk *ch, AstVec l, Op op) {
	compile_node(S, ch, l.vals[1]);
	compile_node(S, ch, l.vals[2]);
	chunk_wbc(S, ch, op);
}

static BuiltinForm builtin_forms[] = {
	{ "puts",  1, false, single_form, OP_PUTS },
	{ "print", 1, false, single_form, OP_PRINT },
	{ "set",   2, false, set_form, 0 },
	{ "do",    1,  true, do_form, 0 },
	{ "if",    3, false, if_form, 0 },
	{ "while", 2,  true, while_form, 0 },
#define ARITH_OP(str, op) \
	{ str, 2, false, arith_form, op },
	ARITH_OP("+", OP_ADD)
	ARITH_OP("-", OP_SUB)
	ARITH_OP("*", OP_MUL)
	ARITH_OP("/", OP_DIV)
	ARITH_OP("=", OP_EQU)
	ARITH_OP("<", OP_CMP)
	ARITH_OP("%", OP_MOD)
#undef ARITH_OP
	{ 0 },
};

typedef struct {
	char *name;
	Op op;
} BuiltinIdent;
static BuiltinIdent builtin_idents[] = {
	{ "true", OP_TRUE },
	{ "false", OP_FALSE },
	{ "nil", OP_NIL },
	{ 0 },
};

static void compile_node(State *S, Chunk *ch, AstNode a) {
	switch (a.ty) {
	case AST_IDENT:;
		char *ident = a.as.str;
		BuiltinIdent *found_builtin = NULL;
		for (BuiltinIdent *b = builtin_idents; b->name != NULL; b++) {
			if (0 == strcmp(b->name, ident)) {
				found_builtin = b;
				break;
			}
		}
		if (found_builtin != NULL) {
			chunk_wbc(S, ch, found_builtin->op);
		} else {
			// read global variable
			ObjString *o = objstring_copy_cstr(S, a.as.str);
			chunk_wbc(S, ch, OP_GETGLOBAL);
			chunk_wbc(S, ch, chunk_wconst(S, ch, VAL_OBJ(o)));
		}
		break;
	case AST_NUM:
		chunk_wbc(S, ch, OP_LOADK);
		chunk_wbc(S, ch, chunk_wconst(S, ch, VAL_NUM(a.as.num)));
		break;
	case AST_STRING: {
		ObjString *o = objstring_copy_cstr(S, a.as.str);
		chunk_wbc(S, ch, OP_LOADK);
		chunk_wbc(S, ch, chunk_wconst(S, ch, VAL_OBJ(o)));
		break;
	}
	case AST_LIST: {
		AstVec l = a.as.list;

		#define CK(cond, msg) if (!(cond)) { fputs(msg "\n", stderr); exit(1); }
		CK(l.len > 0, "can't handle empty list");
		CK(l.vals[0].ty == AST_IDENT, "can only call ops");

		char *head = l.vals[0].as.str;

		BuiltinForm *form = NULL;
		for (BuiltinForm *b = builtin_forms; b->name != NULL; b++) {
			if (0 == strcmp(b->name, head)) {
				form = b;
				break;
			}
		}

		if (form != NULL) {
			if (form->ellipsis && l.len < form->min_params + 1) {
				fprintf(stderr, "%s requires at least %d parameters\n",
					form->name, form->min_params);
				exit(1);
			} else if (!form->ellipsis && l.len != form->min_params + 1) {
				fprintf(stderr, "%s requires exactly %d parameters\n",
					form->name, form->min_params);
				exit(1);
			}

			form->action(S, ch, l, form->op);
		} else {
			fprintf(stderr, "unknown form %s\n", head);
			exit(1);
		}

		break;
	}
	}
}

int main(int argc, char **argv) {
	State st = state_new();
	State *S = &st;
	Thread th = thread_new(S);
	st.th = &th;
	Chunk ch = chunk_new(S);
	th.ch = &ch;

	S->do_disasm = (argc > 1 && 0 == strcmp(argv[1], "-l"));

	char n1[] = "foo";
	char n2[] = "bar";
	char n3[] = "baz";
	ObjString *o1 = objstring_copy_cstr(S, n1);
	ObjString *o2 = objstring_copy_cstr(S, n2);
	ObjString *o3 = objstring_copy_cstr(S, n3);

	ht_put(S, &st.globals, o1, VAL_NUM(69)); 
	ht_put(S, &st.globals, o2, VAL_NUM(2)); 
	ht_put(S, &st.globals, o3, VAL_NUM(3)); 

	AstNode an = read();
	compile_node(S, &ch, an);

	chunk_wbc(S, &ch, OP_PUTS);
	chunk_wbc(S, &ch, OP_RET);

	return runvm(S);
}