/*
 * lex.c — hand-rolled DFA. UTF-8 source, ASCII operators.
 *
 * Comments: //... and (slash-star ... star-slash). Both stripped.
 * Whitespace: space, tab, CR, NL.
 * Identifiers: [A-Za-z_][A-Za-z0-9_]*  — also matches keywords; we
 *   look up the kw table after lexing the run.
 * Integer: 0x[0-9a-fA-F_]+, 0o[0-7_]+, 0b[01_]+, [0-9][0-9_]*
 * Float:   [0-9]+'.'[0-9]+([eE][+-]?[0-9]+)?
 * Rune:    'x' with C-like escapes
 * String:  "..." with C-like escapes
 * Operators: longest match.
 *
 * No automatic semicolon insertion (Hare rule). The lexer only emits
 * what is in the source; the parser is responsible for non-empty rules.
 */
#include "ww.h"
#include <stdlib.h>
#include <string.h>
#include <errno.h>

void
lexinit(Lex *l, Arena *a, const char *file, const char *src, u64 len)
{
	memset(l, 0, sizeof *l);
	l->file = file;
	l->src = src;
	l->srclen = len;
	l->line = 1;
	l->col = 1;
	l->a = a;
}

static int
lpeek(Lex *l, u64 ahead)
{
	u64 p = l->pos + ahead;
	if (p >= l->srclen)
		return -1;
	return (unsigned char)l->src[p];
}

static int
lget(Lex *l)
{
	if (l->pos >= l->srclen)
		return -1;
	int c = (unsigned char)l->src[l->pos++];
	if (c == '\n') {
		l->line++;
		l->col = 1;
	} else {
		l->col++;
	}
	return c;
}

static Pos
lpos(Lex *l)
{
	Pos p = { l->file, l->line, l->col };
	return p;
}

static int
isidstart(int c)
{
	return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c == '_';
}

static int
isidcont(int c)
{
	return isidstart(c) || (c >= '0' && c <= '9');
}

static int
ishex(int c)
{
	return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') ||
	    (c >= 'A' && c <= 'F');
}

/* skip whitespace and comments. returns 0 on EOF, else 1. */
static int
skipws(Lex *l)
{
	for (;;) {
		int c = lpeek(l, 0);
		if (c < 0)
			return 0;
		if (c == ' ' || c == '\t' || c == '\r' || c == '\n') {
			lget(l);
			continue;
		}
		if (c == '/' && lpeek(l, 1) == '/') {
			while ((c = lpeek(l, 0)) >= 0 && c != '\n')
				lget(l);
			continue;
		}
		if (c == '/' && lpeek(l, 1) == '*') {
			lget(l); lget(l);
			int prev = -1;
			for (;;) {
				int x = lget(l);
				if (x < 0) {
					Pos p = lpos(l);
					errorf(p, "unterminated /* comment");
					l->errs++;
					return 0;
				}
				if (prev == '*' && x == '/')
					break;
				prev = x;
			}
			continue;
		}
		return 1;
	}
}

static u64
parseint(const char *s, u64 n, int base, int *ok)
{
	u64 v = 0;
	int got = 0;
	for (u64 i = 0; i < n; i++) {
		int c = (unsigned char)s[i];
		if (c == '_')
			continue;
		int d;
		if (c >= '0' && c <= '9') d = c - '0';
		else if (c >= 'a' && c <= 'f') d = c - 'a' + 10;
		else if (c >= 'A' && c <= 'F') d = c - 'A' + 10;
		else { *ok = 0; return 0; }
		if (d >= base) { *ok = 0; return 0; }
		/* overflow? cheap check */
		if (v > (u64)~0ULL / (u64)base) { *ok = 0; return 0; }
		v = v * (u64)base + (u64)d;
		got = 1;
	}
	*ok = got;
	return v;
}

