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.utilities.structures.ReadOnlyList;

import java.util.ArrayList;

final class PbsExprControlFlowParser {
    @FunctionalInterface
    interface ExpressionParserDelegate {
        PbsAst.Expression parse();
    }

    private final PbsExprParserContext context;
    private final PbsTokenCursor cursor;
    private final ExpressionParserDelegate expressionParserDelegate;
    private final ExpressionParserDelegate precedenceParserDelegate;
    private final PbsExprLiteralSupport literalSupport;

    PbsExprControlFlowParser(
            final PbsExprParserContext context,
            final ExpressionParserDelegate expressionParserDelegate,
            final ExpressionParserDelegate precedenceParserDelegate) {
        this.context = context;
        this.cursor = context.cursor();
        this.expressionParserDelegate = expressionParserDelegate;
        this.precedenceParserDelegate = precedenceParserDelegate;
        this.literalSupport = new PbsExprLiteralSupport(context);
    }

    PbsAst.Expression parseHandle() {
        if (!cursor.match(PbsTokenKind.HANDLE)) {
            return parseElse();
        }

        final var handleToken = cursor.previous();
        final var value = parseElse();
        context.consume(PbsTokenKind.LEFT_BRACE, "Expected '{' after handle expression");

        final var arms = new ArrayList<PbsAst.HandleArm>();
        while (!cursor.check(PbsTokenKind.RIGHT_BRACE) && !cursor.isAtEnd()) {
            final var pattern = parseHandlePattern();
            context.consume(PbsTokenKind.ARROW, "Expected '->' in handle arm");

            PbsAst.ErrorPath remapTarget = null;
            PbsAst.Block block = null;
            long armEnd;
            if (cursor.check(PbsTokenKind.LEFT_BRACE)) {
                block = context.parseSurfaceBlock("Expected handle arm block");
                armEnd = block.span().getEnd();
            } else {
                remapTarget = parseErrorPath();
                armEnd = remapTarget.span().getEnd();
            }

            arms.add(new PbsAst.HandleArm(
                    pattern,
                    remapTarget,
                    block,
                    context.span(pattern.span().getStart(), armEnd)));

            if (!cursor.match(PbsTokenKind.COMMA)) {
                break;
            }
        }

        final var close = context.consume(PbsTokenKind.RIGHT_BRACE, "Expected '}' to close handle map");
        return new PbsAst.HandleExpr(value, ReadOnlyList.wrap(arms), context.span(handleToken.start(), close.end()));
    }

    PbsAst.Expression parseElse() {
        var expression = parseIfExpression();
        if (cursor.match(PbsTokenKind.ELSE)) {
            final var fallback = parseElse();
            expression = new PbsAst.ElseExpr(
                    expression,
                    fallback,
                    context.span(expression.span().getStart(), fallback.span().getEnd()));
        }
        return expression;
    }

    PbsAst.Expression parseIfExpression() {
        if (!cursor.match(PbsTokenKind.IF)) {
            return parseSwitchExpression();
        }
        return parseIfExpressionFromToken(cursor.previous());
    }

    PbsAst.IfExpr parseIfExpressionFromToken(final PbsToken ifToken) {
        final var condition = expressionParserDelegate.parse();
        final var thenBlock = context.parseSurfaceBlock("Expected '{' after if condition");

        if (!cursor.match(PbsTokenKind.ELSE)) {
            context.report(cursor.peek(), ParseErrors.E_PARSE_UNEXPECTED_TOKEN,
                    "If expression requires an else branch");
            final var unit = new PbsAst.UnitExpr(context.span(thenBlock.span().getEnd(), thenBlock.span().getEnd()));
            return new PbsAst.IfExpr(condition, thenBlock, unit, context.span(ifToken.start(), unit.span().getEnd()));
        }

        final PbsAst.Expression elseExpression;
        if (cursor.check(PbsTokenKind.IF)) {
            elseExpression = parseIfExpressionFromToken(cursor.advance());
        } else {
            final var elseBlock = context.parseSurfaceBlock("Expected '{' after else");
            elseExpression = new PbsAst.BlockExpr(elseBlock, elseBlock.span());
        }

        return new PbsAst.IfExpr(
                condition,
                thenBlock,
                elseExpression,
                context.span(ifToken.start(), elseExpression.span().getEnd()));
    }

