package lexer

import "../source"
import "../symbol"
import "../token"

is_identifier_start :: proc(value: byte) -> bool {
	return value == '_' || value >= 'a' && value <= 'z' || value >= 'A' && value <= 'Z'
}

is_identifier_continue :: proc(value: byte) -> bool {
	return is_identifier_start(value) || value >= '0' && value <= '9'
}

keyword_kind :: proc(text: string) -> token.Kind {
	switch text {
	case "func":          return .Keyword_Func
	case "c_func":        return .Keyword_C_Func
	case "struct":        return .Keyword_Struct
	case "c_struct":      return .Keyword_C_Struct
	case "enum":          return .Keyword_Enum
	case "distinct":      return .Keyword_Distinct
	case "alias":         return .Keyword_Alias
	case "import":        return .Keyword_Import
	case "return":        return .Keyword_Return
	case "mut":           return .Keyword_Mut
	case "none":          return .Keyword_None
	case "orelse":        return .Keyword_Orelse
	case "and":           return .Keyword_And
	case "or":            return .Keyword_Or
	case "if":            return .Keyword_If
	case "while":         return .Keyword_While
	case "for":           return .Keyword_For
	case "else":          return .Keyword_Else
	case "true":          return .Keyword_True
	case "false":         return .Keyword_False
	case "void":          return .Keyword_Void
	case "bool":          return .Keyword_Bool
	case "int":           return .Keyword_Int
	case "i8":            return .Keyword_I8
	case "i16":           return .Keyword_I16
	case "i32":           return .Keyword_I32
	case "i64":           return .Keyword_I64
	case "u8":            return .Keyword_U8
	case "u16":           return .Keyword_U16
	case "u32":           return .Keyword_U32
	case "u64":           return .Keyword_U64
	case "isize":         return .Keyword_Isize
	case "usize":         return .Keyword_Usize
	case "f32":           return .Keyword_F32
	case "f64":           return .Keyword_F64
	case "c_char":        return .Keyword_C_Char
	case "c_schar":       return .Keyword_C_Schar
	case "c_uchar":       return .Keyword_C_Uchar
	case "c_short":       return .Keyword_C_Short
	case "c_ushort":      return .Keyword_C_Ushort
	case "c_int":         return .Keyword_C_Int
	case "c_uint":        return .Keyword_C_Uint
	case "c_long":        return .Keyword_C_Long
	case "c_ulong":       return .Keyword_C_Ulong
	case "c_longlong":    return .Keyword_C_Longlong
	case "c_ulonglong":   return .Keyword_C_Ulonglong
	case "c_float":       return .Keyword_C_Float
	case "c_double":      return .Keyword_C_Double
	case "c_longdouble":  return .Keyword_C_Longdouble
	case "_":             return .Underscore
	}
	return .Identifier
}

append_token :: proc(
	stream: ^token.Stream,
	source_file: ^source.Source,
	kind: token.Kind,
	start, end: int,
	id := symbol.INVALID,
	diagnostic := source.INVALID_DIAGNOSTIC,
) {
	append(&stream.items, token.Token{
		kind=kind,
		span=source.Span{file=source_file.id, start=source.Offset(start), end=source.Offset(end)},
		symbol=id,
		diagnostic=diagnostic,
	})
}

