package org.zwobble.hobgoblin.compiler.parser;

import org.zwobble.hobgoblin.compiler.ast.DocComment;
import org.zwobble.hobgoblin.compiler.ast.untyped.*;
import org.zwobble.hobgoblin.compiler.sources.FileFragmentSource;
import org.zwobble.hobgoblin.compiler.sources.Source;
import org.zwobble.hobgoblin.compiler.types.NamespaceName;
import org.zwobble.hobgoblin.compiler.types.QualifiedTypeName;
import org.zwobble.sourcetext.SourcePosition;
import org.zwobble.sourcetext.SourceText;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.function.BooleanSupplier;
import java.util.function.Supplier;

public class Parser {
    private Parser() {
    }

    // == Namespaces ==

    public static UntypedNamespaceNode parseNamespace(
        SourceText sourceText,
        NamespaceName namespaceName
    ) {
        var tokens = Tokenizer.tokenize(sourceText);
        return parseNamespace(tokens.iterator(), namespaceName);
    }

    static UntypedNamespaceNode parseNamespace(
        TokenIterator tokens,
        NamespaceName namespaceName
    ) {
        var start = tokens.startPosition();

        var imports = parseMany(
            () -> !tokens.isNext(TokenType.KEYWORD_IMPORT),
            () -> parseImport(tokens),
            () -> true
        );

        var body = parseMany(
            () -> tokens.isNext(TokenType.END),
            () -> parseNamespaceStatement(tokens),
            () -> true
        );

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedNamespaceNode(
            namespaceName,
            imports,
            body,
            source
        );
    }

    private static UntypedImportNode parseImport(TokenIterator tokens) {
        var start = tokens.startPosition();

        tokens.skip(TokenType.KEYWORD_IMPORT);
        tokens.skip(TokenType.SYMBOL_BRACE_OPEN);

        var importedNames = parseMany(
            () -> tokens.isNext(TokenType.SYMBOL_BRACE_CLOSE),
            () -> parseIdentifier(tokens),
            () -> tokens.trySkip(TokenType.SYMBOL_COMMA)
        );

        tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);
        tokens.skip(TokenType.KEYWORD_FROM);

        var isSibling = tokens.trySkip(TokenType.SYMBOL_DOT);

        Optional<Integer> ancestorDepth;
        if (isSibling) {
            tokens.skip(TokenType.SYMBOL_FORWARD_SLASH);
            ancestorDepth = Optional.of(0);
        } else if (tokens.isNext(TokenType.SYMBOL_DOT_DOT)) {
            int depth = 0;
            while (tokens.trySkip(TokenType.SYMBOL_DOT_DOT)) {
                depth += 1;
                tokens.skip(TokenType.SYMBOL_FORWARD_SLASH);
            }
            ancestorDepth = Optional.of(depth);
        } else {
            ancestorDepth = Optional.empty();
        }

        var namespaceName = parseNamespaceName(tokens);

