const TT_LEFT_PAREN    = 1
const TT_RIGHT_PAREN   = 2
const TT_LEFT_BRACKET  = 3
const TT_RIGHT_BRACKET = 4
const TT_COMMA         = 5
const TT_PLUS          = 6
const TT_PLUS_EQUAL    = 7
const TT_MINUS         = 8
const TT_MINUS_EQUAL   = 9
const TT_STAR          = 10
const TT_SLASH         = 11
const TT_MOD           = 12
const TT_XOR           = 13
const TT_BANG          = 14
const TT_COLON         = 15
const TT_BIT_AND       = 16
const TT_BIT_OR        = 17
const TT_LOGICAL_AND   = 18
const TT_LOGICAL_OR    = 19
const TT_DOUBLE_DOT    = 20
const TT_SHIFT_LEFT    = 21
const TT_SHIFT_RIGHT   = 22
const TT_ARROW         = 23
const TT_TILDE         = 24

const TT_EQUAL         = 25
const TT_DOUBLE_EQUAL  = 26
const TT_NOT_EQUAL     = 27
const TT_GREATER       = 28
const TT_GREATER_EQUAL = 29
const TT_LESS          = 30
const TT_LESS_EQUAL    = 31

const TT_IDENTIFIER     = 32
const TT_STRING_LITERAL = 33
const TT_CHAR_LITERAL   = 34
const TT_INT_LITERAL    = 35
const TT_TRUE           = 36
const TT_FALSE          = 37

const TT_KW_CONST    = 38
const TT_KW_IF       = 39
const TT_KW_ELSE     = 40
const TT_KW_WHILE    = 41
const TT_KW_FOR      = 42
const TT_KW_IN       = 43
const TT_KW_FUNC     = 44
const TT_KW_RETURN   = 45
const TT_KW_BREAK    = 46
const TT_KW_CONTINUE = 47
const TT_KW_EXTERN   = 48
const TT_KW_EXPORT   = 49
const TT_KW_STRUCT   = 50
const TT_KW_NEW      = 51
const TT_KW_AS       = 52

const TT_INDENT = 53
const TT_DEDENT = 54
const TT_EOF    = 55

struct Token
    token_type: i64
    lexeme: str
    loc_filename: str
    loc_line: i64
    loc_column: i64

struct Tokenizer
    source: str
    source_len: i64
    tokens: Array
    indent_stack: Array
    current_indent: i64
    start: i64
    current: i64
    loc_filename: str
    loc_line: i64
    loc_column: i64

func Tokenizer.init[t: Tokenizer, filename: str, source: str] : void
    t->source = source
    t->source_len = str.len(source)
    t->tokens = []
    t->indent_stack = [0]
    t->loc_filename = filename
    t->loc_line = 1
    t->loc_column = 1

func Tokenizer.tokenize[t: Tokenizer] : Array
    while !t->eof()
        t->start = t->current
        t->scan_token()

    t->add_token(TT_EOF, "")
    return t->tokens

