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

import "@source/diagnostics"
import "@source/strpool"

@hide:file
ScanError :: struct {
    code ErrorCode
    end usize
}

@hide:file
ErrorCode :: enum {
    float_must_end_with_digit,
    unterminated_string,
    unterminated_character,
}
SourceError :: enum { source_too_large }

@hide:file
keywords strmap.StaticStringMap(TokenKind) :: strmap.init([
    { "proc", .proc },
    { "c_func", .c_func },
    { "null", .null },
    { "undefined", .undefined },
    { "return", .return },
    { "yield", .yield },
    { "break", .break },
    { "continue", .continue },
    { "unreachable", .unreachable },
    { "try", .try },
    { "catch", .catch },
    { "defer", .defer },
    { "errdefer", .errdefer },
    { "if", .if },
    { "for", .for },
    { "inline", .inline },
    { "else", .else },
    { "while", .while },
    { "true", .true },
    { "false", .false },
    { "and", .and },
    { "or", .or },
    { "xor", .xor },
    { "mut", .mut },
    { "struct", .struct },
    { "c_struct", .c_struct },
    { "opaque", .opaque },
    { "enum", .enum },
    { "union", .union },
    { "distinct", .distinct },
    { "alias", .alias },
    { "match", .match },
    { "import", .import },
    { "test", .test },
])

#! owns token storage and retains its allocator; release with deinit.
#! identifier IDs belong to the compilation's separately owned string pool.
State :: struct {
    tokens arraylist.ArrayList(Token)
    source_id diagnostics.SourceId = diagnostics.SourceId(0)
}

#! retains `allocator` for scanned tokens; the caller owns the returned State.
init proc(allocator mem.Allocator) State {
    return State{ tokens = arraylist.init(allocator) }
}

#! releases token storage with the retained allocator, without releasing identifier strings.
deinit proc(state @mut State) void {
    arraylist.deinit(&state.tokens)
}

