package p.studio.compiler.pbs.parser;

import p.studio.compiler.pbs.ast.PbsAst;
import p.studio.compiler.pbs.lexer.PbsToken;
import p.studio.compiler.pbs.lexer.PbsTokenKind;
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;

/**
 * Parser for `mod.barrel` files.
 */
public final class PbsBarrelParser {
    private static final int MAX_NAMED_TUPLE_ARITY = 6;

    private final PbsTokenCursor cursor;
    private final FileId fileId;
    private final DiagnosticSink diagnostics;

    private PbsBarrelParser(
            final ReadOnlyList<PbsToken> tokens,
            final FileId fileId,
            final DiagnosticSink diagnostics) {
        this.cursor = new PbsTokenCursor(tokens);
        this.fileId = fileId;
        this.diagnostics = diagnostics;
    }

    public static PbsAst.BarrelFile parse(
            final ReadOnlyList<PbsToken> tokens,
            final FileId fileId,
            final DiagnosticSink diagnostics) {
        return new PbsBarrelParser(tokens, fileId, diagnostics).parseFile();
    }

    private PbsAst.BarrelFile parseFile() {
        final var items = new ArrayList<PbsAst.BarrelItem>();

        while (!cursor.isAtEnd()) {
            if (cursor.check(PbsTokenKind.EOF)) {
                break;
            }

            final var visibility = parseVisibility();
            if (visibility == null) {
                synchronizeToSemicolon();
                continue;
            }

            final var item = parseItem(visibility, cursor.previous());
            if (item != null) {
                items.add(item);
            }
        }

        final var eof = cursor.peek();
        return new PbsAst.BarrelFile(ReadOnlyList.wrap(items), span(0, eof.end()));
    }

    private PbsAst.Visibility parseVisibility() {
        if (cursor.match(PbsTokenKind.MOD)) {
            return PbsAst.Visibility.MOD;
        }
        if (cursor.match(PbsTokenKind.PUB)) {
            return PbsAst.Visibility.PUB;
        }

        report(cursor.peek(), ParseErrors.E_PARSE_UNEXPECTED_TOKEN,
                "Expected barrel visibility ('mod' or 'pub')");
        if (!cursor.isAtEnd()) {
            cursor.advance();
        }
        return null;
    }

    private PbsAst.BarrelItem parseItem(final PbsAst.Visibility visibility, final PbsToken visibilityToken) {
        if (cursor.match(PbsTokenKind.FN)) {
            return parseFunctionItem(visibility, visibilityToken);
        }
        if (cursor.match(PbsTokenKind.STRUCT)) {
            return parseSimpleItem(visibility, visibilityToken, PbsTokenKind.STRUCT);
        }
        if (cursor.match(PbsTokenKind.CONTRACT)) {
            return parseSimpleItem(visibility, visibilityToken, PbsTokenKind.CONTRACT);
        }
        if (cursor.match(PbsTokenKind.HOST)) {
            return parseSimpleItem(visibility, visibilityToken, PbsTokenKind.HOST);
        }
        if (cursor.match(PbsTokenKind.ERROR)) {
            return parseSimpleItem(visibility, visibilityToken, PbsTokenKind.ERROR);
        }
        if (cursor.match(PbsTokenKind.ENUM)) {
            return parseSimpleItem(visibility, visibilityToken, PbsTokenKind.ENUM);
        }
        if (cursor.match(PbsTokenKind.SERVICE)) {
            return parseSimpleItem(visibility, visibilityToken, PbsTokenKind.SERVICE);
        }
        if (cursor.match(PbsTokenKind.GLOBAL)) {
            return parseSimpleItem(visibility, visibilityToken, PbsTokenKind.GLOBAL);
        }
        if (cursor.match(PbsTokenKind.CONST)) {
            return parseSimpleItem(visibility, visibilityToken, PbsTokenKind.CONST);
        }
        if (cursor.match(PbsTokenKind.CALLBACK)) {
            return parseSimpleItem(visibility, visibilityToken, PbsTokenKind.CALLBACK);
        }

        report(cursor.peek(), ParseErrors.E_PARSE_UNEXPECTED_TOKEN,
                "Expected barrel declaration kind after visibility");
        synchronizeToSemicolon();
        return null;
    }

