package xyz.ignatz.antlr.java9.recognizer;

import xyz.ignatz.antlr.java9.utils.JavaName;
import xyz.ignatz.flowchart.Flowchart;
import xyz.ignatz.flowchart.module.Block;
import xyz.ignatz.flowchart.module.Node;

import java.util.*;

import xyz.ignatz.antlr.java9.utils.Blocks;


public class Java9FlowchartVisitor extends Java9BaseVisitor<Block> {
    List<Flowchart> flowcharts = new ArrayList<>();
    JavaName name = new JavaName();

    public List<Flowchart> getFlowcharts() {
        return new ArrayList<>(flowcharts);
    }

    @Override
    protected Block aggregateResult(Block block, Block next) {
        if (block == null) return next;
        block.connect(next);
        return block;
    }

    @Override
    public Block visitIfThenStatement(Java9Parser.IfThenStatementContext ctx) {
        Block body = visitStatement(ctx.statement());
        Block block = Blocks.buildJudgement(ctx.expression());
        block.connect(body, "true");
        block.modifyTpType("false", "");
        return block;
    }

    @Override
    public Block visitIfThenElseStatement(Java9Parser.IfThenElseStatementContext ctx) {
        Block bodyTrue = visitStatementNoShortIf(ctx.statementNoShortIf());
        Block bodyFalse = visitStatement(ctx.statement());
        Block block = Blocks.buildJudgement(ctx.expression());
        block.connect(bodyTrue, "true");
        block.connect(bodyFalse, "false");
        return block;
    }

    @Override
    public Block visitIfThenElseStatementNoShortIf(Java9Parser.IfThenElseStatementNoShortIfContext ctx) {
        Block bodyTrue = visitStatementNoShortIf(ctx.statementNoShortIf().get(0));
        Block bodyFalse = visitStatementNoShortIf(ctx.statementNoShortIf().get(1));
        Block block = Blocks.buildJudgement(ctx.expression());
        block.connect(bodyTrue, "true");
        block.connect(bodyFalse, "false");
        return block;
    }

    @Override
    public Block visitSwitchStatement(Java9Parser.SwitchStatementContext ctx) {
        Map<String, Node> map = new HashMap<>();
        Block body = null;
        for (Java9Parser.SwitchBlockStatementGroupContext group : ctx.switchBlock().switchBlockStatementGroup()) {
            Block block = visitBlockStatements(group.blockStatements());
            Node start = block.getStart();
            group.switchLabels().switchLabel().forEach(lab -> map.put(lab.getText(), start));
            body = aggregateResult(body, block);
        }
        Block block = Blocks.build("switch (" + ctx.expression().getText() + ")", ctx.getStart().getLine(), map);
        block.connect(body);
        block.modifyTpType("break", "");
        return block;
    }

    @Override
    public Block visitDoStatement(Java9Parser.DoStatementContext ctx) {
        Block block = visitStatement(ctx.statement());
        Block judge = Blocks.buildJudgement(ctx.expression());
        block.connectLoopJudgement(judge, false);
        return block;
    }

    @Override
    public Block visitBasicForStatement(Java9Parser.BasicForStatementContext ctx) {
        Block init = Blocks.build(ctx.forInit());
        Block judge = Blocks.buildJudgement(ctx.expression(), ctx.getStart().getLine());
        Block update = Blocks.build(ctx.forUpdate());
        Block body = visitStatement(ctx.statement());
        body.connect(update);
        body.connectLoopJudgement(judge, true);
        return aggregateResult(init, body);
    }

    @Override
    public Block visitBasicForStatementNoShortIf(Java9Parser.BasicForStatementNoShortIfContext ctx) {
        Block init = Blocks.build(ctx.forInit());
        Block judge = Blocks.buildJudgement(ctx.expression(), ctx.getStart().getLine());
        Block update = visitForUpdate(ctx.forUpdate());
        Block body = visitStatementNoShortIf(ctx.statementNoShortIf());
        body.connect(update);
        body.connectLoopJudgement(judge, true);
        return aggregateResult(init, body);
    }

