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

#include "pre.h"

#define DATA_STACK_LEN 4096
#define RETURN_STACK_LEN 4096
#define REPL_LINE_BUF_LEN 4096

#define BOX_INTN(n) (Value){.kind = VAL_INT, .intn = n}
#define BOX_STR(n) (Value){.kind = VAL_STR, .str = n}
#define PILA_TRUE ((u64)1)
#define PILA_FALSE ((u64)0)
#define WRDF_IMMEDIATE (1 << 2)

typedef struct UserWord UserWord;
typedef struct Word Word;
typedef struct DictionaryEntry DictionaryEntry;
typedef struct Pila Pila;
typedef struct Str Str;
typedef struct Value Value;
typedef enum ValueKind ValueKind;
typedef enum WordKind WordKind;
typedef void (*NativeWord)(Pila *);

struct Str {
	char *s;
	isize len;
};

enum ValueKind {
	VAL_NIL,
	VAL_INT,
	VAL_STR,
};

struct Value {
	ValueKind kind;
	union {
		u64 intn;
		Str str;
	};
};

enum WordKind {
	WORD_NATIVE,
	WORD_USER,
};

struct UserWord {
	Word **words;
	isize len;
};

struct Word {
	Str name;
	u8 flags;
	WordKind kind;
	union {
		NativeWord nat;
		UserWord uword;
	};
};

struct Pila {
	Value data_stack[DATA_STACK_LEN];
	Word *rts[RETURN_STACK_LEN]; /* "return" stack */
	/* stack, return pointers */
	isize sp, rp;
	/* linked list, last added word is the first */
	DictionaryEntry *dict;
	
	/* state flags XXX: use a bitfield? */
	bool err, compiling;
	/* reader pointers */
	char *rsp; /* reader start pointer */
	char *rep; /* "      end   " */
	/* current word compilation state */
	Word *cur_compw;
	isize cur_compw_cap;
};

struct DictionaryEntry {
	Word *word;
	DictionaryEntry *next;
};

static const Value Nil = (Value){.kind = VAL_NIL};
static const Str StrEmpty = (Str){.s = "", .len = 0};

void
signal_error(Pila *st, const char *err);
void
signal_errorf(Pila *st, const char *err, ...);
Word
*search_word(Pila *st, Str word_name);
Str
read_word_name(char **pstart, char *pend);
bool
read_number(Str ns, u64 *res);
void
skip_whitespace(char **pstart, char *pend);
void
add_word(Pila *st, Word *w);
void
eval_word(Pila *st, Word *w);
void
eval(Pila *st, Str src);

bool
ascii_isdigit(char c)
{
	return '0' <= c && c <= '9';
}

bool
ascii_isspace(char c)
{
	return c == ' ' || ('\t' <= c && c <= '\r');
}

bool
ascii_iswordname(char c)
{
	return ('A' <= c && c <= 'Z')
		|| ('a' <= c && c <= 'z')
		|| (('!' <= c && c <= '\'') || ('*' <= c && c <= '@'))
		|| ascii_isdigit(c);
}

bool
Str_equal(Str s1, Str s2)
{
	/* because passing nil to mem* is UB even if size == 0... */
	return (s1.len == s2.len) && (s1.len == 0 || memcmp(s1.s, s2.s, s1.len) == 0);
}

Str
Str_new(const char *data, isize len)
{
	Assert(len >= 0);
	Str s;
	s.s = calloc(len + 1, sizeof(*s.s));
	s.len = len;
	if (data != nil) {
		memcpy(s.s, data, len);
		s.s[len + 1] = '\0'; /* ensure */
	}
	return s;
}

void
signal_errorf(Pila *st, const char *err, ...)
{
	va_list args;
	va_start(args, err);
	char buf[2048];
	vsnprintf(buf, sizeof(buf), err, args);
	fputs(buf, stderr);
	va_end(args);
	st->err = true;
}

void
signal_error(Pila *st, const char *err)
{
	fputs("error: ", stderr);
	fputs(err, stderr);
	fputc('\n', stderr);
	st->err = true;
}

void
push_val(Pila *st, Value val)
{
	if (st->sp + 1 > DATA_STACK_LEN) {
		signal_error(st, "stack overflow");
		return;
	}
	st->data_stack[st->sp++] = val;
}

bool
top_val(Pila *st, Value *val)
{
	if (st->sp == 0) {
		signal_error(st, "stack underflow");
		return false;
	}
	*val = st->data_stack[st->sp - 1];
	return true;
}

Value
pop_val(Pila *st)
{
	Value val;
	if (!top_val(st, &val))
		return Nil;
	--st->sp;
	return val;
}