    private PbsAst.BarrelItem parseSimpleItem(
            final PbsAst.Visibility visibility,
            final PbsToken visibilityToken,
            final PbsTokenKind kind) {
        final var name = consume(PbsTokenKind.IDENTIFIER, "Expected declaration name in barrel item");
        final var semicolon = consume(PbsTokenKind.SEMICOLON, "Expected ';' after barrel item");
        final var itemSpan = span(visibilityToken.start(), semicolon.end());

        return switch (kind) {
            case STRUCT -> new PbsAst.BarrelStructItem(visibility, name.lexeme(), itemSpan);
            case CONTRACT -> new PbsAst.BarrelContractItem(visibility, name.lexeme(), itemSpan);
            case HOST -> new PbsAst.BarrelHostItem(visibility, name.lexeme(), itemSpan);
            case ERROR -> new PbsAst.BarrelErrorItem(visibility, name.lexeme(), itemSpan);
            case ENUM -> new PbsAst.BarrelEnumItem(visibility, name.lexeme(), itemSpan);
            case SERVICE -> new PbsAst.BarrelServiceItem(visibility, name.lexeme(), itemSpan);
            case GLOBAL -> new PbsAst.BarrelGlobalItem(visibility, name.lexeme(), itemSpan);
            case CONST -> new PbsAst.BarrelConstItem(visibility, name.lexeme(), itemSpan);
            case CALLBACK -> new PbsAst.BarrelCallbackItem(visibility, name.lexeme(), itemSpan);
            default -> null;
        };
    }

    private PbsAst.BarrelFunctionItem parseFunctionItem(
            final PbsAst.Visibility visibility,
            final PbsToken visibilityToken) {
        final var name = consume(PbsTokenKind.IDENTIFIER, "Expected function name in barrel item");
        consume(PbsTokenKind.LEFT_PAREN, "Expected '(' after function name");

        final var parameterTypes = new ArrayList<PbsAst.TypeRef>();
        if (!cursor.check(PbsTokenKind.RIGHT_PAREN)) {
            do {
                consume(PbsTokenKind.IDENTIFIER, "Expected parameter name in barrel function signature");
                consume(PbsTokenKind.COLON, "Expected ':' after parameter name in barrel function signature");
                parameterTypes.add(parseTypeRef());
            } while (cursor.match(PbsTokenKind.COMMA));
        }
        final var rightParen = consume(PbsTokenKind.RIGHT_PAREN, "Expected ')' after barrel function parameters");
        final var returnSpec = parseCallableReturnSpec(rightParen);

        final var semicolon = consume(PbsTokenKind.SEMICOLON, "Expected ';' after barrel function item");
        return new PbsAst.BarrelFunctionItem(
                visibility,
                name.lexeme(),
                ReadOnlyList.wrap(parameterTypes),
                returnSpec.kind(),
                returnSpec.returnType(),
                returnSpec.resultErrorType(),
                span(visibilityToken.start(), semicolon.end()));
    }

    private ParsedReturnSpec parseCallableReturnSpec(final PbsToken anchorToken) {
        if (cursor.match(PbsTokenKind.ARROW)) {
            return parseReturnSpecFromMarker();
        }

        if (cursor.match(PbsTokenKind.COLON)) {
            report(cursor.previous(), ParseErrors.E_PARSE_INVALID_RETURN_ANNOTATION,
                    "Barrel return annotations must use '->' syntax");
            return parseReturnSpecFromMarker();
        }

        final var inferredSpan = span(anchorToken.end(), anchorToken.end());
        return new ParsedReturnSpec(
                PbsAst.ReturnKind.INFERRED_UNIT,
                PbsAst.TypeRef.unit(inferredSpan),
                null);
    }

    private ParsedReturnSpec parseReturnSpecFromMarker() {
        if (cursor.match(PbsTokenKind.RESULT)) {
            final var resultToken = cursor.previous();
            consume(PbsTokenKind.LESS, "Expected '<' after 'result'");
            final var errorType = parseTypeRef();
            consume(PbsTokenKind.GREATER, "Expected '>' after result error type");

            final PbsAst.TypeRef payloadType;
            if (isTypeStart(cursor.peek().kind())) {
                payloadType = parseTypeRef();
            } else {
                payloadType = PbsAst.TypeRef.unit(span(resultToken.end(), resultToken.end()));
            }

            return new ParsedReturnSpec(
                    PbsAst.ReturnKind.RESULT,
                    payloadType,
                    errorType);
        }

        final var typeRef = parseTypeRef();
        if (typeRef.kind() == PbsAst.TypeRefKind.UNIT) {
            return new ParsedReturnSpec(PbsAst.ReturnKind.EXPLICIT_UNIT, typeRef, null);
        }
        return new ParsedReturnSpec(PbsAst.ReturnKind.PLAIN, typeRef, null);
    }

