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 PbsBlockParser {
    private static final boolean ALLOW_TAIL_EXPRESSION = true;

    private final PbsParserContext context;
    private final PbsTokenCursor cursor;
    private final PbsExprParser exprParser;
    private final PbsTypeParser typeParser;

    PbsBlockParser(
            final PbsParserContext context,
            final PbsExprParser exprParser,
            final PbsTypeParser typeParser) {
        this.context = context;
        this.cursor = context.cursor();
        this.exprParser = exprParser;
        this.typeParser = typeParser;
    }

    PbsAst.Block parseBlock() {
        return parseBlock("Expected '{' to start block");
    }

    PbsAst.Block parseBlock(final String message) {
        final var leftBrace = consume(PbsTokenKind.LEFT_BRACE, message);
        final var statements = new ArrayList<PbsAst.Statement>();
        PbsAst.Expression tailExpression = null;
        while (!cursor.check(PbsTokenKind.RIGHT_BRACE) && !cursor.isAtEnd()) {
            if (startsStructuredStatement() || isAssignmentStatementStart()) {
                statements.add(parseStatement());
                continue;
            }

            final var expression = exprParser.parseExpression();
            if (cursor.match(PbsTokenKind.SEMICOLON)) {
                statements.add(new PbsAst.ExpressionStatement(
                        expression,
                        context.span(expression.span().getStart(), cursor.previous().end())));
                continue;
            }

            if (ALLOW_TAIL_EXPRESSION && cursor.check(PbsTokenKind.RIGHT_BRACE)) {
                tailExpression = expression;
                break;
            }

            context.report(cursor.peek(), ParseErrors.E_PARSE_EXPECTED_TOKEN, "Expected ';' after expression");
            if (!cursor.isAtEnd() && !cursor.check(PbsTokenKind.RIGHT_BRACE)) {
                cursor.advance();
            }
        }
        final var rightBrace = consume(PbsTokenKind.RIGHT_BRACE, "Expected '}' to end block");
        return new PbsAst.Block(
                ReadOnlyList.wrap(statements),
                tailExpression,
                context.span(leftBrace.start(), rightBrace.end()));
    }

    private boolean startsStructuredStatement() {
        return cursor.check(PbsTokenKind.LET)
                || cursor.check(PbsTokenKind.IF)
                || cursor.check(PbsTokenKind.FOR)
                || cursor.check(PbsTokenKind.WHILE)
                || cursor.check(PbsTokenKind.BREAK)
                || cursor.check(PbsTokenKind.CONTINUE)
                || cursor.check(PbsTokenKind.RETURN);
    }

    private PbsAst.Statement parseStatement() {
        if (cursor.match(PbsTokenKind.LET)) {
            return parseLetStatement(cursor.previous());
        }
        if (cursor.match(PbsTokenKind.IF)) {
            return parseIfStatement(cursor.previous());
        }
        if (cursor.match(PbsTokenKind.FOR)) {
            return parseForStatement(cursor.previous());
        }
        if (cursor.match(PbsTokenKind.WHILE)) {
            return parseWhileStatement(cursor.previous());
        }
        if (cursor.match(PbsTokenKind.BREAK)) {
            return parseBreakStatement(cursor.previous());
        }
        if (cursor.match(PbsTokenKind.CONTINUE)) {
            return parseContinueStatement(cursor.previous());
        }
        if (cursor.match(PbsTokenKind.RETURN)) {
            return parseReturnStatement(cursor.previous());
        }
        if (isAssignmentStatementStart()) {
            return parseAssignStatement();
        }
        return parseExpressionStatement();
    }

    private PbsAst.Statement parseLetStatement(final PbsToken letToken) {
        final boolean isConst = cursor.match(PbsTokenKind.CONST);
        final var name = consume(PbsTokenKind.IDENTIFIER, "Expected variable name");

        PbsAst.TypeRef explicitType = null;
        if (cursor.match(PbsTokenKind.COLON)) {
            explicitType = typeParser.parseTypeRef();
        }

        consume(PbsTokenKind.EQUAL, "Expected '=' in let statement");
        final var initializer = exprParser.parseExpression();
        final var semicolon = consume(PbsTokenKind.SEMICOLON, "Expected ';' after let statement");

        return new PbsAst.LetStatement(
                isConst,
                name.lexeme(),
                context.span(name.start(), name.end()),
                explicitType,
                initializer,
                context.span(letToken.start(), semicolon.end()));
    }

    private PbsAst.Statement parseAssignStatement() {
        final var lValue = parseLValue();
        final var operator = parseAssignOperator(consumeAssignOperator());
        final var value = exprParser.parseExpression();
        final var semicolon = consume(PbsTokenKind.SEMICOLON, "Expected ';' after assignment");
        return new PbsAst.AssignStatement(
                lValue,
                operator,
                value,
                context.span(lValue.span().getStart(), semicolon.end()));
    }

    private PbsAst.LValue parseLValue() {
        final PbsToken root;
        if (cursor.match(PbsTokenKind.IDENTIFIER, PbsTokenKind.THIS)) {
            root = cursor.previous();
        } else {
            root = consume(PbsTokenKind.IDENTIFIER, "Expected assignment target identifier");
        }
        final var segments = new ArrayList<String>();
        var end = root.end();
        while (cursor.match(PbsTokenKind.DOT)) {
            final var segment = consume(PbsTokenKind.IDENTIFIER, "Expected member name after '.' in assignment target");
            segments.add(segment.lexeme());
            end = segment.end();
        }
        return new PbsAst.LValue(root.lexeme(), ReadOnlyList.wrap(segments), context.span(root.start(), end));
    }

    private PbsToken consumeAssignOperator() {
        if (cursor.match(PbsTokenKind.EQUAL,
                PbsTokenKind.PLUS_EQUAL,
                PbsTokenKind.MINUS_EQUAL,
                PbsTokenKind.STAR_EQUAL,
                PbsTokenKind.SLASH_EQUAL,
                PbsTokenKind.PERCENT_EQUAL)) {
            return cursor.previous();
        }
        final var token = cursor.peek();
        context.report(token, ParseErrors.E_PARSE_INVALID_ASSIGN_TARGET,
                "Expected assignment operator after assignment target");
        if (!cursor.isAtEnd()) {
            return cursor.advance();
        }
        return token;
    }

    private PbsAst.AssignOperator parseAssignOperator(final PbsToken token) {
        return switch (token.kind()) {
            case EQUAL -> PbsAst.AssignOperator.ASSIGN;
            case PLUS_EQUAL -> PbsAst.AssignOperator.ADD_ASSIGN;
            case MINUS_EQUAL -> PbsAst.AssignOperator.SUB_ASSIGN;
            case STAR_EQUAL -> PbsAst.AssignOperator.MUL_ASSIGN;
            case SLASH_EQUAL -> PbsAst.AssignOperator.DIV_ASSIGN;
            case PERCENT_EQUAL -> PbsAst.AssignOperator.MOD_ASSIGN;
            default -> PbsAst.AssignOperator.ASSIGN;
        };
    }

    private PbsAst.IfStatement parseIfStatement(final PbsToken ifToken) {
        final var condition = exprParser.parseExpression();
        final var thenBlock = parseBlock();

        PbsAst.IfStatement elseIf = null;
        PbsAst.Block elseBlock = null;
        if (cursor.match(PbsTokenKind.ELSE)) {
            if (cursor.match(PbsTokenKind.IF)) {
                elseIf = parseIfStatement(cursor.previous());
            } else {
                elseBlock = parseBlock();
            }
        }

        final var end = elseIf != null
                ? elseIf.span().getEnd()
                : (elseBlock != null ? elseBlock.span().getEnd() : thenBlock.span().getEnd());

        return new PbsAst.IfStatement(
                condition,
                thenBlock,
                elseIf,
                elseBlock,
                context.span(ifToken.start(), end));
    }

    private PbsAst.Statement parseForStatement(final PbsToken forToken) {
        final var iterator = consume(PbsTokenKind.IDENTIFIER, "Expected loop iterator name in 'for'");
        consumeForToken(PbsTokenKind.COLON, "Expected ':' after iterator name in 'for'");
        final var iteratorType = typeParser.parseTypeRef();
        consumeForToken(PbsTokenKind.FROM, "Expected 'from' in 'for' statement");
        final var fromExpression = exprParser.parseExpression();
        consumeForToken(PbsTokenKind.UNTIL, "Expected 'until' in 'for' statement");
        final var untilExpression = exprParser.parseExpression();

        PbsAst.Expression stepExpression = null;
        if (cursor.match(PbsTokenKind.STEP)) {
            stepExpression = exprParser.parseExpression();
        }

        context.enterLoop();
        final var body = parseBlock();
        context.exitLoop();

        return new PbsAst.ForStatement(
                iterator.lexeme(),
                iteratorType,
                fromExpression,
                untilExpression,
                stepExpression,
                body,
                context.span(forToken.start(), body.span().getEnd()));
    }

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

    private PbsAst.Statement parseWhileStatement(final PbsToken whileToken) {
        final var condition = exprParser.parseExpression();
        context.enterLoop();
        final var body = parseBlock();
        context.exitLoop();
        return new PbsAst.WhileStatement(condition, body, context.span(whileToken.start(), body.span().getEnd()));
    }

    private PbsAst.Statement parseBreakStatement(final PbsToken breakToken) {
        final var semicolon = consume(PbsTokenKind.SEMICOLON, "Expected ';' after 'break'");
        if (!context.isInsideLoop()) {
            context.report(breakToken, ParseErrors.E_PARSE_LOOP_CONTROL_OUTSIDE_LOOP,
                    "'break' is only valid inside loop contexts");
        }
        return new PbsAst.BreakStatement(context.span(breakToken.start(), semicolon.end()));
    }

    private PbsAst.Statement parseContinueStatement(final PbsToken continueToken) {
        final var semicolon = consume(PbsTokenKind.SEMICOLON, "Expected ';' after 'continue'");
        if (!context.isInsideLoop()) {
            context.report(continueToken, ParseErrors.E_PARSE_LOOP_CONTROL_OUTSIDE_LOOP,
                    "'continue' is only valid inside loop contexts");
        }
        return new PbsAst.ContinueStatement(context.span(continueToken.start(), semicolon.end()));
    }

    private PbsAst.Statement parseReturnStatement(final PbsToken returnToken) {
        PbsAst.Expression value = null;
        if (!cursor.check(PbsTokenKind.SEMICOLON)) {
            value = exprParser.parseExpression();
        }
        final var semicolon = consume(PbsTokenKind.SEMICOLON, "Expected ';' after return");
        return new PbsAst.ReturnStatement(value, context.span(returnToken.start(), semicolon.end()));
    }

    private PbsAst.Statement parseExpressionStatement() {
        final var expression = exprParser.parseExpression();
        final var semicolon = consume(PbsTokenKind.SEMICOLON, "Expected ';' after expression");
        return new PbsAst.ExpressionStatement(expression, context.span(expression.span().getStart(), semicolon.end()));
    }

    private boolean isAssignmentStatementStart() {
        if (!cursor.check(PbsTokenKind.IDENTIFIER) && !cursor.check(PbsTokenKind.THIS)) {
            return false;
        }

        int offset = 1;
        while (cursor.peek(offset).kind() == PbsTokenKind.DOT) {
            if (cursor.peek(offset + 1).kind() != PbsTokenKind.IDENTIFIER) {
                return false;
            }
            if (cursor.peek(offset + 2).kind() == PbsTokenKind.LEFT_PAREN) {
                return false;
            }
            offset += 2;
        }

        return switch (cursor.peek(offset).kind()) {
            case EQUAL, PLUS_EQUAL, MINUS_EQUAL, STAR_EQUAL, SLASH_EQUAL, PERCENT_EQUAL -> true;
            default -> false;
        };
    }

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