#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdarg.h>
#include <ctype.h>
#include <inttypes.h>
#include <errno.h>
#include "defs.h"
#include "dict.h"

#define STACK_SIZE 4096u

#define DEF_NATIVE(name) void name##_fn(void)
#define DEF_ARITH_FN(op, name, ctype)						\
	DEF_NATIVE(name) {										\
		int err = 0;										\
		struct value_box *n2 = pop_stack();					\
		struct value_box *n1 = pop_stack();					\
		if (n1 == NULL || n2 == NULL) return;				\
		ctype v2 = unbox_exb(n2, ctype, CET_NUMBER, &err);	\
		ctype v1 = unbox_exb(n1, ctype, CET_NUMBER, &err);	\
		if (err) return;									\
		push_stack(box_number(v1 op v2));					\
	}

#define DEF_BOX(name, val_type, box_member, type_tag)		\
	struct value_box *box_##name(val_type val) {			\
		struct value_box *tmp = xcalloc(1, sizeof(*tmp));	\
		tmp->type = type_tag;								\
		tmp->box_member = val;								\
		return tmp;											\
	}

struct {
	struct value_box *stack[STACK_SIZE];
	size_t sp;
} data_stack;

struct dict_table *dict;
extern struct dict_ent *dict_last; // pointer to the last added dictionary entry

size_t pc;
char *next_word;
char **program;
size_t program_len;

int compile_mode;
int error_flag;

void error(const char *fmt, ...)
{
	va_list args;
	va_start(args, fmt);
	vfprintf(stderr, fmt, args);
	va_end(args);
	error_flag = 1;
}

intmax_t xstrtoimax(const char *restrict nptr, int base, int *restrict eflag)
{
	errno = 0;
	char *endptr;
	intmax_t val = strtoimax(nptr, &endptr, base);
	if (((val == INTMAX_MAX || val == INTMAX_MIN) && errno == ERANGE) || (!val && errno == EINVAL) || *endptr != '\0') {
		if (eflag != NULL)
			*eflag = 1;
		else
			fatal_error("invalid number on stream\n");
	}
	return val;
}

// unboxes a value, only if it has the same type as specified.
// must be cast to the appropiate C type, since it returns a void *
void *unbox_ex(struct value_box *val, enum comp_ent_type type, int *err)
{
#define SETNNIL(v) if ((v) != NULL) {*(v) = 1;}
	if (val == NULL) {
		SETNNIL(err);
		return NULL;
	}
	if (val->type != type) {
		error("invalid type\n");
		SETNNIL(err);
		return NULL;
	}
	switch (val->type) {
		case CET_NUMBER: return &val->number;
		case CET_WORD:	 return &val->word;
		case CET_ADDR:	 return &val->addr;
		default: UNREACHABLE();
	}
	return NULL;
#undef SETNNIL
}

void push_stack(struct value_box *val)
{
	size_t *sp = &data_stack.sp;
	if ((*sp + 1) > SIZE_ARR(data_stack.stack)) {
		error("stack overflow\n");
		return;
	}
	data_stack.stack[(*sp)++] = val;
}

struct value_box *pop_stack()
{
	size_t *sp = &data_stack.sp;
	if ((*sp - 1) == SIZE_MAX) {
		error("stack underflow\n");
		return NULL;
	}
	return data_stack.stack[--*sp];
}

static inline void print_box(const struct value_box *obj)
{
	switch (obj->type) {
		case CET_NUMBER:
			printf("%jd", obj->number);
			break;
		case CET_WORD:
			printf("<word at %p (%s)>", (void *)obj->word, obj->word->word_name);
			break;
		case CET_ADDR: puts("address"); break;
		default: UNREACHABLE();
	}
}

// non-recursive word dump)
void word_dump(const char *word)
{
	struct dict_ent *ent;
	if ((ent = dict_lookup(dict, word)) != NULL) {
		if (ent->flags & FLAG_NATIVE) {
			printf("<native word at %p (%s)>\n", (void *)ent, ent->word_name);
			return;
		}
		printf(": %s ", word);
		for (size_t i = 0; i < ent->fn_def->len; ++i) {
			struct value_box *cent = &(ent->fn_def->ents[i]);
			switch (cent->type) {
				case CET_ADDR:
					printf("addr\n");
					break;
				case CET_WORD:
					printf("%s ", cent->word->word_name);
					break;
				case CET_NUMBER:
					printf("%jd ", cent->number);
					break;
				default: UNREACHABLE();
			}
		}
		puts(";");
	} else {
		error("no such word `%s`\n", word);
	}
}

