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

static char *
toklit(Arena *a, Tok t)
{
	switch (t.kind) {
	case TK_IDENT: return aprintf(a, "IDENT(%s)", t.text);
	case TK_INT:   return aprintf(a, "INT(%llu)", (unsigned long long)t.v.uval);
	case TK_FLOAT: return aprintf(a, "FLOAT(%g)", t.v.fval);
	case TK_RUNE:  return aprintf(a, "RUNE(%llu)", (unsigned long long)t.v.uval);
	case TK_STR:   return aprintf(a, "STR(%s)", t.text);
	case TK_ERR:   return aprintf(a, "ERR(%s)", t.text);
	default:       return (char *)tokname(t.kind);
	}
}

static int
runrown(const char *src, size_t len, const char *expect, int expecterrs)
{
	Arena *a = newarena();
	Lex l;
	lexinit(&l, a, "<test>", src, len);

	char *got = amalloc(a, 1);
	got[0] = '\0';
	u64 cap = 1, n = 0;

	for (;;) {
		Tok t = lexnext(&l);
		if (t.kind == TK_EOF)
			break;
		const char *piece = toklit(a, t);
		u64 plen = strlen(piece);
		u64 need = n + plen + 2;
		if (need >= cap) {
			u64 nc = need * 2;
			char *nb = amalloc(a, nc);
			memcpy(nb, got, n);
			got = nb;
			cap = nc;
		}
		if (n) got[n++] = ' ';
		memcpy(got + n, piece, plen);
		n += plen;
		got[n] = '\0';
	}

	int ok = strcmp(got, expect) == 0
	    && (expecterrs < 0 || l.errs == expecterrs);
	if (!ok) {
		fprintf(stderr, "lex mismatch:\n  want:   %s (%d errors)\n"
		    "  got:    %s (%d errors)\n", expect, expecterrs, got,
		    l.errs);
	}
	freearena(a);
	return ok;
}

static int
runrow(const char *src, const char *expect)
{
	return runrown(src, strlen(src), expect, -1);
}

struct row { const char *src, *expect; };

static const struct row rows[] = {
	{ "",                "" },
	{ "\xef\xbb\xbf" "package main;", "package IDENT(main) ;" },
	{ "  \t\n  ",        "" },
	{ "// comment\n",    "" },
	{ "/* a /b/ c */",   "" },

	{ "foo",             "IDENT(foo)" },
	{ "fn",              "fn" },
	{ "fn main",         "fn IDENT(main)" },
	{ "let x: i32 = 0;", "let IDENT(x) : IDENT(i32) = INT(0) ;" },
	{ "export fn",       "export fn" },
	{ "if else for switch case return import type struct defer break continue proc chan nil true false package",
	  "if else for switch case return import type struct defer break continue proc chan nil true false package" },

	{ "0",               "INT(0)" },
	{ "42",              "INT(42)" },
	{ "1_000_000",       "INT(1000000)" },
	{ "0xff",            "INT(255)" },
	{ "0xDE_AD_BE_EF",   "INT(3735928559)" },
	{ "0b1010",          "INT(10)" },
	{ "0o777",           "INT(511)" },
	{ "18446744073709551616", "ERR(18446744073709551616)" },
	{ "3.14",            "FLOAT(3.14)" },
	{ "1.5e3",           "FLOAT(1500)" },

	{ "\"hello\"",       "STR(hello)" },
	{ "\"a\\nb\"",       "STR(a\nb)" },
	{ "'A'",             "RUNE(65)" },
	{ "'\\n'",           "RUNE(10)" },
	{ "'\\x7f'",         "RUNE(127)" },

	/* unicode escapes — \u (4 hex) / \U (8 hex) share the \x codepoint
	 * path (ref/hare/hare/lex/lex.ha:347 lex_unicode). Rune literals
	 * carry the raw codepoint; string literals UTF-8-encode it. */
	{ "'\\u00e9'",       "RUNE(233)" },
	{ "'\\u20ac'",       "RUNE(8364)" },
	{ "'\\U0001F600'",   "RUNE(128512)" },
	{ "\"\\u00e9\"",     "STR(\xc3\xa9)" },
	{ "\"\\u20ac\"",     "STR(\xe2\x82\xac)" },
	{ "\"\\U0001F600\"", "STR(\xf0\x9f\x98\x80)" },
	{ "\"caf\\u00e9\"",  "STR(caf\xc3\xa9)" },
	/* \x now also yields a codepoint that UTF-8-encodes in strings:
	 * \xe9 -> U+00E9 -> 0xC3 0xA9 (matches Hare's appendrune). */
	{ "'\\xe9'",         "RUNE(233)" },
	{ "\"\\xe9\"",       "STR(\xc3\xa9)" },
	/* range boundaries: max valid codepoint, and the two edges that
	 * straddle the UTF-16 surrogate gap (0xD7FF ok / 0xE000 ok). */
	{ "'\\U0010FFFF'",   "RUNE(1114111)" },
	{ "'\\uD7FF'",       "RUNE(55295)" },
	{ "'\\uE000'",       "RUNE(57344)" },

	{ "+ - * / % == != < > <= >= && || !",
	  "+ - * / % == != < > <= >= && || !" },
	{ "= += -= *= /= %= &= |= ^= <<= >>=",
	  "= += -= *= /= %= &= |= ^= <<= >>=" },
	{ "<< >> & | ^ ~ ?",
	  "<< >> & | ^ ~ ?" },
	{ "( ) { } [ ] , ; : . ... @",
	  "( ) { } [ ] , ; : . ... @" },
	{ "<- ->",
	  "<- ->" },
	{ "@symbol(\"malloc\")",
	  "@ IDENT(symbol) ( STR(malloc) )" },

	{ "fn add(a: i32, b: i32) i32 = { return a + b; };",
	  "fn IDENT(add) ( IDENT(a) : IDENT(i32) , IDENT(b) : IDENT(i32) ) IDENT(i32) = { return IDENT(a) + IDENT(b) ; } ;" },
};

