package org.shedlang.compiler.parser;

import org.shedlang.compiler.tokeniser.RegexTokeniser;
import org.shedlang.compiler.tokeniser.Token;
import org.shedlang.compiler.tokeniser.TokenIterator;
import org.zwobble.sourcetext.SourceText;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static org.shedlang.compiler.tokeniser.Token.token;

public class Tokeniser {
    private static final Map<String, TokenType> KEYWORDS = Map.ofEntries(
        Map.entry("break", TokenType.KEYWORD_BREAK),
        Map.entry("effect", TokenType.KEYWORD_EFFECT),
        Map.entry("else", TokenType.KEYWORD_ELSE),
        Map.entry("extend", TokenType.KEYWORD_EXTEND),
        Map.entry("false", TokenType.KEYWORD_FALSE),
        Map.entry("fn", TokenType.KEYWORD_FN),
        Map.entry("fun", TokenType.KEYWORD_FUN),
        Map.entry("handle", TokenType.KEYWORD_HANDLE),
        Map.entry("if", TokenType.KEYWORD_IF),
        Map.entry("import", TokenType.KEYWORD_IMPORT),
        Map.entry("is", TokenType.KEYWORD_IS),
        Map.entry("native", TokenType.KEYWORD_NATIVE),
        Map.entry("pub", TokenType.KEYWORD_PUB),
        Map.entry("resume", TokenType.KEYWORD_RESUME),
        Map.entry("return", TokenType.KEYWORD_RETURN),
        Map.entry("set", TokenType.KEYWORD_SET),
        Map.entry("struct", TokenType.KEYWORD_STRUCT),
        Map.entry("sumtype", TokenType.KEYWORD_SUMTYPE),
        Map.entry("true", TokenType.KEYWORD_TRUE),
        Map.entry("try", TokenType.KEYWORD_TRY),
        Map.entry("type", TokenType.KEYWORD_TYPE),
        Map.entry("unit", TokenType.KEYWORD_UNIT),
        Map.entry("val", TokenType.KEYWORD_VAL),
        Map.entry("var", TokenType.KEYWORD_VAR),
        Map.entry("when", TokenType.KEYWORD_WHEN),
        Map.entry("while", TokenType.KEYWORD_WHILE)
    );

    private static final RegexTokeniser<TokenType> tokeniser = new RegexTokeniser<>(
        TokenType.UNKNOWN,
        List.of(
            RegexTokeniser.rule(TokenType.IDENTIFIER, "[A-Za-z][A-Za-z0-9]*"),
            RegexTokeniser.rule(TokenType.INT, "-?[0-9]+"),
            RegexTokeniser.rule(TokenType.STRING, "\"(?:[^\\\\\"\n\r]|\\\\.)*\""),
            RegexTokeniser.rule(TokenType.WHITESPACE, "[\r\n\t ]+"),
            RegexTokeniser.rule(TokenType.COMMENT_SINGLE_LINE, "//[^\r\n]*[\r\n]"),

            RegexTokeniser.rule(TokenType.SYMBOL_DOT_DOT_DOT, "\\.\\.\\."),

            RegexTokeniser.rule(TokenType.SYMBOL_BANG_EQUALS, "!="),
            RegexTokeniser.rule(TokenType.SYMBOL_COLON_EQUALS, ":="),
            RegexTokeniser.rule(TokenType.SYMBOL_EQUALS_EQUALS, "=="),
            RegexTokeniser.rule(TokenType.SYMBOL_GREATER_THAN_OR_EQUAL, ">="),
            RegexTokeniser.rule(TokenType.SYMBOL_LESS_THAN_OR_EQUAL, "<="),
            RegexTokeniser.rule(TokenType.SYMBOL_THIN_ARROW, "->"),
            RegexTokeniser.rule(TokenType.SYMBOL_HASH_PAREN_OPEN, "#\\("),

            RegexTokeniser.rule(TokenType.SYMBOL_ASTERISK, "\\*"),
            RegexTokeniser.rule(TokenType.SYMBOL_BANG, "!"),
            RegexTokeniser.rule(TokenType.SYMBOL_BAR, "\\|"),
            RegexTokeniser.rule(TokenType.SYMBOL_BRACE_OPEN, "\\{"),
            RegexTokeniser.rule(TokenType.SYMBOL_BRACE_CLOSE, "\\}"),
            RegexTokeniser.rule(TokenType.SYMBOL_COLON, ":"),
            RegexTokeniser.rule(TokenType.SYMBOL_COMMA, ","),
            RegexTokeniser.rule(TokenType.SYMBOL_DOT, "\\."),
            RegexTokeniser.rule(TokenType.SYMBOL_EQUALS, "="),
            RegexTokeniser.rule(TokenType.SYMBOL_FORWARD_SLASH, "/"),
            RegexTokeniser.rule(TokenType.SYMBOL_GREATER_THAN, ">"),
            RegexTokeniser.rule(TokenType.SYMBOL_LESS_THAN, "<"),
            RegexTokeniser.rule(TokenType.SYMBOL_MINUS, "-"),
            RegexTokeniser.rule(TokenType.SYMBOL_PAREN_OPEN, "\\("),
            RegexTokeniser.rule(TokenType.SYMBOL_PAREN_CLOSE, "\\)"),
            RegexTokeniser.rule(TokenType.SYMBOL_PLUS, "\\+"),
            RegexTokeniser.rule(TokenType.SYMBOL_SEMICOLON, ";"),
            RegexTokeniser.rule(TokenType.SYMBOL_SQUARE_OPEN, "\\["),
            RegexTokeniser.rule(TokenType.SYMBOL_SQUARE_CLOSE, "\\]"),
            RegexTokeniser.rule(TokenType.SYMBOL_UNDERSCORE, "_")
        )
    );

    public static TokenIterator<TokenType> tokenise(SourceText sourceText) {
        var tokens = tokeniser.tokenise(sourceText).stream()
            .filter(token ->
                token.tokenType() != TokenType.WHITESPACE &&
                    token.tokenType() != TokenType.COMMENT_SINGLE_LINE
            )
            .map(token ->
                token.tokenType() == TokenType.IDENTIFIER
                    ? new Token<>(token.sourceRange(), KEYWORDS.getOrDefault(token.value().toString(), token.tokenType()))
                    : token
            )
            .collect(Collectors.toList());

        var endPosition = sourceText.characterPosition(sourceText.characterLength());
        var endSourceRange = endPosition.to(endPosition);
        var endToken = token(endSourceRange, TokenType.END);

        return new TokenIterator<>(tokens, endToken);
    }
}