DEF_BOX(number, INT_TYPE, number, CET_NUMBER)
DEF_BOX(word, struct dict_ent *, word, CET_WORD)

// --------------- native (primitives) --------------- //
DEF_ARITH_FN(+, add, INT_TYPE)
DEF_ARITH_FN(-, sub, INT_TYPE)
DEF_ARITH_FN(*, mul, INT_TYPE)
DEF_ARITH_FN(/, div, INT_TYPE)
DEF_ARITH_FN(==, equal, INT_TYPE)
DEF_ARITH_FN(<, less, INT_TYPE)

DEF_NATIVE(clearstack) {
	while (data_stack.sp > 0)
		free(data_stack.stack[--data_stack.sp]);
}

// ' <word> returns the address of a dictionary entry, aborting if it does not exist
DEF_NATIVE(quote) {
	++pc;
	if (strnil(next_word))
		error("' requires a name\n");

	struct dict_ent *ent = dict_lookup(dict, next_word);
	if (ent == NULL) {
		error("no such word `%s'\n", next_word);
		return;
	}
	push_stack(box_word(ent));
}

// allot (n -- ) adds n bytes to the parameter field of the most recently defined word
DEF_NATIVE(allot) {
	int err = 0;
	const size_t nbytes = unbox_exsp(size_t, CET_NUMBER, &err);
	if (err)
		return;
	const size_t org_len = dict_last->fn_def->len;
	const size_t len = org_len + nbytes;
	struct value_box **ents = &dict_last->fn_def->ents;

	dict_last->fn_def->capacity = len;
	*ents = reallocarray(*ents, len, sizeof(struct value_box));
	// printf("allot: %d\n", len - dict_last->fn_def->len);
	// clear the newly allocated part
	//memset(*ents + org_len, 0, len - org_len);
}

DEF_NATIVE(comma) {
	struct comp_word *pf = dict_last->fn_def;

	if (pf->ents == NULL) {
		pf->capacity = 4;
		pf->ents = xcalloc(pf->capacity, sizeof(struct value_box));
	} else if (pf->len + 1 > pf->capacity) {
		pf->capacity <<= 1;
		pf->ents = reallocarray(pf->ents, pf->capacity, sizeof(struct value_box));
	}
	struct value_box *val = pop_stack();
	if (val == NULL)
		return;
	pf->ents[pf->len++] = *val;
}

DEF_NATIVE(create) {
	++pc;
	if (strnil(next_word))
		error("create requires a name\n");

	struct comp_word *param = xcalloc(1, sizeof(*param));
	dict_insert(dict, next_word, 0, NULL, param);
}

DEF_NATIVE(forget) {
	++pc;
	if (strnil(next_word))
		error("forget requires a name\n");
	struct dict_ent *ent;
	if ((ent = dict_lookup(dict, next_word)) == NULL) {
		error("no such word `%s'\n", next_word);
		return;
	}

	if (!(ent->flags & FLAG_NATIVE))
		free(ent->fn_def->ents);
	dict_delete(ent);
}

DEF_NATIVE(see) {
	++pc;
	if (strnil(next_word)) {
		error("see requires a name\n");
		return;
	}
	word_dump(next_word);
}

DEF_NATIVE(words) {
	for (size_t i = 0; i < dict->capacity; ++i)
		if (IS_ENTRY(&dict->entries[i]))
			printf("%s ", dict->entries[i].word_name);
	putchar('\n');
}

DEF_NATIVE(compile) {
	compile_mode = 1;
}

DEF_NATIVE(endcompile) {
	compile_mode = 0;
}

DEF_NATIVE(dup) {
	push_stack(data_stack.stack[data_stack.sp - 1]);
}

DEF_NATIVE(drop) {
	free(pop_stack());
}

DEF_NATIVE(over) {
	push_stack(data_stack.stack[data_stack.sp - 2]);
}

DEF_NATIVE(swap) {
	struct value_box *e2 = pop_stack();
	struct value_box *e1 = pop_stack();
	if (e1 == NULL || e2 == NULL)
		return;
	push_stack(e2);
	push_stack(e1);
}

DEF_NATIVE(zero) {
	int err = 0;
	struct value_box *tmp = box_number(!unbox_exsp(INT_TYPE, CET_NUMBER, &err));
	if (err)
		return;
	push_stack(tmp);
}

DEF_NATIVE(emit) {
	int err = 0;
	int chr = unbox_exsp(int, CET_NUMBER, &err);
	if (!err)
		putchar(chr);
}