static int
escape(Lex *l, int *out)
{
	int c = lget(l);
	if (c < 0) return -1;
	switch (c) {
	case 'n':  *out = '\n'; return 0;
	case 't':  *out = '\t'; return 0;
	case 'r':  *out = '\r'; return 0;
	case '\\': *out = '\\'; return 0;
	case '\'': *out = '\''; return 0;
	case '"':  *out = '"';  return 0;
	case '0':  *out = '\0'; return 0;
	case 'a':  *out = '\a'; return 0;
	case 'b':  *out = '\b'; return 0;
	case 'f':  *out = '\f'; return 0;
	case 'v':  *out = '\v'; return 0;
	case 'x': {
		int hi = lget(l), lo = lget(l);
		if (!ishex(hi) || !ishex(lo)) {
			Pos p = lpos(l);
			errorf(p, "bad \\x escape");
			l->errs++;
			return -1;
		}
		int h = (hi <= '9' ? hi - '0' : (hi | 0x20) - 'a' + 10);
		int o = (lo <= '9' ? lo - '0' : (lo | 0x20) - 'a' + 10);
		*out = (h << 4) | o;
		return 0;
	}
	}
	{ Pos p = lpos(l); errorf(p, "bad escape \\%c", c); l->errs++; }
	return -1;
}

static Tok
lexnum(Lex *l, Pos start)
{
	Tok t = (Tok){ TK_INT, start, NULL, 0, {0}, TK_NONE };
	u64 begin = l->pos;
	int base = 10;
	int isfloat = 0;
	int c = lpeek(l, 0);

	if (c == '0' && (lpeek(l, 1) == 'x' || lpeek(l, 1) == 'X')) {
		lget(l); lget(l);
		base = 16;
		while ((c = lpeek(l, 0)) >= 0 && (ishex(c) || c == '_'))
			lget(l);
	} else if (c == '0' && (lpeek(l, 1) == 'b' || lpeek(l, 1) == 'B')) {
		lget(l); lget(l);
		base = 2;
		while ((c = lpeek(l, 0)) >= 0 && (c == '0' || c == '1' || c == '_'))
			lget(l);
	} else if (c == '0' && (lpeek(l, 1) == 'o' || lpeek(l, 1) == 'O')) {
		lget(l); lget(l);
		base = 8;
		while ((c = lpeek(l, 0)) >= 0 && ((c >= '0' && c <= '7') || c == '_'))
			lget(l);
	} else {
		while ((c = lpeek(l, 0)) >= 0 && ((c >= '0' && c <= '9') || c == '_'))
			lget(l);
		if (lpeek(l, 0) == '.' && lpeek(l, 1) >= '0' && lpeek(l, 1) <= '9') {
			isfloat = 1;
			lget(l);
			while ((c = lpeek(l, 0)) >= 0 && ((c >= '0' && c <= '9') || c == '_'))
				lget(l);
			c = lpeek(l, 0);
			if (c == 'e' || c == 'E') {
				lget(l);
				if (lpeek(l, 0) == '+' || lpeek(l, 0) == '-')
					lget(l);
				while ((c = lpeek(l, 0)) >= 0 && c >= '0' && c <= '9')
					lget(l);
			}
		}
	}

	u64 n = l->pos - begin;
	t.text = astrndup(l->a, l->src + begin, n);
	t.tlen = n;

	if (isfloat) {
		t.kind = TK_FLOAT;
		/* strdup with underscores stripped before strtod */
		char *clean = amalloc(l->a, n + 1);
		u64 j = 0;
		for (u64 i = 0; i < n; i++)
			if (l->src[begin + i] != '_')
				clean[j++] = l->src[begin + i];
		clean[j] = '\0';
		errno = 0;
		t.v.fval = strtod(clean, NULL);
		if (errno) {
			errorf(start, "bad float literal '%s'", t.text);
			l->errs++;
		}
	} else {
		const char *digs = l->src + begin;
		u64 dn = n;
		if (base != 10) {
			digs += 2;
			dn -= 2;
		}
		int ok = 0;
		t.v.uval = parseint(digs, dn, base, &ok);
		if (!ok) {
			errorf(start, "bad integer literal '%s'", t.text);
			l->errs++;
			t.kind = TK_ERR;
		}
	}

	/* Typed suffix: i8/i16/i32/i64, u8/u16/u32/u64, f32/f64.
	 * Must be glued (no whitespace) to the digits. We grab the
	 * adjacent identifier-like run and accept it only if it's one
	 * of the recognised type names. */
	if (isidstart(lpeek(l, 0))) {
		u64 sb = l->pos;
		while (isidcont(lpeek(l, 0))) lget(l);
		u64 sl = l->pos - sb;
		const char *names[] = {
			"i8", "i16", "i32", "i64",
			"u8", "u16", "u32", "u64",
			"f32", "f64", NULL
		};
		const char *match = NULL;
		for (int i = 0; names[i]; i++) {
			u64 nl = strlen(names[i]);
			if (nl == sl && memcmp(names[i], l->src + sb, nl) == 0) {
				match = names[i];
				break;
			}
		}
		if (match) {
			t.tsuffix = astrndup(l->a, l->src + sb, sl);
		} else {
			/* not a known suffix — rewind so the run becomes a
			 * separate token. */
			l->pos = sb;
		}
	}
	return t;
}