void
add_nat(Pila *st)
{
	Value y = pop_val(st);
	Value x = pop_val(st);
	push_val(st, BOX_INTN(x.intn + y.intn));
}

void
sub_nat(Pila *st)
{
	Value y = pop_val(st);
	Value x = pop_val(st);
	push_val(st, BOX_INTN(x.intn - y.intn));
}

void
mul_nat(Pila *st)
{
	Value y = pop_val(st);
	Value x = pop_val(st);
	push_val(st, BOX_INTN(x.intn * y.intn));
}

void
div_nat(Pila *st)
{
	Value y = pop_val(st);
	Value x = pop_val(st);
	if (y.intn == 0 || x.intn == 0) {
		signal_error(st, "division by zero");
		return;
	}
	push_val(st, BOX_INTN(x.intn / y.intn));
}

void
lt_nat(Pila *st)
{
	Value rhs = pop_val(st);
	Value lhs = pop_val(st);
	push_val(st, BOX_INTN(lhs.intn < rhs.intn));
}

void
eq_nat(Pila *st)
{
	Value rhs = pop_val(st);
	Value lhs = pop_val(st);
	push_val(st, BOX_INTN(lhs.intn == rhs.intn));
}

void
not_nat(Pila *st)
{
	Value lhs = pop_val(st);
	push_val(st, BOX_INTN(!lhs.intn));
}

void
dup_nat(Pila *st)
{
	Value val;
	top_val(st, &val);
	push_val(st, val);
}

void
swap_nat(Pila *st)
{
	if (st->sp < 2)
		return;

	Value top = st->data_stack[st->sp - 1];
	st->data_stack[st->sp - 1] = st->data_stack[st->sp - 2];
	st->data_stack[st->sp - 2] = top;
}

void
drop_nat(Pila *st)
{
	pop_val(st);
}

void
bye_nat(Pila *st)
{
	(void)st;
	exit(0);
}

void
print_value(Pila *st, Value val, bool reader_fmt)
{
	(void)st;
	switch (val.kind) {
		case VAL_NIL:
			fputs("nil", stdout);
			break;
		case VAL_INT:
			printf("%lu", val.intn);
			break;
		case VAL_STR:
			printf(reader_fmt ? "\"%.*s\"": "%.*s", (int)val.str.len, val.str.s);
			break;
	}
}

void
stack_contents_nat(Pila *st)
{
	printf("<%lu> ", st->sp);
	if (st->sp == 0) {
		putchar('\n');
		return;
	}

	for (isize i = 0; i < st->sp; ++i) {
		Value val = st->data_stack[i];
		print_value(st, val, true);
		putchar(' ');
	}
	putchar('\n');
}

void
print_nat(Pila *st)
{
	Value top = pop_val(st);
	print_value(st, top, false);
	putchar('\n');
}

void
getch_nat(Pila *st)
{
	char c = getchar();
	push_val(st, BOX_INTN((u64)c));
}

void
emit_nat(Pila *st)
{
	Value c = pop_val(st);
	printf("%c", (char)c.intn);
}

void
compile_start(Pila *st, Str wn)
{
	Word *w = malloc(sizeof(*w));
	w->name = wn;
	w->kind = WORD_USER;

	UserWord uw;
	uw.words = calloc(16, sizeof(Word));
	uw.len = 0;

	w->uword = uw;
	st->cur_compw_cap = 16;
	st->cur_compw = w;
	st->compiling = true;

	add_word(st, w);
}

void
compile_start_nat(Pila *st)
{
	Str wn;
	skip_whitespace(&st->rsp, st->rep);
	wn = read_word_name(&st->rsp, st->rep);
	wn = Str_new(wn.s, wn.len);

	compile_start(st, wn);
}

void
compile_anon_start_nat(Pila *st)
{
	compile_start(st, StrEmpty);
}

void
compile_end_nat(Pila *st)
{
	if (!st->compiling) {
		signal_error(st, "not in compiling mode");
		return;
	}
	Assert(st->cur_compw != nil);
	if (!Str_empty(st->cur_compw->name))
		add_word(st, st->cur_compw);
	else /* for an anonymous word */
		push_val(st, BOX_INTN((u64)st->cur_compw));

	st->compiling = false;
}

Str
next_word(Pila *st)
{
	skip_whitespace(&st->rsp, st->rep);
	return read_word_name(&st->rsp, st->rep);
}

Word *
next_word_read_get(Pila *st)
{
	Str wn;
	skip_whitespace(&st->rsp, st->rep);
	wn = read_word_name(&st->rsp, st->rep);
	wn = Str_new(wn.s, wn.len);
	Word *w = search_word(st, wn);
	if (w == nil)
		return nil;
	return w;
}