        tokens.skip(TokenType.SYMBOL_SEMICOLON);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedImportNode(importedNames, ancestorDepth, namespaceName, source);
    }

    // == Namespace statements ==

    static UntypedNamespaceStatementNode parseNamespaceStatement(
        TokenIterator tokens
    ) {
        var docComment = parseDocComment(tokens);

        if (tokens.isNext(TokenType.IDENTIFIER, "databaseSchema")) {
            return parseDatabaseSchemaDefinition(tokens, docComment);
        } else if (tokens.isNext(TokenType.KEYWORD_ENUM)) {
            return parseEnumDefinition(tokens, docComment);
        } else if (tokens.isNext(TokenType.KEYWORD_NATIVE)) {
            return parseNativeTypeDefinition(tokens, docComment);
        } else if (tokens.isNext(TokenType.KEYWORD_STRUCT)) {
            return parseStructDefinition(tokens, docComment);
        } else if (tokens.isNext(TokenType.KEYWORD_SUM)) {
            return parseSumDefinition(tokens, docComment);
        } else {
            throw new UnexpectedTokenError("namespace statement", tokens.peek());
        }
    }

    private static UntypedDatabaseSchemaDefinitionNode parseDatabaseSchemaDefinition(
        TokenIterator tokens,
        DocComment docComment
    ) {
        var start = tokens.startPosition();

        tokens.skip(TokenType.IDENTIFIER, "databaseSchema");
        var name = parseIdentifier(tokens);
        tokens.skip(TokenType.SYMBOL_BRACE_OPEN);

        var body = parseMany(
            () -> tokens.isNext(TokenType.SYMBOL_BRACE_CLOSE),
            () -> parseDatabaseSchemaStatement(tokens),
            () -> true
        );

        tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedDatabaseSchemaDefinitionNode(
            name,
            body,
            docComment,
            source
        );
    }

    private static UntypedEnumDefinitionNode parseEnumDefinition(
        TokenIterator tokens,
        DocComment docComment
    ) {
        var start = tokens.startPosition();

        tokens.skip(TokenType.KEYWORD_ENUM);
        var name = tokens.next(TokenType.IDENTIFIER);
        tokens.skip(TokenType.SYMBOL_BRACE_OPEN);

        var variants = parseMany(
            () -> !tokens.isNext(TokenType.KEYWORD_VARIANT),
            () -> parseEnumVariantDefinition(tokens),
            () -> true
        );

        tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedEnumDefinitionNode(
            name.charSequence().toString(),
            variants,
            docComment,
            source
        );
    }

    private static UntypedEnumVariantDefinitionNode parseEnumVariantDefinition(TokenIterator tokens) {
        var start = tokens.startPosition();

        tokens.skip(TokenType.KEYWORD_VARIANT);
        var name = parseIdentifier(tokens);

        tokens.skip(TokenType.SYMBOL_SEMICOLON);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedEnumVariantDefinitionNode(name, source);
    }

    private static UntypedNativeTypeDefinitionNode parseNativeTypeDefinition(
        TokenIterator tokens,
        DocComment docComment
    ) {
        var start = tokens.startPosition();

        tokens.skip(TokenType.KEYWORD_NATIVE);
        var name = tokens.next(TokenType.IDENTIFIER);

        var typeParams = parseTypeParams(tokens);

        tokens.skip(TokenType.SYMBOL_BRACE_OPEN);

        tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedNativeTypeDefinitionNode(
            name.charSequence().toString(),
            typeParams,
            docComment,
            source
        );
    }

    private static DocComment parseDocComment(TokenIterator tokens) {
        var docComment = new StringBuilder();

        while (tokens.isNext(TokenType.DOC_COMMENT)) {
            var docCommentToken = tokens.next(TokenType.DOC_COMMENT);
            var docCommentLine = docCommentToken.charSequence();

            var firstCharIndex = 3;
            while (docCommentLine.charAt(firstCharIndex) == ' ') {
                firstCharIndex++;
            }

            var docCommentContents = docCommentLine.subSequence(firstCharIndex, docCommentLine.length());
            docComment.append(docCommentContents);
        }

        return new DocComment(docComment.toString());
    }

    private static UntypedStructDefinitionNode parseStructDefinition(
        TokenIterator tokens,
        DocComment docComment
    ) {
        var start = tokens.startPosition();

        tokens.skip(TokenType.KEYWORD_STRUCT);
        var name = tokens.next(TokenType.IDENTIFIER);
        var typeParams = parseTypeParams(tokens);
        tokens.skip(TokenType.SYMBOL_BRACE_OPEN);

        var fields = parseStructDefinitionFields(tokens);

        tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedStructDefinitionNode(
            name.charSequence().toString(),
            typeParams,
            fields,
            docComment,
            source
        );
    }

    private static Optional<List<UntypedStructFieldDefinitionNode>> parseStructDefinitionFields(TokenIterator tokens) {
        var isSingleton = tokens.trySkip(TokenType.IDENTIFIER, "singleton");
        if (isSingleton) {
            tokens.skip(TokenType.SYMBOL_SEMICOLON);
            return Optional.empty();
        } else {
            var fields = parseMany(
                () -> tokens.isNext(TokenType.SYMBOL_BRACE_CLOSE),
                () -> parseStructFieldDefinition(tokens),
                () -> true
            );
            return Optional.of(fields);
        }
    }

    private static UntypedStructFieldDefinitionNode parseStructFieldDefinition(
        TokenIterator tokens
    ) {
        var start = tokens.startPosition();

        tokens.skip(TokenType.KEYWORD_FIELD);
        var name = tokens.next(TokenType.IDENTIFIER);
        tokens.skip(TokenType.SYMBOL_COLON);
        var type = parseTypeLevelExpression(tokens);
        tokens.skip(TokenType.SYMBOL_SEMICOLON);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedStructFieldDefinitionNode(
            name.charSequence().toString(),
            type,
            source
        );
    }

    private static UntypedSumDefinitionNode parseSumDefinition(
        TokenIterator tokens,
        DocComment docComment
    ) {
        var start = tokens.startPosition();

        tokens.skip(TokenType.KEYWORD_SUM);
        var name = tokens.next(TokenType.IDENTIFIER);
        tokens.skip(TokenType.SYMBOL_BRACE_OPEN);

        var variants = parseMany(
            () -> !tokens.isNext(TokenType.KEYWORD_VARIANT),
            () -> parseSumVariantDefinition(tokens),
            () -> true
        );

        var fields = parseMany(
            () -> !tokens.isNext(TokenType.KEYWORD_FIELD),
            () -> parseStructFieldDefinition(tokens),
            () -> true
        );

        tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedSumDefinitionNode(
            name.charSequence().toString(),
            variants,
            fields,
            docComment,
            source
        );
    }

    private static UntypedSumVariantDefinitionNode parseSumVariantDefinition(
        TokenIterator tokens
    ) {
        var start = tokens.startPosition();

        tokens.skip(TokenType.KEYWORD_VARIANT);
        var type = parseTypeLevelExpression(tokens);
        tokens.skip(TokenType.SYMBOL_SEMICOLON);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedSumVariantDefinitionNode(
            type,
            source
        );
    }

    // == Database schema statements ==

    static UntypedDatabaseSchemaStatementNode parseDatabaseSchemaStatement(TokenIterator tokens) {
        var start = tokens.startPosition();

        var docComment = parseDocComment(tokens);
        tokens.skip(TokenType.IDENTIFIER, "table");
        var name = parseIdentifier(tokens);
        tokens.skip(TokenType.SYMBOL_BRACE_OPEN);

        tokens.skip(TokenType.IDENTIFIER, "row");
        var rowType = parseTypeLevelExpression(tokens);
        tokens.skip(TokenType.SYMBOL_SEMICOLON);

        tokens.skip(TokenType.IDENTIFIER, "primaryKey");
        var primaryKey = parseIdentifier(tokens);
        tokens.skip(TokenType.SYMBOL_SEMICOLON);

        tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedTableDefinitionNode(
            name,
            rowType,
            primaryKey,
            docComment,
            source
        );
    }

    // == Type params ==

    private static Optional<List<UntypedTypeParamNode>> parseTypeParams(TokenIterator tokens) {
        if (!tokens.trySkip(TokenType.SYMBOL_SQUARE_OPEN)) {
            return Optional.empty();
        }

        var typeParams = parseMany(
            () -> tokens.isNext(TokenType.SYMBOL_SQUARE_CLOSE),
            () -> {
                var start = tokens.startPosition();

                var name = parseIdentifier(tokens);

                var end = tokens.endPosition();
                var source = createSource(start, end);

                return new UntypedTypeParamNode(name, source);
            },
            () -> tokens.trySkip(TokenType.SYMBOL_COMMA)
        );

        tokens.skip(TokenType.SYMBOL_SQUARE_CLOSE);

        return Optional.of(typeParams);
    }

    // == Type-level expressions ==

    static UntypedTypeLevelExpressionNode parseTypeLevelExpression(
        TokenIterator tokens
    ) {
        var start = tokens.startPosition();

        UntypedTypeLevelExpressionNode expression = parsePrimaryTypeLevelExpression(tokens);

        while (tokens.trySkip(TokenType.SYMBOL_SQUARE_OPEN)) {
            var args = parseMany(
                () -> tokens.isNext(TokenType.SYMBOL_SQUARE_CLOSE),
                () -> parseTypeLevelExpression(tokens),
                () -> tokens.trySkip(TokenType.SYMBOL_COMMA)
            );
            tokens.skip(TokenType.SYMBOL_SQUARE_CLOSE);

            var end = tokens.endPosition();
            var source = createSource(start, end);

            expression = new UntypedConstructedTypeNode(expression, args, source);
        }

        return expression;
    }

    private static UntypedTypeLevelReferenceNode parsePrimaryTypeLevelExpression(TokenIterator tokens) {
        var start = tokens.startPosition();

        var name = parseIdentifier(tokens);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedTypeLevelReferenceNode(
            name,
            source
        );
    }

    // == Query expressions ==

    static UntypedQueryExpressionNode parseQueryExpression(TokenIterator tokens) {
        var start = tokens.startPosition();

        var name = parseIdentifier(tokens);

        var end = tokens.endPosition();
        var source = createSource(start, end);

        return new UntypedQueryExpressionReferenceNode(name, source);
    }

    // == Qualified type names ==

    public static QualifiedTypeName parseQualifiedTypeName(SourceText sourceText) {
        var tokens = Tokenizer.tokenize(sourceText);
        var tokenIterator = tokens.iterator();
        var namespaceName = parseQualifiedTypeName(tokenIterator);
        tokenIterator.skip(TokenType.END);
        return namespaceName;
    }

    static QualifiedTypeName parseQualifiedTypeName(TokenIterator tokens) {
        var namespaceName = parseNamespaceName(tokens);
        tokens.skip(TokenType.SYMBOL_DOT);
        var typeName = parseIdentifier(tokens);
        return new QualifiedTypeName(namespaceName, typeName);
    }

    public static NamespaceName parseNamespaceName(SourceText sourceText) {
        var tokens = Tokenizer.tokenize(sourceText);
        var tokenIterator = tokens.iterator();
        var namespaceName = parseNamespaceName(tokenIterator);
        tokenIterator.skip(TokenType.END);
        return namespaceName;
    }

    static NamespaceName parseNamespaceName(TokenIterator tokens) {
        var namespaceParts = parseMany(
            () -> !tokens.isNext(TokenType.IDENTIFIER),
            () -> parseIdentifier(tokens),
            () -> tokens.trySkip(TokenType.SYMBOL_FORWARD_SLASH)
        );

        return new NamespaceName(namespaceParts);
    }

    private static String parseIdentifier(TokenIterator tokens) {
        return tokens.next(TokenType.IDENTIFIER).charSequence().toString();
    }

    /**
     * Parse zero, one, or many instances of {@code T} using {@code parseOne}.
     * Each instance is expected to be separated by tokens consumed by
     * {@code parseSeparator}. Parsing stops either if {@code parseSeparator}
     * returns false, or if {@code stop} returns true. Trailing separators
     * are permitted.
     */
    private static <T> List<T> parseMany(
        BooleanSupplier stop,
        Supplier<T> parseOne,
        BooleanSupplier parseSeparator
    ) {
        var values = new ArrayList<T>();

        while (!stop.getAsBoolean()) {
            var value = parseOne.get();
            values.add(value);
            if (!parseSeparator.getAsBoolean()) {
                return values;
            }
        }

        return values;
    }

    private static Source createSource(SourcePosition start, SourcePosition end) {
        return new FileFragmentSource(start.to(end));
    }
}