package submit;

import parser.CminusBaseVisitor;
import parser.CminusParser;
import submit.ast.*;

import java.util.ArrayList;
import java.util.List;
import java.util.logging.Logger;
import java.util.stream.Collectors;

public class ASTVisitor extends CminusBaseVisitor<Node> {
    private final Logger LOGGER;
    private SymbolTable symbolTable;

    public ASTVisitor(Logger LOGGER) {
        this.LOGGER = LOGGER;
    }

    private VarType getVarType(CminusParser.TypeSpecifierContext ctx) {
        final String t = ctx.getText();
        return (t.equals("int")) ? VarType.INT : (t.equals("bool")) ? VarType.BOOL : VarType.CHAR;
    }

    @Override
    public Node visitProgram(CminusParser.ProgramContext ctx) {
        symbolTable = new SymbolTable();
        List<Declaration> decls = new ArrayList<>();
        for (CminusParser.DeclarationContext d : ctx.declaration())
            decls.add((Declaration) visitDeclaration(d));
        return new Program(decls);
    }

    @Override
    public Node visitVarDeclaration(CminusParser.VarDeclarationContext ctx) {
        VarType type = getVarType(ctx.typeSpecifier());
        List<String> ids = new ArrayList<>();
        List<Integer> arraySizes = new ArrayList<>();
        for (CminusParser.VarDeclIdContext v : ctx.varDeclId()) {
            String id = v.ID().getText();
            ids.add(id);
            symbolTable.addSymbol(id, new SymbolInfo(id, type, false));
            if (v.NUMCONST() != null)
                arraySizes.add(Integer.parseInt(v.NUMCONST().getText()));
            else
                arraySizes.add(-1);
        }
        final boolean isStatic = false;
        return new VarDeclaration(type, ids, arraySizes, isStatic);
    }

    @Override
    public Node visitReturnStmt(CminusParser.ReturnStmtContext ctx) {
        if (ctx.expression() != null)
            return new Return((Expression) visitExpression(ctx.expression()));
        return new Return(null);
    }

    @Override
    public Node visitConstant(CminusParser.ConstantContext ctx) {
        final Node node;
        if (ctx.NUMCONST() != null)
            node = new NumConstant(Integer.parseInt(ctx.NUMCONST().getText()));
        else if (ctx.CHARCONST() != null)
            node = new CharConstant(ctx.CHARCONST().getText().charAt(0));
        else if (ctx.STRINGCONST() != null)
            node = new StringConstant(ctx.STRINGCONST().getText());
        else
            node = new BoolConstant(ctx.getText().equals("true"));
        return node;
    }

    @Override
    public Node visitFunDeclaration(CminusParser.FunDeclarationContext ctx) {
        VarType rType = VarType.VOID;
        if (ctx.typeSpecifier() != null)
            rType = VarType.fromString(ctx.typeSpecifier().getText());
        String id = ctx.ID().getText();

        symbolTable.addSymbol(id, new SymbolInfo(id, rType, true));
        symbolTable = symbolTable.createChild();

        List<Param> params = ctx.param().stream().map(pCtx -> (Param) visitParam(pCtx)).collect(Collectors.toList());
        Statement statement = (Statement) visitStatement(ctx.statement());

        symbolTable = symbolTable.getParent();
        return new FunctionDeclaration(rType, id, params, statement);
    }

    @Override
    public Node visitParam(CminusParser.ParamContext ctx) {
        VarType pType = getVarType(ctx.typeSpecifier());

        String pId = ctx.paramId().getText();
        Boolean isArray = pId.endsWith("[]");
        if (isArray)
            pId = pId.substring(0, pId.length() - 2);

        symbolTable.addSymbol(pId, new SymbolInfo(pId, pType, false));

        return new Param(pType, pId, isArray);
    }

    @Override
    public Node visitStatement(CminusParser.StatementContext ctx) {
        if (ctx.expressionStmt() != null)
            return new ExpressionStatement((Expression) visitExpression(ctx.expressionStmt().expression()));
        if (ctx.compoundStmt() != null) {
            CminusParser.CompoundStmtContext cCtx = ctx.compoundStmt();

            symbolTable = symbolTable.createChild();

            List<VarDeclaration> varDecls = cCtx.varDeclaration().stream()
                    .map(declCtx -> (VarDeclaration) visitVarDeclaration(declCtx)).collect(Collectors.toList());

            List<Statement> statements = cCtx.statement().stream().map(stmtCtx -> (Statement) visitStatement(stmtCtx))
                    .collect(Collectors.toList());

            symbolTable = symbolTable.getParent();

            return new CompoundStatement(varDecls, statements);
        }
        if (ctx.ifStmt() != null) {
            List<Statement> statements = ctx.ifStmt().statement().stream()
                    .map(stmtCtx -> (Statement) visitStatement(stmtCtx)).collect(Collectors.toList());
            return new IfStatement((Expression) visitSimpleExpression(ctx.ifStmt().simpleExpression()), statements);
        }
        if (ctx.whileStmt() != null)
            return new WhileStatement((Expression) visitSimpleExpression(ctx.whileStmt().simpleExpression()),
                    (Statement) visitStatement(ctx.whileStmt().statement()));
        if (ctx.returnStmt() != null) {
            if (ctx.returnStmt().expression() != null)
                return new Return((Expression) visitExpression(ctx.returnStmt().expression()));
            return new Return(null);
        }
        if (ctx.breakStmt() != null)
            return new Break();
        return null;
    }

