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

import "@source/strpool"

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

TokenIndex :: alias usize

ScanDiagnostic :: struct {
    token TokenIndex
    message []u8
}

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

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 arraylist.append(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 {}
            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.NoId

            try arraylist.append(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 :: tokens.items.len
	            try arraylist.append(tokens, Token{ kind = .invalid, start = start })
	            try arraylist.append(diagnostics, ScanDiagnostic{ token = token, message = "float must end with a digit" })
	            continue
            }

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

            if has_decimal {
                try arraylist.append(tokens, Token{ kind = .float, start = start })
            } else {
                try arraylist.append(tokens, Token{ kind = .int, 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 arraylist.append(tokens, Token{ kind = .string, start = start })
            } else {
                try arraylist.append(tokens, Token{ kind = .invalid, start = start })
            }

            continue
        }

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

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

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

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

        # invalid character
        token :: tokens.items.len
        try arraylist.append(tokens, Token{ kind = .invalid, start = cursor })
        try arraylist.append(diagnostics, ScanDiagnostic{ token = token, message = "invalid character" })
        cursor += 1
    }

    try arraylist.append(tokens, Token{ kind = .eof, start = 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
    }
}