#! borrows `program`, `pool`, and `diagnostic_store`; tokens remain state-owned.
#! identifier storage uses the pool's allocator; keep the pool alive while using their IDs.
#! `source_id` is copied into every token and scanner diagnostic span.
scan proc(
    state @mut State,
    program []u8,
    pool @mut strpool.StringPool,
    diagnostic_store @mut diagnostics.Store,
) void ! (mem.AllocError | strpool.InternError | diagnostics.CapacityError | SourceError) {
    if (program.len > usize(maxval!(u32))) return .source_too_large
    source :: state.source_id
    cursor := 0
    while cursor < program.len {
        char :: program[cursor]

        if char == '\n' {
            try add_token(state, Token{
                span = make_span(source, cursor, cursor + 1),
                kind = .newline,
            })
            cursor += 1
            continue
        }
        if is_whitespace(char) {
            cursor += 1
            continue
        }

        if char == '#' {
            while (cursor < program.len and program[cursor] != '\n') cursor += 1
            continue
        }

        if is_alpha(char) or char == '_' {
            start :: cursor
            result :: try scan_ident(start, program, pool)
            compound :: result.kind == .xor and result.end < program.len and program[result.end] == '='
            extra :: if (compound) usize(1) else 0
            end :: result.end + extra
            kind :: if (compound) TokenKind.xor_equal else result.kind
            try add_token(state, Token{
                span = make_span(source, start, end),
                kind = kind,
                data = TokenData{ str_id = result.str_id },
            })
            cursor = end
            continue
        }

        if is_digit(char) {
            start :: cursor
            member :: (
                state.tokens.items.len > 0
                and state.tokens.items[state.tokens.items.len - 1].kind == .dot
                and usize(state.tokens.items[state.tokens.items.len - 1].span.end) == start
            )
            result :: scan_number(start, program, !member) catch |err| {
                id :: try diagnostics.add(
                    diagnostic_store,
                    diagnostic_code(err.code),
                    make_span(source, start, err.end),
                )
                try add_token(state, Token{
                    span = make_span(source, start, err.end),
                    kind = .invalid,
                    data = TokenData{ diagnostic = id },
                })
                cursor = err.end
                continue
            }
            kind :: if (result.has_decimal) .float else .int
            try add_token(state, Token{
                span = make_span(source, start, result.end),
                kind = kind,
            })
            cursor = result.end
            continue
        }

        if (char == '"' or char == '\'') {
            start :: cursor
            result :: scan_quoted(start, program, char) catch |err| {
                id :: try diagnostics.add(
                    diagnostic_store,
                    diagnostic_code(err.code),
                    make_span(source, start, err.end),
                )
                try add_token(state, Token{
                    span = make_span(source, start, err.end),
                    kind = .invalid,
                    data = TokenData{ diagnostic = id },
                })
                cursor = err.end
                continue
            }
            kind :: if (char == '\'') TokenKind.character else .string
            try add_token(state, Token{
                span = make_span(source, start, result.end),
                kind = kind,
            })
            cursor = result.end
            continue
        }

        if (char == '.' and cursor + 1 < program.len and program[cursor + 1] == '.') {
            inclusive :: cursor + 2 < program.len and program[cursor + 2] == '='
            variadic :: cursor + 2 < program.len and program[cursor + 2] == '.'
            width :: if (inclusive or variadic) 3 else 2
            range_kind :: if (inclusive) .range_inclusive else if (variadic) .ellipsis else .range
            try add_token(state, Token{
                span = make_span(source, cursor, cursor + width),
                kind = range_kind,
            })
            cursor += width
            continue
        }
        if (char == '.') {
            try add_token(state, Token{
                span = make_span(source, cursor, cursor + 1),
                kind = .dot,
            })
            cursor += 1
            continue
        }

        if char == ':' {
            if cursor + 1 < program.len and program[cursor + 1] == ':' {
                try add_token(state, Token{
                    span = make_span(source, cursor, cursor + 2),
                    kind = .double_colon,
                })
                cursor += 2
                continue
            }
            if cursor + 1 < program.len and program[cursor + 1] == '=' {
                try add_token(state, Token{
                    span = make_span(source, cursor, cursor + 2),
                    kind = .colon_equal,
                })
                cursor += 2
                continue
            }
            try add_token(state, Token{
                span = make_span(source, cursor, cursor + 1),
                kind = .colon,
            })
            cursor += 1
            continue
        }

        if char == '+' {
            if cursor + 1 < program.len and program[cursor + 1] == '=' {
                try add_token(state, Token{
                    span = make_span(source, cursor, cursor + 2),
                    kind = .plus_equal,
                })
                cursor += 2
                continue
            }
            try add_token(state, Token{
                span = make_span(source, cursor, cursor + 1),
                kind = .plus,
            })
            cursor += 1
            continue
        }
        if char == '-' {
            if cursor + 1 < program.len and program[cursor + 1] == '=' {
                try add_token(state, Token{
                    span = make_span(source, cursor, cursor + 2),
                    kind = .minus_equal,
                })
                cursor += 2
                continue
            }
            try add_token(state, Token{
                span = make_span(source, cursor, cursor + 1),
                kind = .minus,
            })
            cursor += 1
            continue
        }
        if char == '*' {
            if cursor + 1 < program.len and program[cursor + 1] == '=' {
                try add_token(state, Token{
                    span = make_span(source, cursor, cursor + 2),
                    kind = .star_equal,
                })
                cursor += 2
                continue
            }
            try add_token(state, Token{
                span = make_span(source, cursor, cursor + 1),
                kind = .star,
            })
            cursor += 1
            continue
        }
        if char == '/' {
            if cursor + 1 < program.len and program[cursor + 1] == '=' {
                try add_token(state, Token{
                    span = make_span(source, cursor, cursor + 2),
                    kind = .slash_equal,
                })
                cursor += 2
                continue
            }
            try add_token(state, Token{
                span = make_span(source, cursor, cursor + 1),
                kind = .slash,
            })
            cursor += 1
            continue
        }

        if (
            (char == '<' or char == '>') and cursor + 1 < program.len
            and program[cursor + 1] == char
        ) {
            saturating :: char == '<' and cursor + 2 < program.len and program[cursor + 2] == '|'
            width :: if (saturating) usize(3) else 2
            compound :: cursor + width < program.len and program[cursor + width] == '='
            kind TokenKind := .shift_right
            if (saturating) {
                kind = if (compound) .shift_left_saturating_equal else .shift_left_saturating
            } else if (char == '<') {
                kind = if (compound) .shift_left_equal else .shift_left
            } else if (compound) kind = .shift_right_equal
            extra :: if (compound) usize(1) else 0
            end :: cursor + width + extra
            try add_token(state, Token{
                span = make_span(source, cursor, end),
                kind = kind,
            })
            cursor = end
            continue
        }

        if (cursor + 1 < program.len and program[cursor + 1] == '=') {
            comparison ?TokenKind :: match char {
                '=': .equal_equal
                '!': .not_equal
                '<': .less_equal
                '>': .greater_equal
                '&': .ampersand_equal
                '|': .pipe_equal
                else: null
            }
            if (comparison) |value| {
                try add_token(state, Token{
                    span = make_span(source, cursor, cursor + 2),
                    kind = value,
                })
                cursor += 2
                continue
            }
        }

        kind ?TokenKind :: match char {
            '=': .equal
            '!': .bang
            '~': .tilde
            '<': .less
            '>': .greater
            ',': .comma
            '|': .pipe
            '@': .at
            '$': .dollar
            '&': .ampersand
            '^': .caret
            '?': .question
            ';': .semicolon
            '(': .open_paren
            ')': .close_paren
            '[': .open_bracket
            ']': .close_bracket
            '{': .open_curly
            '}': .close_curly
            else: null
        }
        if (kind) |value| {
            try add_token(state, Token{
                span = make_span(source, cursor, cursor + 1),
                kind = value,
            })
            cursor += 1
            continue
        }

        id :: try diagnostics.add(
            diagnostic_store,
            .invalid_character,
            make_span(source, cursor, cursor + 1),
        )
        try add_token(state, Token{
            span = make_span(source, cursor, cursor + 1),
            kind = .invalid,
            data = TokenData{ diagnostic = id },
        })
        cursor += 1
    }

    try add_token(state, Token{
        span = make_span(source, cursor, cursor),
        kind = .eof,
    })
}

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