    @Override
    public Block visitWhileStatement(Java9Parser.WhileStatementContext ctx) {
        Block judge = Blocks.buildJudgement(ctx.expression());
        Block body = visitStatement(ctx.statement());
        body.connectLoopJudgement(judge, true);
        return body;
    }

    @Override
    public Block visitWhileStatementNoShortIf(Java9Parser.WhileStatementNoShortIfContext ctx) {
        Block judge = Blocks.buildJudgement(ctx.expression());
        Block body = visitStatementNoShortIf(ctx.statementNoShortIf());
        body.connectLoopJudgement(judge, true);
        return body;
    }

    @Override
    public Block visitNormalClassDeclaration(Java9Parser.NormalClassDeclarationContext ctx) {
        name.push('$', ctx.identifier().getText());
        super.visitNormalClassDeclaration(ctx);
        name.pop('$');
        return null;
    }

    @Override
    public Block visitEnumDeclaration(Java9Parser.EnumDeclarationContext ctx) {
        name.push('$', ctx.identifier().getText());
        super.visitEnumDeclaration(ctx);
        name.pop('$');
        return null;
    }

    @Override
    public Block visitNormalInterfaceDeclaration(Java9Parser.NormalInterfaceDeclarationContext ctx) {
        name.push('$', ctx.identifier().getText());
        super.visitNormalInterfaceDeclaration(ctx);
        name.pop('$');
        return null;
    }

    @Override
    public Block visitAnnotationTypeDeclaration(Java9Parser.AnnotationTypeDeclarationContext ctx) {
        name.push('$', ctx.identifier().getText());
        super.visitAnnotationTypeDeclaration(ctx);
        name.pop('$');
        return null;
    }

    @Override
    public Block visitMethodDeclaration(Java9Parser.MethodDeclarationContext ctx) {
        name.push('#', ctx.methodHeader().methodDeclarator().identifier().getText());
        Block block = Blocks.build(name.toString(), ctx.getStart().getLine());
        Block body = visitMethodBody(ctx.methodBody());
        block.connect(body);
        flowcharts.add(new Flowchart(name.toString(), block));
        name.pop('#');
        return null;
    }

    @Override
    public Block visitLocalVariableDeclaration(Java9Parser.LocalVariableDeclarationContext ctx) {
        return Blocks.build(ctx);
    }

    @Override
    public Block visitStatementExpressionList(Java9Parser.StatementExpressionListContext ctx) {
        return Blocks.build(ctx);
    }

    @Override
    public Block visitEmptyStatement(Java9Parser.EmptyStatementContext ctx) {
        return Blocks.build(ctx);
    }

    @Override
    public Block visitExpressionStatement(Java9Parser.ExpressionStatementContext ctx) {
        return Blocks.build(ctx);
    }

    @Override
    public Block visitAssertStatement(Java9Parser.AssertStatementContext ctx) {
        List<Java9Parser.ExpressionContext> exp = ctx.expression();
        Block judge = Blocks.buildJudgement(exp.get(0));
        String info = exp.size() == 1 ? "assert !(" + exp.get(0).getText() + ")" : "assert " + exp.get(1).getText();
        Block ast  = Blocks.build(info, ctx.start.getLine(), "assert failed");
        judge.modifyTpType("true", "");
        judge.connect(ast, "false");
        return judge;
        // TODO: assertStatement
    }

    @Override
    public Block visitBreakStatement(Java9Parser.BreakStatementContext ctx) {
        String label = ctx.identifier() == null ? "" : ctx.identifier().getText();
        Block block = Blocks.build(ctx, label);
        block.modifyTpType(label, "break");
        return block;
        // TODO: break label
    }

    @Override
    public Block visitContinueStatement(Java9Parser.ContinueStatementContext ctx) {
        String label = ctx.identifier() == null ? "" : ctx.identifier().getText();
        Block block = Blocks.build(ctx, label);
        block.modifyTpType(label, "continue");
        return block;
        // TODO: continue label
    }

    @Override
    public Block visitReturnStatement(Java9Parser.ReturnStatementContext ctx) {
        return Blocks.build(ctx, "return");
    }
}