    PbsAst.Expression parseSwitchExpression() {
        if (!cursor.match(PbsTokenKind.SWITCH)) {
            return precedenceParserDelegate.parse();
        }

        final var switchToken = cursor.previous();
        final var selector = expressionParserDelegate.parse();
        context.consume(PbsTokenKind.LEFT_BRACE, "Expected '{' after switch selector");

        final var arms = new ArrayList<PbsAst.SwitchArm>();
        while (!cursor.check(PbsTokenKind.RIGHT_BRACE) && !cursor.isAtEnd()) {
            final var pattern = parseSwitchPattern();
            context.consume(PbsTokenKind.COLON, "Expected ':' after switch pattern");
            final var block = context.parseSurfaceBlock("Expected switch arm block");
            arms.add(new PbsAst.SwitchArm(pattern, block, context.span(pattern.span().getStart(), block.span().getEnd())));

            if (!cursor.match(PbsTokenKind.COMMA)) {
                break;
            }
        }

        final var close = context.consume(PbsTokenKind.RIGHT_BRACE, "Expected '}' to close switch expression");
        return new PbsAst.SwitchExpr(selector, ReadOnlyList.wrap(arms), context.span(switchToken.start(), close.end()));
    }

    PbsAst.SwitchPattern parseSwitchPattern() {
        if (cursor.match(PbsTokenKind.DEFAULT)) {
            return new PbsAst.WildcardSwitchPattern(
                    PbsAst.WildcardSwitchKind.DEFAULT,
                    context.span(cursor.previous().start(), cursor.previous().end()));
        }

        if (cursor.check(PbsTokenKind.IDENTIFIER) && "_".equals(cursor.peek().lexeme())) {
            final var wildcard = cursor.advance();
            return new PbsAst.WildcardSwitchPattern(
                    PbsAst.WildcardSwitchKind.UNDERSCORE,
                    context.span(wildcard.start(), wildcard.end()));
        }

        if (cursor.check(PbsTokenKind.IDENTIFIER)
                && cursor.peek(1).kind() == PbsTokenKind.DOT
                && cursor.peek(2).kind() == PbsTokenKind.IDENTIFIER) {
            final var path = parseErrorPath();
            return new PbsAst.EnumCaseSwitchPattern(path, path.span());
        }

        final var literal = parseLiteralPatternExpression();
        if (literal != null) {
            return new PbsAst.LiteralSwitchPattern(literal, literal.span());
        }

        final var token = cursor.peek();
        context.report(token, ParseErrors.E_PARSE_INVALID_SWITCH_FORM, "Invalid switch pattern");
        cursor.advance();
        return new PbsAst.WildcardSwitchPattern(
                PbsAst.WildcardSwitchKind.RECOVERY,
                context.span(token.start(), token.end()));
    }

    PbsAst.Expression parseLiteralPatternExpression() {
        return literalSupport.parseLiteralExpression();
    }

    PbsAst.HandlePattern parseHandlePattern() {
        if (cursor.check(PbsTokenKind.IDENTIFIER) && "_".equals(cursor.peek().lexeme())) {
            final var wildcard = cursor.advance();
            return new PbsAst.WildcardHandlePattern(context.span(wildcard.start(), wildcard.end()));
        }

        if (cursor.check(PbsTokenKind.IDENTIFIER)) {
            final var path = parseErrorPath();
            return new PbsAst.ErrorPathHandlePattern(path, path.span());
        }

        final var token = cursor.peek();
        context.report(token, ParseErrors.E_PARSE_INVALID_HANDLE_FORM, "Invalid handle pattern");
        if (!cursor.isAtEnd()) {
            cursor.advance();
        }
        return new PbsAst.WildcardHandlePattern(context.span(token.start(), token.end()));
    }

    PbsAst.ErrorPath parseErrorPath() {
        final var first = context.consume(PbsTokenKind.IDENTIFIER, "Expected identifier in error path");
        final var segments = new ArrayList<String>();
        segments.add(first.lexeme());
        var end = first.end();
        while (cursor.match(PbsTokenKind.DOT)) {
            final var segment = context.consume(PbsTokenKind.IDENTIFIER, "Expected identifier after '.' in error path");
            segments.add(segment.lexeme());
            end = segment.end();
        }
        return new PbsAst.ErrorPath(ReadOnlyList.wrap(segments), context.span(first.start(), end));
    }

}