void
call_nat(Pila *st)
{
	Value waddr_w = pop_val(st);
	if (waddr_w.kind == VAL_NIL) {
		return;
	}
	eval_word(st, (Word *)waddr_w.intn); /* :DDDD */
}

void
eval_nat(Pila *st)
{
	Value pstr = pop_val(st);
	if (pstr.kind == VAL_NIL || pstr.kind != VAL_STR) {
		signal_error(st, "cannot eval non string Value");
		return;
	}
	eval(st, pstr.str);
}

void
to_number_nat(Pila *st)
{
	Value ns = pop_val(st);
	if (ns.kind == VAL_NIL || ns.kind != VAL_STR) {
		signal_error(st, "cannot parse non string value");
		return;
	}
	u64 n;
	if (!read_number(ns.str, &n)) {
		signal_error(st, "invalid number");
		return;
	}
	push_val(st, BOX_INTN(n));
}

void
parse_word_nat(Pila *st)
{
	Str wn = next_word(st);
	wn = Str_new(wn.s, wn.len);
	push_val(st, BOX_STR(wn));
}

void
find_word_nat(Pila *st)
{
	Value wnw = pop_val(st);
	if (wnw.kind != VAL_STR) {
		signal_error(st, "expected string");
		return;
	}
	Word *w = search_word(st, wnw.str);
	if (w == nil) {
		push_val(st, Nil);
		return;
	}
	push_val(st, BOX_INTN((u64)w));
}

void
immediate_nat(Pila *st)
{
	if (st->cur_compw == nil) {
		signal_error(st, "no word to mark as immediate");
		return;
	}
	st->cur_compw->flags |= WRDF_IMMEDIATE;
}

void
branch_nat(Pila *st)
{
	Value predw = pop_val(st);
	if (predw.kind != VAL_INT) {
		return;
	}
	Str wn = next_word(st);
	if (predw.intn == PILA_TRUE) {
		Word *w = search_word(st, wn);
		if (w == nil) {
			signal_errorf(st, "no such word '%.*s'", (int)wn.len, wn.s);
			return;
		}
		eval_word(st, w);
	}
}

void
add_word(Pila *st, Word *w)
{
	DictionaryEntry *de = malloc(sizeof(*de));
	de->word = w;
	de->next = st->dict != nil ? st->dict : nil;
	st->dict = de;
}

Word *
search_word(Pila *st, Str word_name)
{
	DictionaryEntry *de = st->dict;
	while (de != nil) {
		if (Str_equal(de->word->name, word_name))
			return de->word;
		de = de->next;
	}
	return nil;
}

bool
read_number(Str ns, u64 *res)
{
	u64 n = 0;
	char *p = ns.s;
	char *pend = ns.s + ns.len;

	while (p < pend && ascii_isdigit(*p)) {
		u64 cv = *p - '0';
		n = n * 10 + cv;
		++p;
	}
	if (p < pend) /* trailling characters, invalid number */
		return false;

	*res = n;
	return true;
}

Str
read_word_name(char **pstart, char *pend)
{
	char *p = *pstart;
	Str wn = { .s = p };
	while (p < pend && ascii_iswordname(*p))
		++p;
	wn.len = p - *pstart;
	*pstart = p;
	return wn;
}

void
skip_whitespace(char **pstart, char *pend)
{
	char *p = *pstart;
	while (p < pend && ascii_isspace(*p))
		++p;
	*pstart = p;
}

void
read_comment(char **pstart, char *pend)
{
	char *p = *pstart;
	switch (*p) {
	case '\\':
		while (p < pend && *p++ != '\n');
		break;
	case '(':
		while (p < pend && *p++ != ')');
		break;
	default:
		break;
	}
	*pstart = p;
}

Str
read_string_lit(char **pstart, char *pend)
{
	isize i = 0;
	char buf[4096] = {0};
	char *p = *pstart;

	while (p < pend && *p != '"')
		buf[i++] = *p++;
	++p; /* move past end " */
	*pstart = p;
	return Str_new(buf, i);
}

void
eval_word(Pila *st, Word *w)
{
	/* Here, as an implementation detail, the "return stack" behaves
	 * more like a queue. */
	st->rts[st->rp++] = w;

	while (st->rp > 0) {
		Word *tw = st->rts[--st->rp];
		switch (tw->kind) {
		case WORD_NATIVE:
			tw->nat(st);
			break;
		case WORD_USER:
			for (isize i = tw->uword.len - 1; i >= 0; --i)
				st->rts[st->rp++] = tw->uword.words[i];
			break;
		}
	}
}