/* A bad escape sets l.errs (the string/rune still lexes with the
 * offending codepoint zeroed), so detection is via the error counter,
 * not the token stream. The verbatim Hare messages live at lexunicode
 * in cmd/wcc/lex.c. */
static int
runerrn(const char *src, size_t len)
{
	Arena *a = newarena();
	Lex l;
	lexinit(&l, a, "<test>", src, len);
	for (;;) {
		Tok t = lexnext(&l);
		if (t.kind == TK_EOF)
			break;
	}
	int ok = l.errs > 0;
	if (!ok)
		fprintf(stderr, "expected escape error, none raised:\n  src: %s\n",
		    src);
	freearena(a);
	return ok;
}

static int
runerr(const char *src)
{
	return runerrn(src, strlen(src));
}

static int
runrewind(void)
{
	const char *src = "1foo bar";
	Arena *a = newarena();
	Lex l;
	lexinit(&l, a, "<test>", src, strlen(src));
	Tok n = lexnext(&l);
	Tok f = lexnext(&l);
	Tok b = lexnext(&l);
	int ok = n.kind == TK_INT && n.pos.col == 1 && n.v.uval == 1
	    && f.kind == TK_IDENT && f.pos.col == 2
	    && strcmp(f.text, "foo") == 0
	    && b.kind == TK_IDENT && b.pos.col == 6
	    && strcmp(b.text, "bar") == 0;
	if (!ok)
		fprintf(stderr, "suffix rewind columns: want 1,2,6 got %d,%d,%d\n",
		    n.pos.col, f.pos.col, b.pos.col);
	freearena(a);
	return ok;
}

static int
runbompos(void)
{
	const char *src = "\xef\xbb\xbf" "package main;";
	Arena *a = newarena();
	Lex l;
	lexinit(&l, a, "<test>", src, strlen(src));
	Tok t = lexnext(&l);
	int ok = t.kind == TK_MODULE && t.pos.line == 1 && t.pos.col == 4;
	if (!ok)
		fprintf(stderr, "leading BOM position: want 1:4 got %d:%d\n",
		    t.pos.line, t.pos.col);
	freearena(a);
	return ok;
}

static const char *const errrows[] = {
	"\xef\xbb\xbf" "\xef\xbb\xbf" "package main;",
	"package \xef\xbb\xbf" "main;",
	"// \xef\xbb\xbf\n",
	"/* \xef\xbb\xbf */",
	"\"\xef\xbb\xbf\"",
	"'\xef\xbb\xbf'",
	"'\\uZ'",		/* unexpected rune scanning for escape */
	"\"\\u00g0\"",		/* non-hex digit inside a string escape */
	"'\\u00",		/* unexpected EOF scanning for escape */
	"'\\UFFFFFFFF'",	/* codepoint > 0x10FFFF (high bit set) */
	"'\\U00110000'",	/* exactly one past U+10FFFF */
	"'\\uD800'",		/* bottom of the UTF-16 surrogate range */
	"'\\uDFFF'",		/* top of the UTF-16 surrogate range */
};

