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

import "@source/strpool"

ScanError :: struct {
    code ErrorCode
    end usize
}

ErrorCode :: enum {
    invalid_character,
    float_must_end_with_digit,
    unterminated_string,
}

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 {
    token TokenId
    code ErrorCode
}

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, program []u8) void ! (mem.AllocError | strpool.InternError) {
    tokens :: &state.tokens
    diagnostics :: &state.diagnostics

    cursor := 0
    while cursor < program.len {
        char :: program[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 < program.len and program[cursor] != '\n') cursor += 1
            cursor += 1 # also skip newline
            continue
        }

        # identifiers and keywords
        if is_alpha(char) or char == '_' {
            start :: cursor
            result :: scan_ident(start, program)
            try add_token(tokens, Token{
                kind = result.kind,
                start = start,
                str_id = result.str_id,
            })
            cursor = result.end
            continue
        }

        # numeric literals
        if is_digit(char) {
            start :: cursor
            result :: scan_number(start, program) catch |err| {
                token :: token_id(tokens.items.len)
                try add_token(tokens, Token{ kind = .invalid, start = start })
                try arraylist.append(diagnostics, Diagnostic{ token = token, code = err.code })
                cursor = err.end
                continue
            }

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

        # string literals
        if char == '"' {
            start :: cursor
            result :: scan_string(start, program) catch |err| {
                token :: token_id(tokens.items.len)
                try add_token(tokens, Token{ kind = .invalid, start = start })
                try arraylist.append(diagnostics, Diagnostic{ token = token, code = err.code })
                cursor = err.end
                continue
            }

            try add_token(tokens, Token{ kind = .string, start = start })
            cursor = result.end
            continue
        }

        # immutable assignment or single colon
        if char == ':' {
            if cursor + 1 < program.len and program[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
        }

        # mutable reassignment
        if char == '=' {
            try add_token(tokens, Token{ kind = .equal, 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 })
}

ScanIdentResult :: struct {
    end usize
    kind TokenKind
    str_id strpool.StringId
}

scan_ident proc(start usize, program []u8) ScanIdentResult {
    cursor := start + 1

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

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

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

    return ScanIdentResult{
        end = cursor,
        kind = kind,
        str_id = str_id,
    }
}

ScanNumResult :: struct {
    end usize
    has_decimal bool
}

scan_number proc(start usize, program []u8) ScanNumResult ! ScanError {
    cursor := start
    has_decimal := false

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

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

    # assert non-terminating decimal
    if has_decimal and (cursor >= program.len or !is_digit(program[cursor])) {
        return ScanError{
            code = .float_must_end_with_digit,
            end = cursor + 1,
        }
    }

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

    return ScanNumResult{
        end = cursor,
        has_decimal = has_decimal,
    }
}

ScanStrResult :: struct { end usize }

scan_string proc(start usize, program []u8) ScanStrResult ! ScanError {
    cursor := start + 1 # skip first `"`

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

    # assert string terminal
    if (cursor >= program.len or program[cursor] != '"') return ScanError{
        code = .unterminated_string,
        end = cursor,
    }

    return ScanStrResult{
        end = cursor + 1, # skip last `"`
    }
}

@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)
}