    @Override
    public Node visitExpression(CminusParser.ExpressionContext ctx) {
        if (ctx.mutable() != null && ctx.expression() != null)
            return new ExpressionExpression((Mutable) visitMutable(ctx.mutable()),
                    (Expression) visitExpression(ctx.expression()), ctx.getChild(1).getText());
        else if (ctx.mutable() != null)
            return new ExpressionExpression((Mutable) visitMutable(ctx.mutable()), ctx.getChild(1).getText());
        return visitSimpleExpression(ctx.simpleExpression());
    }

    @Override
    public Node visitMutable(CminusParser.MutableContext ctx) {
        String id = ctx.ID().getText();

        if (symbolTable.find(id) == null)
            LOGGER.warning("Undefined symbol on line " + ctx.getStart().getLine() + ": " + id);

        return new Mutable(id,
                ctx.expression() != null ? (Expression) visitExpression(ctx.expression()) : null);
    }

    @Override
    public Node visitOrExpression(CminusParser.OrExpressionContext ctx) {
        List<AndExpression> andExpressions = ctx.andExpression().stream().map(andExpression -> {
            List<UnaryRelExpression> uRE = andExpression.unaryRelExpression().stream()
                    .map(unaryRelExpressionCtx -> (UnaryRelExpression) visitUnaryRelExpression(unaryRelExpressionCtx))
                    .collect(Collectors.toList());
            return new AndExpression(uRE);
        }).collect(Collectors.toList());

        return new OrExpression(andExpressions);
    }

    @Override
    public Node visitUnaryRelExpression(CminusParser.UnaryRelExpressionContext ctx) {
        int bangs = ctx.BANG().size();
        RelExpression relExpression = (RelExpression) visitRelExpression(ctx.relExpression());
        return new UnaryRelExpression(bangs, relExpression);
    }

    @Override
    public Node visitRelExpression(CminusParser.RelExpressionContext ctx) {
        List<BinaryOperatorType> relOps = ctx.relop().stream()
                .map(operatorContext -> BinaryOperatorType.fromString(operatorContext.getText()))
                .collect(Collectors.toList());
        List<SumExpression> sumExpressions = ctx.sumExpression().stream().map(sumExpressionCtx -> {
            List<BinaryOperatorType> sumOps = sumExpressionCtx.sumop().stream()
                    .map(operatorContext -> BinaryOperatorType.fromString(operatorContext.getText()))
                    .collect(Collectors.toList());
            List<TermExpression> termExprs = sumExpressionCtx.termExpression().stream()
                    .map(termExprCtx -> (TermExpression) visitTermExpression(termExprCtx)).collect(Collectors.toList());
            return new SumExpression(sumOps, termExprs);
        }).collect(Collectors.toList());

        return new RelExpression(sumExpressions, relOps);
    }

    public Node visitTermExpression(CminusParser.TermExpressionContext ctx) {
        List<BinaryOperatorType> mulOps = ctx.mulop().stream()
                .map(operatorContext -> BinaryOperatorType.fromString(operatorContext.getText()))
                .collect(Collectors.toList());
        List<UnaryExpression> unaryExpressions = ctx.unaryExpression().stream().map(unaryExpressionCtx -> {
            List<String> unaryOps = unaryExpressionCtx.unaryop().stream()
                    .map(operatorContext -> operatorContext.getText())
                    .collect(Collectors.toList());
            Factor factor = (Factor) visitFactor(unaryExpressionCtx.factor());

            return new UnaryExpression(unaryOps, factor);
        }).collect(Collectors.toList());

        return new TermExpression(unaryExpressions, mulOps);
    }

    public Node visitFactor(CminusParser.FactorContext ctx) {
        if (ctx.immutable() != null)
            return new Factor((Immutable) visitImmutable(ctx.immutable()));
        return new Factor((Mutable) visitMutable(ctx.mutable()));
    }

    public Node visitImmutable(CminusParser.ImmutableContext ctx) {
        if (ctx.constant() != null)
            return new Immutable((Constant) visitConstant(ctx.constant()));

        if (ctx.call() != null)
            return new Immutable((Call) visitCall(ctx.call()), true);

        return new Immutable((Expression) visitExpression(ctx.expression()));
    }

    public Node visitCall(CminusParser.CallContext ctx) {
        String id = ctx.ID().getText();

        if (symbolTable.find(id) == null)
            LOGGER.warning("Undefined symbol on line " + ctx.getStart().getLine() + ": " + id);

        return new Call(id, ctx.expression().stream()
                .map(exprCtx -> (Expression) visitExpression(exprCtx)).collect(Collectors.toList()));
    }
}