static Tok
lexident(Lex *l, Pos start)
{
	u64 begin = l->pos;
	while (isidcont(lpeek(l, 0)))
		lget(l);
	u64 n = l->pos - begin;
	const char *p = l->src + begin;
	Tkind k = kwlookup(p, n);
	Tok t = (Tok){ k != TK_NONE ? k : TK_IDENT, start,
	    astrndup(l->a, p, n), n, {0}, TK_NONE };
	return t;
}

static Tok
lexstr(Lex *l, Pos start)
{
	/* opening quote already consumed by caller */
	u64 cap = 32, n = 0;
	char *buf = amalloc(l->a, cap);
	for (;;) {
		int c = lpeek(l, 0);
		if (c < 0) {
			errorf(start, "unterminated string");
			l->errs++;
			Tok t = (Tok){ TK_ERR, start, astrndup(l->a, "", 0), 0, {0}, TK_NONE };
			return t;
		}
		if (c == '"') { lget(l); break; }
		int ch;
		if (c == '\\') {
			lget(l);
			if (escape(l, &ch) < 0)
				ch = 0;
		} else {
			ch = lget(l);
		}
		if (n + 1 >= cap) {
			u64 ncap = cap * 2;
			char *nb = amalloc(l->a, ncap);
			memcpy(nb, buf, n);
			buf = nb;
			cap = ncap;
		}
		buf[n++] = (char)ch;
	}
	buf[n] = '\0';
	Tok t = (Tok){ TK_STR, start, buf, n, {0}, TK_NONE };
	return t;
}

static Tok
lexrune(Lex *l, Pos start)
{
	int ch;
	int c = lpeek(l, 0);
	if (c < 0) {
		errorf(start, "unterminated rune");
		l->errs++;
		return (Tok){ TK_ERR, start, "", 0, {0}, TK_NONE };
	}
	if (c == '\\') {
		lget(l);
		if (escape(l, &ch) < 0)
			ch = 0;
	} else {
		ch = lget(l);
	}
	if (lpeek(l, 0) != '\'') {
		errorf(start, "rune literal missing closing '");
		l->errs++;
		return (Tok){ TK_ERR, start, "", 0, {0}, TK_NONE };
	}
	lget(l);
	Tok t = (Tok){ TK_RUNE, start, NULL, 0, {0}, TK_NONE };
	t.v.uval = (u64)(u32)ch;
	t.text = aprintf(l->a, "%d", ch);
	t.tlen = strlen(t.text);
	return t;
}

