import "@std"
import "@std/mem"
import "@std/strmap"
import "@std/enums/enummap"
import "@std/arraylist"
import "@std/debug"

import "@source/strpool"

ErrorCode :: enum {
    invalid_character,
    float_must_end_with_digit,
}

ErrorDetails :: struct {
    name []u8
    message []u8
}

error_msg_map std.EnumMap(ErrorCode, ErrorDetails) :: enummap.init({
    invalid_character = ErrorDetails {
        name = "L0",
        message = "invalid character",
    },
    float_must_end_with_digit = ErrorDetails {
        name = "L1",
        message = "float must end with a digit",
    },
})

keywords std.StringMap(TokenKind) :: strmap.init([
    { "proc", .proc },
    { "return", .return },
    { "if", .if },
    { "for", .for },
    { "else", .else },
    { "while", .while },
])

TokenId :: distinct u32

Diagnostic :: struct {
    code ErrorCode
    token TokenId
}

State :: struct {
    tokens std.ArrayList(Token)
    diagnostics std.ArrayList(Diagnostic)
}

init proc(allocator mem.Allocator) State {
    return State {
        tokens = arraylist.init(allocator),
        diagnostics = arraylist.init(allocator),
    }
}

deinit proc(state @mut State) void {
    arraylist.deinit(&state.tokens)
    arraylist.deinit(&state.diagnostics)
}

scan proc(state @mut State, input []u8) void ! (mem.AllocError | strpool.InternError) {
    tokens :: &state.tokens
    diagnostics :: &state.diagnostics

    cursor usize = 0
    while cursor < input.len {
        char :: input[cursor]

        # whitespace
        if char == '\n' {
            try add_token(tokens, Token{ kind = .newline, start = cursor })
            cursor += 1
            continue
        } else if is_whitespace(char) {
            cursor += 1
            continue
        }

        # comments
        if char == '#' {
            while (cursor < input.len and input[cursor] != '\n') cursor += 1
            cursor += 1 # also skip newline
            continue
        }

        # identifiers and keywords
        if is_alpha(char) or char == '_' {
            start :: cursor
            cursor += 1

            # scan whole identifier
            while (cursor < input.len and (
                is_alpha(input[cursor]) or
                is_digit(input[cursor]) or
                input[cursor] == '_'
            )) cursor += 1

            kind :: strmap.get(&keywords, input[start..cursor]) orelse .ident

            # don't intern keywords (already O(1) lookup via token kind)
            str_id :: if (kind == .ident)
                try strpool.intern(&strpool.strings, input[start..cursor])
            else
                strpool.NO_ID

            try add_token(tokens, Token{
                kind = kind,
                start = start,
                str_id = str_id
            })
            continue
        }

        # numeric literals
        if is_digit(char) {
            start :: cursor
            has_decimal bool = false

            # scan integer part
            while (cursor < input.len and is_digit(input[cursor])) cursor += 1

            # check for decimal point
            if cursor < input.len and input[cursor] == '.' {
                has_decimal = true
                cursor += 1
            }

            # assert that decimals follow the decimal point
            if has_decimal and (cursor >= input.len or !is_digit(input[cursor])) {
                token :: token_id(tokens.items.len)
                try arraylist.append(diagnostics, Diagnostic{ token = token, code = .float_must_end_with_digit })
                try add_token(tokens, Token{ kind = .invalid, start = start })
                continue
            }

            # scan decimal part
            while (cursor < input.len and is_digit(input[cursor])) cursor += 1

            kind :: if (has_decimal) .float else .int
            try add_token(tokens, Token{ kind = kind, start = start })

            continue
        }

        # string literals
        if char == '"' {
            start :: cursor
            cursor += 1

            while cursor < input.len and input[cursor] != '"' and input[cursor] != '\n' : cursor += 1 {
                # ignore escaped characters
                if (input[cursor] == '\\' and cursor + 1 < input.len) cursor += 1
            }

            if cursor < input.len and input[cursor] == '"' {
                cursor += 1
                try add_token(tokens, Token{ kind = .string, start = start })
            } else {
                try add_token(tokens, Token{ kind = .invalid, start = start })
            }

            continue
        }

        # mutable assignment
        if char == '=' {
            try add_token(tokens, Token{ kind = .equal, start = cursor })
            cursor += 1
            continue
        }

        # immutable assignment or single colon
        if char == ':' {
            if cursor + 1 < input.len and input[cursor + 1] == ':' {
                try add_token(tokens, Token{ kind = .double_colon, start = cursor })
                cursor += 2
                continue
            }

            try add_token(tokens, Token{ kind = .colon, start = cursor })
            cursor += 1
            continue
        }

        # parentheses
        if char == '(' {
            try add_token(tokens, Token{ kind = .open_paren, start = cursor })
            cursor += 1
            continue
        } else if char == ')' {
            try add_token(tokens, Token{ kind = .close_paren, start = cursor })
            cursor += 1
            continue
        }

        # curly braces
        if char == '{' {
            try add_token(tokens, Token{ kind = .open_curly, start = cursor })
            cursor += 1
            continue
        } else if char == '}' {
            try add_token(tokens, Token{ kind = .close_curly, start = cursor })
            cursor += 1
            continue
        }

        # invalid character
        token :: token_id(tokens.items.len)
        try arraylist.append(diagnostics, Diagnostic{ token = token, code = .invalid_character })
        try add_token(tokens, Token{ kind = .invalid, start = cursor })
        cursor += 1
    }

    try add_token(tokens, Token{ kind = .eof, start = cursor })
}

#! returns the cursor position after scanning a number.
scan_number proc(start usize, input []u8) usize {
    debug.assert(is_digit(input[start]))

    has_decimal bool = false
    cursor usize = start

    # scan integer part
    while cursor < input.len and is_digit(input[cursor]) : cursor += 1 {}

    # check for decimal point
    if cursor < input.len and input[cursor] == '.' {
        has_decimal = true
        cursor += 1
    }

    # assert that decimals follow the decimal point
    debug.assert(!(has_decimal and (cursor >= input.len or !is_digit(input[cursor]))))

    # scan decimal part
    while cursor < input.len and is_digit(input[cursor]) : cursor += 1 {}

    return cursor
}

#! returns the cursor position after scanning a string.
scan_string proc(start usize, input []u8) usize {
    debug.assert(input[start] == '"')

    cursor usize = start + 1
    while cursor < input.len and input[cursor] != '"' and input[cursor] != '\n' : cursor += 1 {
        # ignore escaped characters
        if (input[cursor] == '\\' and cursor + 1 < input.len) cursor += 1
    }

    debug.assert(cursor < input.len and input[cursor] == '"')
    cursor += 1

    return cursor
}

hide is_whitespace proc(char u8) bool {
    return char == ' ' or char == '\t' or char == '\n' or char == '\r'
}

hide is_alpha proc(char u8) bool {
    return match char {
        'a'..='z', 'A'..='Z': true
        else: false
    }
}

hide is_digit proc(char u8) bool {
    return match char {
        '0'..='9': true
        else: false
    }
}

hide add_token proc(tokens @mut std.ArrayList(Token), token Token) void ! mem.AllocError {
    _ = token_id(tokens.items.len)
    try arraylist.append(tokens, token)
}

hide token_id proc(idx uint) TokenId {
    debug.assert(u64(idx) < u64(maxval!(TokenId)))
    return TokenId(idx)
}