    private PbsAst.TypeRef parseTypeRef() {
        if (cursor.match(PbsTokenKind.OPTIONAL)) {
            final var optionalToken = cursor.previous();
            final var inner = parseTypeRef();
            if (inner.kind() == PbsAst.TypeRefKind.UNIT) {
                report(optionalToken, ParseErrors.E_PARSE_INVALID_TYPE_SURFACE,
                        "'optional void' is not a valid type surface");
            }
            return PbsAst.TypeRef.optional(inner, span(optionalToken.start(), inner.span().getEnd()));
        }

        if (cursor.match(PbsTokenKind.VOID)) {
            final var token = cursor.previous();
            return PbsAst.TypeRef.unit(span(token.start(), token.end()));
        }

        if (cursor.match(PbsTokenKind.SELF)) {
            final var token = cursor.previous();
            return PbsAst.TypeRef.self(span(token.start(), token.end()));
        }

        if (cursor.match(PbsTokenKind.IDENTIFIER)) {
            final var token = cursor.previous();
            return PbsAst.TypeRef.simple(token.lexeme(), span(token.start(), token.end()));
        }

        if (cursor.match(PbsTokenKind.LEFT_PAREN)) {
            final var open = cursor.previous();
            if (cursor.match(PbsTokenKind.RIGHT_PAREN)) {
                return PbsAst.TypeRef.unit(span(open.start(), cursor.previous().end()));
            }

            if (cursor.check(PbsTokenKind.IDENTIFIER) && cursor.checkNext(PbsTokenKind.COLON)) {
                final var fields = new ArrayList<PbsAst.NamedTypeField>();
                do {
                    final var label = consume(PbsTokenKind.IDENTIFIER, "Expected tuple field label");
                    consume(PbsTokenKind.COLON, "Expected ':' after tuple field label");
                    final var type = parseTypeRef();
                    fields.add(new PbsAst.NamedTypeField(label.lexeme(), type, span(label.start(), type.span().getEnd())));
                    if (fields.size() == MAX_NAMED_TUPLE_ARITY + 1) {
                        report(label, ParseErrors.E_PARSE_INVALID_TYPE_SURFACE,
                                "Named tuple type arity must be between 1 and 6 fields");
                    }
                } while (cursor.match(PbsTokenKind.COMMA) && !cursor.check(PbsTokenKind.RIGHT_PAREN));

                final var close = consume(PbsTokenKind.RIGHT_PAREN, "Expected ')' after named tuple type");
                return PbsAst.TypeRef.namedTuple(ReadOnlyList.wrap(fields), span(open.start(), close.end()));
            }

            final var inner = parseTypeRef();
            final var close = consume(PbsTokenKind.RIGHT_PAREN, "Expected ')' after grouped type");
            return PbsAst.TypeRef.group(inner, span(open.start(), close.end()));
        }

        final var token = cursor.peek();
        report(token, ParseErrors.E_PARSE_EXPECTED_TOKEN, "Expected type surface");
        if (!cursor.isAtEnd()) {
            cursor.advance();
        }
        return PbsAst.TypeRef.error(span(token.start(), token.end()));
    }

    private boolean isTypeStart(final PbsTokenKind kind) {
        return switch (kind) {
            case OPTIONAL, VOID, SELF, IDENTIFIER, LEFT_PAREN -> true;
            default -> false;
        };
    }

    private void synchronizeToSemicolon() {
        while (!cursor.isAtEnd()) {
            if (cursor.match(PbsTokenKind.SEMICOLON)) {
                return;
            }
            if (cursor.check(PbsTokenKind.MOD) || cursor.check(PbsTokenKind.PUB)) {
                return;
            }
            cursor.advance();
        }
    }

    private PbsToken consume(final PbsTokenKind kind, final String message) {
        if (cursor.check(kind)) {
            return cursor.advance();
        }
        final var token = cursor.peek();
        report(token, ParseErrors.E_PARSE_EXPECTED_TOKEN, message + ", found " + token.kind());
        if (!cursor.isAtEnd()) {
            return cursor.advance();
        }
        return token;
    }

    private Span span(final long start, final long end) {
        return new Span(fileId, start, end);
    }

    private void report(final PbsToken token, final ParseErrors parseErrors, final String message) {
        p.studio.compiler.source.diagnostics.Diagnostics.error(diagnostics, parseErrors.name(), message, new Span(fileId, token.start(), token.end()));
    }

    private record ParsedReturnSpec(
            PbsAst.ReturnKind kind,
            PbsAst.TypeRef returnType,
            PbsAst.TypeRef resultErrorType) {
    }
}
