package p.studio.compiler.pbs.lexer;

import p.studio.compiler.source.Span;
import p.studio.compiler.source.diagnostics.DiagnosticSink;
import p.studio.compiler.source.identifiers.FileId;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;

public final class PbsLexer {
    private static final Map<String, PbsTokenKind> KEYWORDS = buildKeywords();

    private enum LexerState {
        DEFAULT,
        IDENTIFIER,
        NUMBER,
        STRING,
        LINE_COMMENT
    }

    private final String source;
    private final FileId fileId;
    private final DiagnosticSink diagnostics;
    private final ArrayList<PbsToken> tokens = new ArrayList<>();

    private int start;
    private int current;
    private int startByte;
    private int currentByte;
    private LexerState state = LexerState.DEFAULT;

    private PbsLexer(final String source, final FileId fileId, final DiagnosticSink diagnostics) {
        this.source = source == null ? "" : source;
        this.fileId = fileId;
        this.diagnostics = diagnostics;
    }

    public static ReadOnlyList<PbsToken> lex(
            final String source,
            final FileId fileId,
            final DiagnosticSink diagnostics) {
        final var lexer = new PbsLexer(source, fileId, diagnostics);
        return lexer.lexInternal();
    }

    private ReadOnlyList<PbsToken> lexInternal() {
        while (!isAtEnd() || state != LexerState.DEFAULT) {
            switch (state) {
                case DEFAULT -> scanDefaultState();
                case IDENTIFIER -> scanIdentifierState();
                case NUMBER -> scanNumberState();
                case STRING -> scanStringState();
                case LINE_COMMENT -> scanLineCommentState();
            }
        }
        tokens.add(new PbsToken(PbsTokenKind.EOF, "", currentByte, currentByte));
        return ReadOnlyList.wrap(tokens);
    }

    private void scanDefaultState() {
        if (isAtEnd()) {
            return;
        }
        start = current;
        startByte = currentByte;
        final int c = advance();
        switch (c) {
            case ' ', '\r', '\t', '\n' -> {
                // Deliberately ignored.
            }
            case '(' -> addToken(PbsTokenKind.LEFT_PAREN);
            case ')' -> addToken(PbsTokenKind.RIGHT_PAREN);
            case '{' -> addToken(PbsTokenKind.LEFT_BRACE);
            case '}' -> addToken(PbsTokenKind.RIGHT_BRACE);
            case '[' -> addToken(PbsTokenKind.LEFT_BRACKET);
            case ']' -> addToken(PbsTokenKind.RIGHT_BRACKET);
            case ',' -> addToken(PbsTokenKind.COMMA);
            case ':' -> addToken(PbsTokenKind.COLON);
            case ';' -> addToken(PbsTokenKind.SEMICOLON);
            case '@' -> addToken(PbsTokenKind.AT);
            case '?' -> addToken(PbsTokenKind.QUESTION);
            case '+' -> addToken(match('=') ? PbsTokenKind.PLUS_EQUAL : PbsTokenKind.PLUS);
            case '-' -> {
                if (match('>')) {
                    addToken(PbsTokenKind.ARROW);
                    return;
                }
                addToken(match('=') ? PbsTokenKind.MINUS_EQUAL : PbsTokenKind.MINUS);
            }
            case '*' -> addToken(match('=') ? PbsTokenKind.STAR_EQUAL : PbsTokenKind.STAR);
            case '%' -> addToken(match('=') ? PbsTokenKind.PERCENT_EQUAL : PbsTokenKind.PERCENT);
            case '.' -> {
                if (match('.')) {
                    addToken(PbsTokenKind.DOT_DOT);
                    return;
                }
                addToken(PbsTokenKind.DOT);
            }
            case '!' -> addToken(match('=') ? PbsTokenKind.BANG_EQUAL : PbsTokenKind.BANG);
            case '=' -> addToken(match('=') ? PbsTokenKind.EQUAL_EQUAL : PbsTokenKind.EQUAL);
            case '<' -> addToken(match('=') ? PbsTokenKind.LESS_EQUAL : PbsTokenKind.LESS);
            case '>' -> addToken(match('=') ? PbsTokenKind.GREATER_EQUAL : PbsTokenKind.GREATER);
            case '&' -> {
                if (match('&')) {
                    addToken(PbsTokenKind.AND_AND);
                    return;
                }
                report(LexErrors.E_LEX_INVALID_CHAR, "Unexpected '&'");
            }
            case '|' -> {
                if (match('|')) {
                    addToken(PbsTokenKind.OR_OR);
                    return;
                }
                report(LexErrors.E_LEX_INVALID_CHAR, "Unexpected '|'");
            }
            case '/' -> {
                if (match('/')) {
                    state = LexerState.LINE_COMMENT;
                    return;
                }
                addToken(match('=') ? PbsTokenKind.SLASH_EQUAL : PbsTokenKind.SLASH);
            }
            case '"' -> state = LexerState.STRING;
            default -> {
                if (isDigit(c)) {
                    state = LexerState.NUMBER;
                    return;
                }
                if (isIdentifierStart(c)) {
                    state = LexerState.IDENTIFIER;
                    return;
                }
                report(LexErrors.E_LEX_INVALID_CHAR, "Invalid character: '%s'".formatted(Character.toString(c)));
            }
        }
    }