#define EMIT(K)	do { Tok _t = (Tok){ (K), start, NULL, 0, {0}, TK_NONE }; \
		     _t.text = tokname(K); _t.tlen = strlen(_t.text); return _t; } while (0)

Tok
lexnext(Lex *l)
{
	if (!skipws(l)) {
		Pos p = lpos(l);
		Tok t = (Tok){ TK_EOF, p, "", 0, {0}, TK_NONE };
		return t;
	}
	Pos start = lpos(l);
	int c = lpeek(l, 0);

	if (isidstart(c))
		return lexident(l, start);
	if (c >= '0' && c <= '9')
		return lexnum(l, start);

	if (c == '"') { lget(l); return lexstr(l, start); }
	if (c == '\'') { lget(l); return lexrune(l, start); }

	lget(l);
	switch (c) {
	case '(':	EMIT(TK_LPAREN);
	case ')':	EMIT(TK_RPAREN);
	case '{':	EMIT(TK_LBRACE);
	case '}':	EMIT(TK_RBRACE);
	case '[':	EMIT(TK_LBRACK);
	case ']':	EMIT(TK_RBRACK);
	case ',':	EMIT(TK_COMMA);
	case ';':	EMIT(TK_SEMI);
	case ':':	EMIT(TK_COLON);
	case '@':	EMIT(TK_AT);
	case '?':	EMIT(TK_QUESTION);
	case '~':	EMIT(TK_TILDE);
	case '.':
		if (lpeek(l, 0) == '.' && lpeek(l, 1) == '.') {
			lget(l); lget(l);
			EMIT(TK_ELLIPSIS);
		}
		if (lpeek(l, 0) == '.') {
			lget(l);
			EMIT(TK_DOTDOT);
		}
		EMIT(TK_DOT);
	case '+':
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_PLUSEQ); }
		EMIT(TK_PLUS);
	case '-':
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_MINUSEQ); }
		if (lpeek(l, 0) == '>') { lget(l); EMIT(TK_ARROW); }
		EMIT(TK_MINUS);
	case '*':
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_STAREQ); }
		EMIT(TK_STAR);
	case '/':
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_SLASHEQ); }
		EMIT(TK_SLASH);
	case '%':
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_PERCENTEQ); }
		EMIT(TK_PERCENT);
	case '&':
		if (lpeek(l, 0) == '&') { lget(l); EMIT(TK_AND); }
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_AMPEQ); }
		EMIT(TK_AMP);
	case '|':
		if (lpeek(l, 0) == '|') { lget(l); EMIT(TK_OR); }
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_PIPEEQ); }
		EMIT(TK_PIPE);
	case '^':
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_CARETEQ); }
		EMIT(TK_CARET);
	case '=':
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_EQ); }
		if (lpeek(l, 0) == '>') { lget(l); EMIT(TK_FATARROW); }
		EMIT(TK_ASSIGN);
	case '!':
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_NEQ); }
		EMIT(TK_NOT);
	case '<':
		if (lpeek(l, 0) == '<') {
			lget(l);
			if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_LSHIFTEQ); }
			EMIT(TK_LSHIFT);
		}
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_LE); }
		if (lpeek(l, 0) == '-') { lget(l); EMIT(TK_LARROW); }
		EMIT(TK_LT);
	case '>':
		if (lpeek(l, 0) == '>') {
			lget(l);
			if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_RSHIFTEQ); }
			EMIT(TK_RSHIFT);
		}
		if (lpeek(l, 0) == '=') { lget(l); EMIT(TK_GE); }
		EMIT(TK_GT);
	}
	errorf(start, "unexpected character 0x%02x", c);
	l->errs++;
	Tok t = (Tok){ TK_ERR, start, NULL, 0, {0}, TK_NONE };
	t.text = astrndup(l->a, (const char[]){ (char)c }, 1);
	t.tlen = 1;
	return t;
}