DEF_NATIVE(print) {
	struct value_box *val = pop_stack();
	if (val == NULL)
		return;
	print_box(val);
	putchar('\n');
	free(val);
}

DEF_NATIVE(prints) {
	if (data_stack.sp) {
		size_t c = 0;
		printf("stack dump: <%zu> ", data_stack.sp);
		while (c < data_stack.sp) {
			print_box(data_stack.stack[c++]);
			putchar(' ');
		}
		putchar('\n');
	}
}

static inline void append_to_code(struct value_box *slot, const char *current_word)
{
	int errn = 0;
	intmax_t val = xstrtoimax(current_word, 0, &errn);

	if (!errn) {
		slot->type = CET_NUMBER;
		slot->number = val;
	} else {
		struct dict_ent *ent = dict_lookup(dict, current_word);
		if (ent == NULL)
			fatal_error("no such word `%s`\n", current_word);
		slot->type = CET_WORD;
		slot->word = ent;
	}
}

void eval_def(struct comp_word *def);
void if_fn()
{
	size_t start_pc = ++pc;
	size_t end_pc;
	while (pc < program_len && strcmp(program[pc], "end") != 0)
		++pc;
	if (pc >= program_len)
		fatal_error("expected `end`, reached EOP\n");
	end_pc = pc;

	int err = 0;
	int cond = unbox_exsp(int, CET_NUMBER, &err);
	if (err)
		fatal_error("if requires a condition\n");

	if (cond) {
		char **program_slice = program + start_pc;
		size_t len = end_pc - start_pc;

		struct value_box *code = xcalloc(len, sizeof(*code));
		struct comp_word c = {.len = len, .capacity = len, .ents = code};

		for (size_t i = 0; i < len; ++i)
			append_to_code(&code[i], program_slice[i]);
		eval_def(&c);
		free(code);
	}
}

// --------------- interpretation --------------- //
void eval_def(struct comp_word *def)
{
	for (size_t i = 0; i < def->len; ++i) {
		const struct value_box *cent = &(def->ents[i]);
		struct dict_ent *ent;

		switch (cent->type) {
			case CET_WORD:
				ent = cent->word;
				if (ent->flags & FLAG_NATIVE)
					ent->nat_fn();
				else
					eval_def(ent->fn_def);
				break;
			case CET_NUMBER:
				push_stack(box_number(cent->number));
				break;
			default: UNREACHABLE();
		}
	}
}

void interpret(char **prog, size_t prog_len)
{
	struct comp_word *c_word = NULL;
	char *fn_name = NULL;

	for (; pc < prog_len; ++pc) {
		if (error_flag) {
			pc = prog_len;
			break;
		}
		char *word = prog[pc];
		next_word = (pc + 1) < prog_len ? prog[pc + 1] : NULL;

		if (compile_mode) {
			if (*word == ';') {
				dict_insert(dict, fn_name, 0, NULL, c_word);
				c_word = NULL;
				fn_name = NULL;
				compile_mode = 0;
				continue;
			}
			if (c_word == NULL) {
				c_word = xcalloc(1, sizeof(*c_word));
				c_word->ents = xcalloc(256, sizeof(*c_word->ents));
			}
			if (fn_name == NULL) {
				fn_name = word;
				continue;
			}
			append_to_code(&(c_word->ents[c_word->len++]), word);
		} else {
			int nef = 0;
			INT_TYPE num = xstrtoimax(word, 0, &nef);
			if (!nef) {
				push_stack(box_number(num));
				continue;
			}

			struct dict_ent *ent;
			if ((ent = dict_lookup(dict, word)) != NULL) {
				if (ent->flags & FLAG_NATIVE)
					ent->nat_fn();
				else
					eval_def(ent->fn_def);
			} else {
				fatal_error("interpretation: no such word `%s`\n", word);
			}
		}
	}
}

// --------------- Parsing --------------- //
void tokenize(char *s, char **buf, size_t buflen, size_t *i)
{
	const char delim = ' ';
	char *p = s, *wp = s;

	while (*p != '\0') {
		if (*p == '(') { // eat block comment
			while (*p != ')' && *p != '\0')
				++p;
			p += 2; // jump past the ) and space
			wp = p;
			continue;
		} else if (*p == '\\') {
			while (*p++ != '\0');
			wp = p;
			continue;
		}
		if (*p == delim) {
			if (*i + 1 > buflen)
				return;
			*p = '\0';
			buf[(*i)++] = strdup(wp);
			wp = ++p;
		} else {
			++p;
		}
	}
	if (wp < p && *i < buflen) // handle last element
		buf[(*i)++] = strdup(wp);
}