    private void scanIdentifierState() {
        while (!isAtEnd() && isIdentifierPart(peek())) {
            advance();
        }

        final String text = source.substring(start, current);
        final PbsTokenKind kind = KEYWORDS.getOrDefault(text, PbsTokenKind.IDENTIFIER);
        addToken(kind);
        state = LexerState.DEFAULT;
    }

    private void scanNumberState() {
        while (!isAtEnd() && isDigit(peek())) {
            advance();
        }

        var isFloat = false;
        if (!isAtEnd() && peek() == '.' && peekNext() != '.') {
            isFloat = true;
            advance();
            while (!isAtEnd() && isDigit(peek())) {
                advance();
            }
        }

        if (!isFloat && !isAtEnd() && peek() == 'b') {
            advance();
            addToken(PbsTokenKind.BOUNDED_LITERAL);
            state = LexerState.DEFAULT;
            return;
        }

        addToken(isFloat ? PbsTokenKind.FLOAT_LITERAL : PbsTokenKind.INT_LITERAL);
        state = LexerState.DEFAULT;
    }

    private void scanStringState() {
        while (!isAtEnd() && peek() != '"') {
            if (peek() == '\\') {
                final var escapeStart = currentByte;
                advance();
                if (!isAtEnd()) {
                    final var escaped = advance();
                    if (!isValidStringEscape(escaped)) {
                        report(
                                LexErrors.E_LEX_INVALID_STRING_ESCAPE,
                                "Invalid string escape '\\%s'".formatted(Character.toString(escaped)),
                                escapeStart,
                                currentByte);
                    }
                }
                continue;
            }
            advance();
        }

        if (isAtEnd()) {
            report(LexErrors.E_LEX_UNTERMINATED_STRING, "Unterminated string literal");
            state = LexerState.DEFAULT;
            return;
        }

        // Closing quote.
        advance();
        addToken(PbsTokenKind.STRING_LITERAL);
        state = LexerState.DEFAULT;
    }

    private void scanLineCommentState() {
        while (!isAtEnd() && peek() != '\n') {
            advance();
        }
        addToken(PbsTokenKind.COMMENT);
        state = LexerState.DEFAULT;
    }

    private void addToken(final PbsTokenKind kind) {
        final String lexeme = source.substring(start, current);
        tokens.add(new PbsToken(kind, lexeme, startByte, currentByte));
    }

    private boolean match(final char expected) {
        if (isAtEnd()) return false;
        final var codePoint = source.codePointAt(current);
        if (Character.charCount(codePoint) != 1 || codePoint != expected) {
            return false;
        }
        current++;
        currentByte++;
        return true;
    }

    private int advance() {
        final int codePoint = source.codePointAt(current);
        current += Character.charCount(codePoint);
        currentByte += utf8Bytes(codePoint);
        return codePoint;
    }

    private int peek() {
        if (isAtEnd()) return '\0';
        return source.codePointAt(current);
    }

    private int peekNext() {
        if (isAtEnd()) return '\0';
        final int first = source.codePointAt(current);
        final int nextIndex = current + Character.charCount(first);
        if (nextIndex >= source.length()) return '\0';
        return source.codePointAt(nextIndex);
    }

    private boolean isAtEnd() {
        return current >= source.length();
    }

    private boolean isDigit(final int c) {
        return c >= '0' && c <= '9';
    }