static const struct {
	const char *src;
	size_t len;
	const char *expect;
	int errs;
} nulrows[] = {
	{ "\0package main;", sizeof "\0package main;" - 1,
	  "package IDENT(main) ;", 1 },
	{ "pack\0age main;", sizeof "pack\0age main;" - 1,
	  "package IDENT(main) ;", 1 },
	{ "1\0_0", sizeof "1\0_0" - 1, "INT(10)", 1 },
	{ "1f\0oo", sizeof "1f\0oo" - 1, "INT(1) IDENT(foo)", 1 },
	{ "=\0=", sizeof "=\0=" - 1, "==", 1 },
	{ "/\0/ comment\nfn", sizeof "/\0/ comment\nfn" - 1, "fn", 1 },
	{ "/* end *\0/ fn", sizeof "/* end *\0/ fn" - 1, "fn", 1 },
	{ "\"a\0b\"", sizeof "\"a\0b\"" - 1, "STR(ab)", 1 },
	{ "'\0A'", sizeof "'\0A'" - 1, "RUNE(65)", 1 },
	{ "fn\0", sizeof "fn\0" - 1, "fn", 1 },
	{ "f\0\0n", sizeof "f\0\0n" - 1, "fn", 2 },
};

static int
runescapenul(void)
{
	static const char line[] = "\"\\\0n\"";
	static const char hex[] = "\"\\x0" "\0" "0\"";
	Arena *a = newarena();
	Lex l;
	lexinit(&l, a, "<test>", line, sizeof line - 1);
	Tok t = lexnext(&l);
	int ok = t.kind == TK_STR && t.tlen == 1 && t.text[0] == '\n'
	    && l.errs == 1 && lexnext(&l).kind == TK_EOF;
	lexinit(&l, a, "<test>", hex, sizeof hex - 1);
	t = lexnext(&l);
	ok = ok && t.kind == TK_STR && t.tlen == 1 && t.text[0] == '\0'
	    && l.errs == 1 && lexnext(&l).kind == TK_EOF;
	if (!ok)
		fprintf(stderr, "NUL escape recovery failed\n");
	freearena(a);
	return ok;
}

static int
runsuffixnul(void)
{
	static const char src[] = "1i" "\0" "8";
	Arena *a = newarena();
	Lex l;
	lexinit(&l, a, "<test>", src, sizeof src - 1);
	Tok t = lexnext(&l);
	int ok = t.kind == TK_INT && t.v.uval == 1 && t.tsuffix != NULL
	    && strcmp(t.tsuffix, "i8") == 0 && l.errs == 1
	    && lexnext(&l).kind == TK_EOF;
	if (!ok)
		fprintf(stderr, "NUL typed-suffix recovery failed\n");
	freearena(a);
	return ok;
}

static int
runlongdirective(void)
{
	static const char pre[] = "//ww:module ";
	static const char tail[] = "\nfn";
	const size_t pathlen = 16384;
	const size_t len = sizeof pre - 1 + pathlen + sizeof tail - 1;
	char *src = malloc(len);
	if (!src)
		return 0;
	memcpy(src, pre, sizeof pre - 1);
	memset(src + sizeof pre - 1, 'a', pathlen);
	memcpy(src + sizeof pre - 1 + pathlen, tail, sizeof tail - 1);
	Arena *a = newarena();
	Lex l;
	lexinit(&l, a, "<test>", src, len);
	Tok m = lexnext(&l);
	Tok f = lexnext(&l);
	int ok = m.kind == TK_MODPATH && m.tlen == pathlen
	    && m.text[0] == 'a' && m.text[pathlen - 1] == 'a'
	    && f.kind == TK_FN && l.errs == 0;
	if (!ok)
		fprintf(stderr, "long module directive failed\n");
	freearena(a);
	free(src);
	return ok;
}

int
main(void)
{
	int fail = 0;
	for (size_t i = 0; i < sizeof rows / sizeof rows[0]; i++) {
		if (!runrow(rows[i].src, rows[i].expect)) {
			fprintf(stderr, "row %zu failed\n", i);
			fail++;
		}
	}
	if (!runrewind())
		fail++;
	if (!runbompos())
		fail++;
	for (size_t i = 0; i < sizeof errrows / sizeof errrows[0]; i++) {
		if (!runerr(errrows[i])) {
			fprintf(stderr, "errrow %zu failed\n", i);
			fail++;
		}
	}
	for (size_t i = 0; i < sizeof nulrows / sizeof nulrows[0]; i++) {
		if (!runrown(nulrows[i].src, nulrows[i].len,
		    nulrows[i].expect, nulrows[i].errs)) {
			fprintf(stderr, "NUL row %zu failed\n", i);
			fail++;
		}
	}
	if (!runescapenul())
		fail++;
	if (!runsuffixnul())
		fail++;
	if (!runlongdirective())
		fail++;
	if (fail) {
		fprintf(stderr, "%d/%zu lex tests failed\n", fail,
		    sizeof rows / sizeof rows[0]);
		return 1;
	}
	printf("lex: %zu/%zu ok\n", sizeof rows / sizeof rows[0],
	    sizeof rows / sizeof rows[0]);
	return 0;
}