void
eval(Pila *st, Str src)
{
	st->rsp = src.s;
	st->rep = src.s + src.len;

	while (st->rsp < st->rep) {
		if (*st->rsp == '"') {
			++st->rsp;
			Str s = read_string_lit(&st->rsp, st->rep);
			push_val(st, BOX_STR(s));
		} else if (ascii_iswordname(*st->rsp)) {
			Str wn = read_word_name(&st->rsp, st->rep);
			Word *w = search_word(st, wn);
			if (w == nil) {
				u64 val;
				if (read_number(wn, &val)) {
					push_val(st, BOX_INTN(val));
					continue;
				} else {
					signal_errorf(st, "no such word: '%.*s'\n", (int)wn.len, wn.s);
					continue;
				}
			}

			if (!st->compiling)
				eval_word(st, w);
			else {
				if (w->flags & WRDF_IMMEDIATE) {
					eval_word(st, w);
				} else {
					st->cur_compw->uword.words[st->cur_compw->uword.len++] = w;
				}
			}
		} else if (*st->rsp == '\\' || *st->rsp == '(') {
			read_comment(&st->rsp, st->rep);
		} else {
			++st->rsp;
		}
	}
}

void
repl(Pila *st)
{
	char buf[REPL_LINE_BUF_LEN] = {0};
	for (;;) {
		fprintf(stdout, "%s> ", st->err ? "err" : "ok");
		fflush(stdout);
		if (fgets(buf, sizeof(buf), stdin) == nil) {
			break;
		}
		Str line = Str_from_c(buf);
		if (line.len == 0)
			continue;
		eval(st, line);
	}
}

int
main(int argc, char **argv)
{
	Pila st = {0};
	Word natws[] = {
		{.name = Sl("+"), .kind = WORD_NATIVE, .nat = add_nat },
		{.name = Sl("-"), .kind = WORD_NATIVE, .nat = sub_nat },
		{.name = Sl("*"), .kind = WORD_NATIVE, .nat = mul_nat },
		{.name = Sl("/"), .kind = WORD_NATIVE, .nat = div_nat },
		{.name = Sl("<"), .kind = WORD_NATIVE, .nat = lt_nat },
		{.name = Sl("="), .kind = WORD_NATIVE, .nat = eq_nat },
		{.name = Sl("!"), .kind = WORD_NATIVE, .nat = not_nat },
		{.name = Sl(":"), .kind = WORD_NATIVE, .nat = compile_start_nat },
		{.name = Sl(":?"), .kind = WORD_NATIVE, .nat = compile_anon_start_nat },
		{	.name = Sl(";"), .flags = WRDF_IMMEDIATE,
			.kind = WORD_NATIVE, .nat = compile_end_nat
		},
		{.name = Sl(">#"), .kind = WORD_NATIVE, .nat = to_number_nat },
		{.name = Sl("parse-word"), .kind = WORD_NATIVE, .nat = parse_word_nat },
		{.name = Sl("find-word"), .kind = WORD_NATIVE, .nat = find_word_nat },
		{.name = Sl("immediate"), .kind = WORD_NATIVE, .nat = immediate_nat },
		{.name = Sl("call"), .kind = WORD_NATIVE, .nat = call_nat},
		{.name = Sl("branch"), .kind = WORD_NATIVE, .nat = branch_nat },
		{.name = Sl("eval"), .kind = WORD_NATIVE, .nat = eval_nat},
		{.name = Sl("dup"), .kind = WORD_NATIVE, .nat = dup_nat },
		{.name = Sl("drop"), .kind = WORD_NATIVE, .nat = drop_nat },
		{.name = Sl("swap"), .kind = WORD_NATIVE, .nat = swap_nat },
		{.name = Sl("bye"), .kind = WORD_NATIVE, .nat = bye_nat },
		{.name = Sl("s?"), .kind = WORD_NATIVE, .nat = stack_contents_nat },
		{.name = Sl("emit"), .kind = WORD_NATIVE, .nat = emit_nat },
		{.name = Sl("."), .kind = WORD_NATIVE, .nat = print_nat },
		{.name = Sl(","), .kind = WORD_NATIVE, .nat = getch_nat },
	};
	for (isize i = 0; i < countof(natws); ++i)
		add_word(&st, &natws[i]);

	eval(&st, Sl(": ' parse-word find-word ;"));
	/* isolate prelude state */
	st.cur_compw = nil;
	st.cur_compw_cap = 0;

	if (argc > 1) {
		eval(&st, Str_from_c(argv[1]));
	}
	repl(&st);

	return 0;
}