    private boolean isIdentifierStart(final int c) {
        return c == '_' || Character.isAlphabetic(c);
    }

    private boolean isIdentifierPart(final int c) {
        return c == '_' || Character.isAlphabetic(c) || Character.isDigit(c);
    }

    private boolean isValidStringEscape(final int escaped) {
        return escaped == '\\'
                || escaped == '"'
                || escaped == 'n'
                || escaped == 'r'
                || escaped == 't';
    }

    private void report(final LexErrors lexErrors, final String message) {
        report(lexErrors, message, startByte, currentByte);
    }

    private void report(
            final LexErrors lexErrors,
            final String message,
            final int spanStart,
            final int spanEnd) {
        p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, lexErrors.name(), message, new Span(fileId, spanStart, spanEnd));
    }

    private int utf8Bytes(final int codePoint) {
        if (codePoint <= 0x7F) {
            return 1;
        }
        if (codePoint <= 0x7FF) {
            return 2;
        }
        if (codePoint <= 0xFFFF) {
            return 3;
        }
        return 4;
    }

    private static Map<String, PbsTokenKind> buildKeywords() {
        final var map = new HashMap<String, PbsTokenKind>();
        map.put("import", PbsTokenKind.IMPORT);
        map.put("from", PbsTokenKind.FROM);
        map.put("as", PbsTokenKind.AS);

        map.put("service", PbsTokenKind.SERVICE);
        map.put("host", PbsTokenKind.HOST);
        map.put("fn", PbsTokenKind.FN);
        map.put("apply", PbsTokenKind.APPLY);
        map.put("bind", PbsTokenKind.BIND);
        map.put("new", PbsTokenKind.NEW);
        map.put("implements", PbsTokenKind.IMPLEMENTS);
        map.put("using", PbsTokenKind.USING);
        map.put("ctor", PbsTokenKind.CTOR);
        map.put("let", PbsTokenKind.LET);
        map.put("const", PbsTokenKind.CONST);
        map.put("global", PbsTokenKind.GLOBAL);
        map.put("declare", PbsTokenKind.DECLARE);
        map.put("struct", PbsTokenKind.STRUCT);
        map.put("contract", PbsTokenKind.CONTRACT);
        map.put("error", PbsTokenKind.ERROR);
        map.put("enum", PbsTokenKind.ENUM);
        map.put("callback", PbsTokenKind.CALLBACK);
        map.put("builtin", PbsTokenKind.BUILTIN);
        map.put("Self", PbsTokenKind.SELF);
        map.put("this", PbsTokenKind.THIS);
        map.put("pub", PbsTokenKind.PUB);
        map.put("mut", PbsTokenKind.MUT);
        map.put("mod", PbsTokenKind.MOD);
        map.put("type", PbsTokenKind.TYPE);

        map.put("if", PbsTokenKind.IF);
        map.put("else", PbsTokenKind.ELSE);
        map.put("switch", PbsTokenKind.SWITCH);
        map.put("default", PbsTokenKind.DEFAULT);
        map.put("for", PbsTokenKind.FOR);
        map.put("until", PbsTokenKind.UNTIL);
        map.put("step", PbsTokenKind.STEP);
        map.put("while", PbsTokenKind.WHILE);
        map.put("break", PbsTokenKind.BREAK);
        map.put("continue", PbsTokenKind.CONTINUE);
        map.put("return", PbsTokenKind.RETURN);
        map.put("void", PbsTokenKind.VOID);
        map.put("optional", PbsTokenKind.OPTIONAL);
        map.put("result", PbsTokenKind.RESULT);
        map.put("some", PbsTokenKind.SOME);
        map.put("none", PbsTokenKind.NONE);
        map.put("ok", PbsTokenKind.OK);
        map.put("err", PbsTokenKind.ERR);
        map.put("handle", PbsTokenKind.HANDLE);

        map.put("true", PbsTokenKind.TRUE);
        map.put("false", PbsTokenKind.FALSE);
        map.put("and", PbsTokenKind.AND);
        map.put("or", PbsTokenKind.OR);
        map.put("not", PbsTokenKind.NOT);

        map.put("spawn", PbsTokenKind.SPAWN);
        map.put("yield", PbsTokenKind.YIELD);
        map.put("sleep", PbsTokenKind.SLEEP);
        map.put("match", PbsTokenKind.MATCH);
        return map;
    }
}