lex :: proc(
	source_file: ^source.Source,
	diagnostics: ^source.Diagnostics,
	symbols: ^symbol.Table,
	allocator := context.allocator,
) -> token.Stream {
	stream: token.Stream
	stream.items.allocator = allocator
	bytes := transmute([]byte)source_file.text
	cursor := 0

	for cursor < len(bytes) {
		value := bytes[cursor]
		switch value {
		case ' ', '\t', '\r':
			cursor += 1
		case '\n':
			append_token(&stream, source_file, .Newline, cursor, cursor+1)
			cursor += 1
		case '#':
			for cursor < len(bytes) && bytes[cursor] != '\n' {
				cursor += 1
			}
		case ':':
			start := cursor
			cursor += 1
			if cursor < len(bytes) && bytes[cursor] == ':' {
				cursor += 1
				append_token(&stream, source_file, .Colon_Colon, start, cursor)
			} else {
				append_token(&stream, source_file, .Colon, start, cursor)
			}
		case '=':
			start := cursor
			cursor += 1
			if cursor < len(bytes) && bytes[cursor] == '=' {
				cursor += 1
				append_token(&stream, source_file, .Equal_Equal, start, cursor)
			} else {
				append_token(&stream, source_file, .Equal, start, cursor)
			}
		case '!':
			start := cursor
			cursor += 1
			if cursor < len(bytes) && bytes[cursor] == '=' {
				cursor += 1
				append_token(&stream, source_file, .Bang_Equal, start, cursor)
			} else {
				append_token(&stream, source_file, .Bang, start, cursor)
			}
		case '<':
			start := cursor
			cursor += 1
			if cursor < len(bytes) && bytes[cursor] == '=' {
				cursor += 1
				append_token(&stream, source_file, .Less_Equal, start, cursor)
			} else {
				append_token(&stream, source_file, .Less, start, cursor)
			}
		case '>':
			start := cursor
			cursor += 1
			if cursor < len(bytes) && bytes[cursor] == '=' {
				cursor += 1
				append_token(&stream, source_file, .Greater_Equal, start, cursor)
			} else {
				append_token(&stream, source_file, .Greater, start, cursor)
			}
		case '+':
			start := cursor
			cursor += 1
			if cursor < len(bytes) && bytes[cursor] == '=' {
				cursor += 1
				append_token(&stream, source_file, .Plus_Equal, start, cursor)
			} else {
				append_token(&stream, source_file, .Plus, start, cursor)
			}
		case '-':
			start := cursor
			cursor += 1
			if cursor < len(bytes) && bytes[cursor] == '=' {
				cursor += 1
				append_token(&stream, source_file, .Minus_Equal, start, cursor)
			} else {
				append_token(&stream, source_file, .Minus, start, cursor)
			}
		case '.':
			start := cursor
			cursor += 1
			if cursor < len(bytes) && bytes[cursor] == '.' {
				cursor += 1
				if cursor < len(bytes) && bytes[cursor] == '.' {
					cursor += 1
					append_token(&stream, source_file, .Ellipsis, start, cursor)
				} else if cursor < len(bytes) && bytes[cursor] == '=' {
					cursor += 1
					append_token(&stream, source_file, .Range_Inclusive, start, cursor)
				} else {
					append_token(&stream, source_file, .Range, start, cursor)
				}
			} else {
				append_token(&stream, source_file, .Dot, start, cursor)
			}
		case '@':
			append_token(&stream, source_file, .At, cursor, cursor+1)
			cursor += 1
		case '*':
			start := cursor
			cursor += 1
			if cursor < len(bytes) && bytes[cursor] == '=' {
				cursor += 1
				append_token(&stream, source_file, .Star_Equal, start, cursor)
			} else {
				append_token(&stream, source_file, .Star, start, cursor)
			}
		case '/':
			start := cursor
			cursor += 1
			if cursor < len(bytes) && bytes[cursor] == '=' {
				cursor += 1
				append_token(&stream, source_file, .Slash_Equal, start, cursor)
			} else {
				append_token(&stream, source_file, .Slash, start, cursor)
			}
		case '&':
			append_token(&stream, source_file, .Ampersand, cursor, cursor+1)
			cursor += 1
		case '^':
			append_token(&stream, source_file, .Caret, cursor, cursor+1)
			cursor += 1
		case '?':
			append_token(&stream, source_file, .Question, cursor, cursor+1)
			cursor += 1
		case '[':
			append_token(&stream, source_file, .Left_Bracket, cursor, cursor+1)
			cursor += 1
		case ']':
			append_token(&stream, source_file, .Right_Bracket, cursor, cursor+1)
			cursor += 1
		case '(':
			append_token(&stream, source_file, .Left_Paren, cursor, cursor+1)
			cursor += 1
		case ')':
			append_token(&stream, source_file, .Right_Paren, cursor, cursor+1)
			cursor += 1
		case '{':
			append_token(&stream, source_file, .Left_Brace, cursor, cursor+1)
			cursor += 1
		case '}':
			append_token(&stream, source_file, .Right_Brace, cursor, cursor+1)
			cursor += 1
		case ',':
			append_token(&stream, source_file, .Comma, cursor, cursor+1)
			cursor += 1
		case '|':
			append_token(&stream, source_file, .Pipe, cursor, cursor+1)
			cursor += 1
		case '"':
			start := cursor
			cursor += 1
			valid := true
			for cursor < len(bytes) && bytes[cursor] != '"' && bytes[cursor] != '\n' {
				if bytes[cursor] == '\\' {
					cursor += 1
					if cursor >= len(bytes) ||
					   (bytes[cursor] != '\\' && bytes[cursor] != '"' && bytes[cursor] != 'n' &&
					    bytes[cursor] != 'r' && bytes[cursor] != 't' && bytes[cursor] != '0') {
						source.add(
							diagnostics,
							source.Span{file=source_file.id, start=source.Offset(max(cursor-1, start)), end=source.Offset(min(cursor+1, len(bytes)))},
							"strings only support '\\\\', '\\\"', '\\n', '\\r', '\\t', and '\\0' escapes",
						)
						valid = false
					}
				}
				if cursor < len(bytes) && bytes[cursor] != '\n' {
					cursor += 1
				}
			}
			if cursor < len(bytes) && bytes[cursor] == '"' {
				cursor += 1
				append_token(&stream, source_file, .String if valid else .Invalid, start, cursor)
			} else {
				id := source.add(
					diagnostics,
					source.Span{file=source_file.id, start=source.Offset(start), end=source.Offset(cursor)},
					"unterminated import string",
				)
				append_token(&stream, source_file, .Invalid, start, cursor, diagnostic=id)
			}
		case ';':
			append_token(&stream, source_file, .Semicolon, cursor, cursor+1)
			cursor += 1
		case '\'':
			start := cursor
			cursor += 1
			for cursor < len(bytes) && bytes[cursor] != '\'' && bytes[cursor] != '\n' {
				if bytes[cursor] == '\\' && cursor+1 < len(bytes) {
					cursor += 1
				}
				cursor += 1
			}
			if cursor < len(bytes) && bytes[cursor] == '\'' {
				cursor += 1
				append_token(&stream, source_file, .Character, start, cursor)
			} else {
				id := source.add(
					diagnostics,
					source.Span{file=source_file.id, start=source.Offset(start), end=source.Offset(cursor)},
					"unterminated character literal",
				)
				append_token(&stream, source_file, .Invalid, start, cursor, diagnostic=id)
			}
		case:
			if value >= '0' && value <= '9' {
				start := cursor
				for cursor < len(bytes) && bytes[cursor] >= '0' && bytes[cursor] <= '9' {
					cursor += 1
				}
				kind := token.Kind.Integer
				if cursor+1 < len(bytes) && bytes[cursor] == '.' &&
				   bytes[cursor+1] != '.' && bytes[cursor+1] >= '0' && bytes[cursor+1] <= '9' {
					kind = .Float
					cursor += 1
					for cursor < len(bytes) && bytes[cursor] >= '0' && bytes[cursor] <= '9' {
						cursor += 1
					}
				}
				append_token(&stream, source_file, kind, start, cursor)
			} else if is_identifier_start(value) {
				start := cursor
				for cursor < len(bytes) && is_identifier_continue(bytes[cursor]) {
					cursor += 1
				}
				text := source_file.text[start:cursor]
				kind := keyword_kind(text)
				id := symbol.INVALID
				if kind == .Identifier || kind == .Underscore {
					id = symbol.intern(symbols, text)
				}
				append_token(&stream, source_file, kind, start, cursor, id)
			} else {
				id := source.addf(
					diagnostics,
					source.Span{file=source_file.id, start=source.Offset(cursor), end=source.Offset(cursor+1)},
					"invalid source byte 0x%02x",
					value,
				)
				append_token(&stream, source_file, .Invalid, cursor, cursor+1, diagnostic=id)
				cursor += 1
			}
		}
	}

	append_token(&stream, source_file, .Eof, len(bytes), len(bytes))
	return stream
}