@hide:file
scan_ident proc(
    start usize,
    program []u8,
    pool @mut strpool.StringPool,
) ScanIdentResult ! (mem.AllocError | strpool.InternError) {
    cursor := start + 1
    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
    # member contexts accept keywords and therefore need their interned spelling too.
    str_id :: try strpool.intern(pool, program[start..cursor])
    return ScanIdentResult{ end = cursor, kind = kind, str_id = str_id }
}

@hide:file
ScanNumResult :: struct {
    end usize
    has_decimal bool
}

@hide:file
scan_number proc(start usize, program []u8, allow_decimal bool) ScanNumResult ! ScanError {
    cursor := start
    has_decimal := false
    while (cursor < program.len and is_digit(program[cursor])) cursor += 1
    if (
        allow_decimal and cursor < program.len and program[cursor] == '.'
        and !(cursor + 1 < program.len and program[cursor + 1] == '.')
    ) {
        has_decimal = true
        cursor += 1
    }
    if has_decimal and (cursor >= program.len or !is_digit(program[cursor])) {
        return ScanError{
            code = .float_must_end_with_digit,
            end = cursor,
        }
    }
    while (cursor < program.len and is_digit(program[cursor])) cursor += 1
    return ScanNumResult{ end = cursor, has_decimal = has_decimal }
}

@hide:file
ScanStrResult :: struct { end usize }

@hide:file
scan_quoted proc(start usize, program []u8, quote u8) ScanStrResult ! ScanError {
    cursor := start + 1
    while (cursor < program.len and program[cursor] != quote and program[cursor] != '\n') : cursor += 1 {
        # an escaped delimiter is content, but a newline still ends an invalid literal.
        if (program[cursor] == '\\' and cursor + 1 < program.len and program[cursor + 1] != '\n') cursor += 1
    }
    if (cursor >= program.len or program[cursor] != quote) {
        code :: if (quote == '\'') ErrorCode.unterminated_character else .unterminated_string
        return ScanError{
            code = code,
            end = cursor,
        }
    }
    return ScanStrResult{ end = cursor + 1 }
}

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

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

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

@hide:file
diagnostic_codes :: $enummap.init(ErrorCode, diagnostics.Code, {
    float_must_end_with_digit = .float_must_end_with_digit,
    unterminated_string = .unterminated_string,
    unterminated_character = .unterminated_character,
})

@hide:file
diagnostic_code proc(code ErrorCode) diagnostics.Code {
    return enummap.get(&diagnostic_codes, code)?
}

@hide:file
make_span proc(source diagnostics.SourceId, start, end usize) diagnostics.Span {
    return diagnostics.Span{
        start = u32(start),
        end = u32(end),
        source = source,
    }
}

@hide:file
add_token proc(state @mut State, token Token) void ! mem.AllocError {
    _ = token_id(state.tokens.items.len)
    try arraylist.append(&state.tokens, token)
}