func Tokenizer.scan_token[t: Tokenizer] : void
    c := t->advance()
    if c == '('
        t->add_token(TT_LEFT_PAREN, "(")
    else if c == ')'
        t->add_token(TT_RIGHT_PAREN, ")")
    else if c == '['
        t->add_token(TT_LEFT_BRACKET, "[")
    else if c == ']'
        t->add_token(TT_RIGHT_BRACKET, "]")
    else if c == '+'
        if t->match_char('=')
            t->add_token(TT_PLUS_EQUAL, "+=")
        else
            t->add_token(TT_PLUS, "+")
    else if c == '*'
        t->add_token(TT_STAR, "*")
    else if c == ','
        t->add_token(TT_COMMA, ",")
    else if c == '%'
        t->add_token(TT_MOD, "%")
    else if c == '^'
        t->add_token(TT_XOR, "^")
    else if c == ':'
        t->add_token(TT_COLON, ":")
    else if c == '~'
        t->add_token(TT_TILDE, "~")
    else if c == '-'
        if t->match_char('=')
            t->add_token(TT_MINUS_EQUAL, "-=")
        else if t->match_char('>')
            t->add_token(TT_ARROW, "->")
        else
            t->add_token(TT_MINUS, "-")
    else if c == '.'
        if t->match_char('.')
            t->add_token(TT_DOUBLE_DOT, "..")
        else
            t->error("expected '.' after '.'")
    else if c == '/'
        if t->match_char('/')
            while !t->eof() && t->peek() != '\n'
                t->advance()
        else
            t->add_token(TT_SLASH, "/")
    else if c == '&'
        if t->match_char('&')
            t->add_token(TT_LOGICAL_AND, "&&")
        else
            t->add_token(TT_BIT_AND, "&")
    else if c == '|'
        if t->match_char('|')
            t->add_token(TT_LOGICAL_OR, "||")
        else
            t->add_token(TT_BIT_OR, "|")
    else if c == '!'
        if t->match_char('=')
            t->add_token(TT_NOT_EQUAL, "!=")
        else
            t->add_token(TT_BANG, "!")
    else if c == '='
        if t->match_char('=')
            t->add_token(TT_DOUBLE_EQUAL, "==")
        else
            t->add_token(TT_EQUAL, "=")
    else if c == '>'
        if t->match_char('>')
            t->add_token(TT_SHIFT_RIGHT, ">>")
        else if t->match_char('=')
            t->add_token(TT_GREATER_EQUAL, ">=")
        else
            t->add_token(TT_GREATER, ">")
    else if c == '<'
        if t->match_char('<')
            t->add_token(TT_SHIFT_LEFT, "<<")
        else if t->match_char('=')
            t->add_token(TT_LESS_EQUAL, "<=")
        else
            t->add_token(TT_LESS, "<")
    else if c == '\''
        if t->eof()
            t->error("unterminated char literal")

        t->match_char('\\')
        t->advance()

        if !t->match_char('\'')
            t->error("expected ' after char literal")

        lexeme := t->source->substr(t->start, t->current - t->start)
        t->add_token(TT_CHAR_LITERAL, lexeme)
    else if c == '"'
        start_filename := t->loc_filename
        start_line := t->loc_line
        start_column := t->loc_column

        while !t->eof()
            if t->peek() == '\\'
                t->advance()
            else if t->peek() == '"'
                break
            else if t->peek() == '\n'
                t->loc_line += 1
                t->loc_column = 0

            t->advance()

        if t->eof()
            t->loc_filename = start_filename
            t->loc_line = start_line
            t->loc_column = start_column
            t->error("unterminated string")

        t->advance()
        lexeme := t->source->substr(t->start, t->current - t->start)
        t->add_token(TT_STRING_LITERAL, lexeme)
    else if c == ' ' || c == '\r'
        // ignore
    else if c == '\n'
        t->loc_line += 1
        t->loc_column = 1
        t->handle_indentation()
    else if c >= '0' && c <= '9'
        t->scan_number()
    else if (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || c == '_'
        t->scan_identifier()
    else
        t->error("unexpected character")

func Tokenizer.handle_indentation[t: Tokenizer] : void
    if t->peek() == '\n'
        return 0
    
    new_indent := t->count_indentation()

    if new_indent > t->current_indent
        t->indent_stack->push(new_indent)
        t->add_token(TT_INDENT, "")
    else if new_indent < t->current_indent
        while t->indent_stack->size > 0 && t->indent_stack->last() > new_indent
            t->indent_stack->pop()
            t->add_token(TT_DEDENT, "")
        if t->indent_stack->size == 0 || t->indent_stack->last() != new_indent
            t->error("invalid indentation")
    
    t->current_indent = new_indent

func Tokenizer.count_indentation[t: Tokenizer] : i64
    count := 0
    while t->peek() == ' '
        count += 1
        t->advance()
    return count

func Tokenizer.scan_number[t: Tokenizer] : void
    if t->match_char('x')
        while t->peek()->is_hex_digit()
            t->advance()
    else if t->match_char('o')
        while '0' <= t->peek() && t->peek() <= '7'
            t->advance()
    else
        while t->peek()->is_digit()
            t->advance()
    
    lexeme := t->source->substr(t->start, t->current - t->start)
    t->add_token(TT_INT_LITERAL, lexeme)

func str.match_at[s: str, start: i64, keyword: str] : bool
    i := 0
    while keyword[i] != 0
        if s[start + i] != keyword[i]
            return false
        i += 1
    return true

func Tokenizer.scan_identifier[t: Tokenizer] : void
    while t->peek()->is_alphanumeric() || t->peek() == '_' || t->peek() == '.'
        t->advance()

    len := t->current - t->start

    if len == 5 && t->source->match_at(t->start, "const")
        t->add_token(TT_KW_CONST, "const")
    else if len == 2 && t->source->match_at(t->start, "if")
        t->add_token(TT_KW_IF, "if")
    else if len == 4 && t->source->match_at(t->start, "else")
        t->add_token(TT_KW_ELSE, "else")
    else if len == 5 && t->source->match_at(t->start, "while")
        t->add_token(TT_KW_WHILE, "while")
    else if len == 3 && t->source->match_at(t->start, "for")
        t->add_token(TT_KW_FOR, "for")
    else if len == 2 && t->source->match_at(t->start, "in")
        t->add_token(TT_KW_IN, "in")
    else if len == 4 && t->source->match_at(t->start, "func")
        t->add_token(TT_KW_FUNC, "func")
    else if len == 6 && t->source->match_at(t->start, "return")
        t->add_token(TT_KW_RETURN, "return")
    else if len == 5 && t->source->match_at(t->start, "break")
        t->add_token(TT_KW_BREAK, "break")
    else if len == 8 && t->source->match_at(t->start, "continue")
        t->add_token(TT_KW_CONTINUE, "continue")
    else if len == 6 && t->source->match_at(t->start, "extern")
        t->add_token(TT_KW_EXTERN, "extern")
    else if len == 6 && t->source->match_at(t->start, "export")
        t->add_token(TT_KW_EXPORT, "export")
    else if len == 6 && t->source->match_at(t->start, "struct")
        t->add_token(TT_KW_STRUCT, "struct")
    else if len == 3 && t->source->match_at(t->start, "new")
        t->add_token(TT_KW_NEW, "new")
    else if len == 2 && t->source->match_at(t->start, "as")
        t->add_token(TT_KW_AS, "as")
    else if len == 4 && t->source->match_at(t->start, "true")
        t->add_token(TT_TRUE, "true")
    else if len == 5 && t->source->match_at(t->start, "false")
        t->add_token(TT_FALSE, "false")
    else
        lexeme := t->source->substr(t->start, len)
        t->add_token(TT_IDENTIFIER, lexeme)

func Tokenizer.match_char[t: Tokenizer, expected: u8] : bool
    if t->eof() || t->peek() != expected
        return false
    t->current += 1
    t->loc_column += 1
    return true

func Tokenizer.add_token[t: Tokenizer, token_type: i64, lexeme: str] : void
    if token_type == TT_STRING_LITERAL || token_type == TT_CHAR_LITERAL
        lexeme = t->unescape(lexeme)

    token := new* Token
    token->token_type = token_type
    token->lexeme = lexeme
    token->loc_filename = t->loc_filename
    token->loc_line = t->loc_line
    token->loc_column = t->loc_column
    t->tokens->push(token)

func Tokenizer.unescape[t: Tokenizer, s: str] : str
    b := new str.Builder
    i := 0
    len := s->len()
    while i < len
        if s[i] == '\\'
            i += 1
            if i >= len
                t->error("unexpected end of escape sequence")

            if s[i] == 'n'
                b->append_char('\n')
            else if s[i] == 'r'
                b->append_char('\r')
            else if s[i] == 't'
                b->append_char('\t')
            else if s[i] == '0'
                b->append_char('\0')
            else if s[i] == '\\'
                b->append_char('\\')
            else if s[i] == '\''
                b->append_char('\'')
            else if s[i] == '"'
                b->append_char('"')
            else
                t->error("unknown escape sequence")
        else
            b->append_char(s[i])
        i += 1

    result := b->build()
    b->destroy()
    return result

func Tokenizer.advance[t: Tokenizer] : u8
    c := t->source[t->current]
    t->current += 1
    t->loc_column += 1
    return c

func Tokenizer.peek[t: Tokenizer] : u8
    if t->eof()
        return '\0'
    return t->source[t->current]

func Tokenizer.eof[t: Tokenizer] : bool
    return t->current >= t->source_len

func Tokenizer.error[t: Tokenizer, msg: str] : void
    io.print(t->loc_filename)
    io.print(":")
    io.print_i64(t->loc_line)
    io.print(":")
    io.print_i64(t->loc_column)
    io.print(" ERROR: ")
    io.println(msg)
    os.exit(1)

func main[argc: i64, argv: ptr] : i64
    path := mem.read64(argv + 8) as str
    ~source, ok := io.read_text_file(path)
    if !ok
        panic("file not found")
    
    filename := os.basename(path)

    tokenizer := new Tokenizer
    tokenizer->init(filename, source)

    tokens := tokenizer->tokenize()
    for i in 0..tokens->size
        t := tokens->nth(i) as Token
        io.printf("Token(token_type=%d, lexeme=\"%s\", loc=\"%s:%d:%d\")\n", t->token_type, t->lexeme, t->loc_filename, t->loc_line, t->loc_column)