size_t parse_program(FILE *fp, char **buf, size_t buflen)
{
	size_t i = 0;
	char line_buf[1024] = {0};
	while (fgets(line_buf, sizeof(line_buf), fp) != NULL) {
		STRIP_FINAL_NL(line_buf);
		tokenize(line_buf, buf, buflen, &i);
	}
	return i;
}

void repl(void)
{
	char *prog[1024] = {0};
	size_t prog_len = 0;

	char *line_buf = NULL;
	size_t line_buf_len = 0;
	ssize_t rb = 0;
	int prompt_comp_mode = 0;

	errno = 0;
	for (;;) {
		printf("[%s]%s ", error_flag ? "err" : "ok", !prompt_comp_mode ? ">" : ">>");
		fflush(stdout);
		error_flag = 0;

		if ((rb = getline(&line_buf, &line_buf_len, stdin)) < 0)
			break;

		STRIP_FINAL_NL(line_buf);
		tokenize(line_buf, prog, SIZE_ARR(prog), &prog_len);
		if (strchr(line_buf, ':') != NULL)
			prompt_comp_mode = 1;
		else if (strchr(line_buf, ';') != NULL)
			prompt_comp_mode = 0;

		if (!prompt_comp_mode)
			interpret(prog, prog_len);
	}
	if (rb < 0 && errno)
		fatal_error("getline err\n");
	free(line_buf);
}

void init()
{
	dict = dict_new_env(NULL);
	dict_insert(dict, "+", FLAG_NATIVE, add_fn, NULL);
	dict_insert(dict, "-", FLAG_NATIVE, sub_fn, NULL);
	dict_insert(dict, "*", FLAG_NATIVE, mul_fn, NULL);
	dict_insert(dict, "=", FLAG_NATIVE, equal_fn, NULL);
	dict_insert(dict, "<", FLAG_NATIVE, less_fn, NULL);
	dict_insert(dict, ".", FLAG_NATIVE, print_fn, NULL);
	dict_insert(dict, ".s", FLAG_NATIVE, prints_fn, NULL);
	dict_insert(dict, "zero?", FLAG_NATIVE, zero_fn, NULL);
	dict_insert(dict, ":", FLAG_NATIVE, compile_fn, NULL);
	dict_insert(dict, ";", FLAG_NATIVE | FLAG_IMMEDIATE, endcompile_fn, NULL);
	dict_insert(dict, "'", FLAG_NATIVE, quote_fn, NULL);
	dict_insert(dict, ",", FLAG_NATIVE, comma_fn, NULL);
	dict_insert(dict, "print", FLAG_NATIVE, print_fn, NULL);
	dict_insert(dict, "see", FLAG_NATIVE, see_fn, NULL);
	dict_insert(dict, "dup", FLAG_NATIVE, dup_fn, NULL);
	dict_insert(dict, "drop", FLAG_NATIVE, drop_fn, NULL);
	dict_insert(dict, "over", FLAG_NATIVE, over_fn, NULL);
	dict_insert(dict, "swap", FLAG_NATIVE, swap_fn, NULL);
	dict_insert(dict, "create", FLAG_NATIVE, create_fn, NULL);
	dict_insert(dict, "forget", FLAG_NATIVE, forget_fn, NULL);
	dict_insert(dict, "allot", FLAG_NATIVE, allot_fn, NULL);
	dict_insert(dict, "cst", FLAG_NATIVE, clearstack_fn, NULL);
	dict_insert(dict, "words", FLAG_NATIVE, words_fn, NULL);
	dict_insert(dict, "emit", FLAG_NATIVE, emit_fn, NULL);
	dict_insert(dict, "if", FLAG_NATIVE, if_fn, NULL);
}

int main(int argc, char **argv)
{
	init();
	if (argc < 2) {
		printf("nsforth interactive repl (build: %s, %s)\n", __DATE__, __TIME__);
		repl();
	} else {
		FILE *fp = fopen(*++argv, "r");
		if (fp == NULL)
			fatal_error("error opening '%s': %s\n", *argv, strerror(errno));

		char *prog[1024] = {0};
		size_t prog_len = parse_program(fp, prog, SIZE_ARR(prog));
		fclose(fp);
		program = prog;
		program_len = prog_len;
		interpret(prog, prog_len);
	}

	dict_destroy_env(dict);
	return EXIT_SUCCESS;
}