package cd.frontend.parser;

import cd.frontend.parser.JavaliParser.*;
import cd.ir.Ast;
import cd.ir.Ast.*;
import cd.ir.Ast.BinaryOp.BOp;
import cd.ir.Ast.UnaryOp.UOp;
import org.antlr.v4.runtime.tree.ParseTree;
import org.antlr.v4.runtime.tree.TerminalNode;

import java.util.ArrayList;
import java.util.List;

public final class JavaliAstVisitor extends JavaliBaseVisitor<Ast> {

	public List<ClassDecl> classDecls = new ArrayList<>();

	@Override
	public Ast visitUnit(UnitContext ctx) {
		for (ClassDeclContext classDeclContext : ctx.classDecl())
			classDecls.add((ClassDecl) visit(classDeclContext));
		return new Seq(new ArrayList<>(classDecls));
	}

	@Override
	public Ast visitClassDecl(ClassDeclContext ctx) {
        String name = ctx.Ident(0).getText();
        String superClass = "Object"; // Common superclass
        if (ctx.Ident().size() == 2)
            superClass = ctx.Ident(1).getText();
        Seq members = (Seq) visit(ctx.memberList());
		return new ClassDecl(name, superClass, members.children());
	}

	@Override
	public Ast visitMemberList(MemberListContext ctx) {
		if (ctx.children == null) {
			return new Seq(new ArrayList<>());
		}
		List<Ast> list = new ArrayList<>(ctx.children.size());
		for (ParseTree parseTree : ctx.children) {
			if (parseTree instanceof VarDeclContext) {
				Seq seqVars = (Seq) visit(parseTree);
				list.addAll(seqVars.children());
			} else {
				assert parseTree instanceof MethodDeclContext;
				list.add(visit(parseTree));
			}
		}
		return new Seq(list);
	}

	@Override
	public Ast visitVarDecl(VarDeclContext ctx) {
		List<Ast> list = new ArrayList<>(ctx.Ident().size());
		String type = ctx.type().getText();
		for (TerminalNode n : ctx.Ident())
			list.add(new VarDecl(type, n.getText()));
		return new Seq(list);
	}

	@Override
	public Ast visitMethodDecl(MethodDeclContext ctx) {
        List<String> argumentTypes = new ArrayList<>();
        List<String> argumentNames = new ArrayList<>();
        if (ctx.formalParamList() != null) {
            Seq paramList = (Seq) visit(ctx.formalParamList());
            for (Ast ast : paramList.children()) {
                VarDecl var = (VarDecl) ast;
                argumentNames.add(var.name);
                argumentTypes.add(var.type);
            }
        }

        List<Ast> decls = new ArrayList<>(ctx.varDecl().size());
        for (VarDeclContext varDecl : ctx.varDecl()) {
            Seq declarationSeq = (Seq) visit(varDecl);
            decls.addAll(declarationSeq.children());
        }

        List<Ast> stmts = new ArrayList<>(ctx.stmt().size());
        for (StmtContext s : ctx.stmt())
            stmts.add(visit(s));

        return new MethodDecl(
				ctx.type() == null ? "void" : ctx.type().getText(),
				ctx.Ident().getText(),
                argumentTypes,
                argumentNames,
                new Seq(decls),
                new Seq(stmts)
        );
	}

	@Override
	public Ast visitFormalParamList(FormalParamListContext ctx) {
		List<Ast> list = new ArrayList<>(ctx.type().size());
		for (int i = 0; i < ctx.type().size(); i++) {
		    String type = ctx.type(i).getText();
		    String name = ctx.Ident(i).getText();
		    list.add(new VarDecl(type, name));
        }
		return new Seq(list);
	}

	@Override
	public Ast visitStmt(StmtContext ctx) {
		return visit(ctx.children.get(0));
	}

	@Override
	public Ast visitStmtBlock(StmtBlockContext ctx) {
		List<Ast> list = new ArrayList<>(ctx.stmt().size());
		for (StmtContext stmtContext : ctx.stmt())
			list.add(visit(stmtContext));
        return new Seq(list);
	}

	@Override
	public Ast visitLocalMethodCallStmt(LocalMethodCallStmtContext ctx) {
		return new MethodCall(new MethodCallExpr(
				new ThisRef(),
				ctx.Ident().getText(),
				getParams(ctx.actualParamList())
		));
	}

	@Override
	public Ast visitObjectMethodCallStmt(ObjectMethodCallStmtContext ctx) {
        return new MethodCall(new MethodCallExpr(
				(Expr) visit(ctx.identAccess()),
				ctx.Ident().getText(),
				getParams(ctx.actualParamList())
		));
	}

	@Override
	public Ast visitAssignmentStmt(AssignmentStmtContext ctx) {
		Expr left  = (Expr) visit(ctx.children.get(0));
		Expr right = (Expr) visit(ctx.children.get(2));
		return new Assign(left, right);
	}

	@Override
	public Ast visitWrite(WriteContext ctx) {
		return new BuiltInWrite((Expr) visit(ctx.expr()));
	}

	@Override
	public Ast visitWriteLn(WriteLnContext ctx) {
		return new BuiltInWriteln();
	}

	@Override
	public Ast visitIfStmt(IfStmtContext ctx) {
		Expr condition = (Expr) visit(ctx.expr());
		Ast then = visit(ctx.stmtBlock(0));
		Ast otherwise = new Nop();
		if (ctx.stmtBlock().size() == 2)
			otherwise = visit(ctx.stmtBlock(1));
		return new IfElse(condition, then, otherwise);
	}

	@Override
	public Ast visitWhileStmt(WhileStmtContext ctx) {
		Expr condition = (Expr) visit(ctx.expr());
		Ast body = visit(ctx.stmtBlock());
		return new WhileLoop(condition, body);
	}

	@Override
	public Ast visitReturnStmt(ReturnStmtContext ctx) {
		if (ctx.expr() != null)
			return new ReturnStmt((Expr) visit(ctx.expr()));
		else
			return new ReturnStmt(null);
	}

	@Override
	public Ast visitNewObject(NewObjectContext ctx) {
		return new NewObject(ctx.Ident().getText());
	}

	@Override
	public Ast visitNewObjectArray(NewObjectArrayContext ctx) {
		String type = ctx.Ident().getText();
		Expr size = (Expr) visit(ctx.expr());
        return new NewArray(type + "[]", size);
	}

	@Override
	public Ast visitNewPrimitiveArray(NewPrimitiveArrayContext ctx) {
		String type = ctx.primitiveType().getText();
		Expr size = (Expr) visit(ctx.expr());
        return new NewArray(type + "[]", size);
	}

	@Override
	public Ast visitReadExpr(ReadExprContext ctx) {
        return new BuiltInRead();
	}

	@Override
	public Ast visitActualParamList(ActualParamListContext ctx) {
		List<Ast> list = new ArrayList<>(ctx.expr().size());
		for (ExprContext exprContext : ctx.expr())
			list.add(visit(exprContext));
        return new Seq(list);
	}

	@Override
	public Ast visitAccessThis(AccessThisContext ctx) {
		return new ThisRef();
	}

	@Override
	public Ast visitAccessLocalField(AccessLocalFieldContext ctx) {
        return new Var(ctx.Ident().getText());
	}

	@Override
	public Ast visitAccessObjectField(AccessObjectFieldContext ctx) {
		Expr arg = (Expr) visit(ctx.identAccess());
		String fieldName = ctx.Ident().getText();
        return new Field(arg, fieldName);
	}

	@Override
	public Ast visitAccessLocalMethod(AccessLocalMethodContext ctx) {
		return new MethodCallExpr(
        		new ThisRef(),
				ctx.Ident().getText(),
				getParams(ctx.actualParamList())
		);
	}

	@Override
	public Ast visitAccessArray(AccessArrayContext ctx) {
		Expr array = (Expr) visit(ctx.identAccess());
		Expr index = (Expr) visit(ctx.expr());
		return new Index(array, index);
	}

	@Override
	public Ast visitAccessObjectMethod(AccessObjectMethodContext ctx) {
		return new MethodCallExpr(
				(Expr) visit(ctx.identAccess()),
				ctx.Ident().getText(),
				getParams(ctx.actualParamList())
		);
	}

	@Override
	public Ast visitExprBinary(ExprBinaryContext ctx) {
		Expr left  = (Expr) visit(ctx.expr(0));
		Expr right = (Expr) visit(ctx.expr(1));
		String op  = ctx.getChild(1).getText();
		for (BOp bop : BOp.values())
			if (bop.repr.equals(op))
				return new BinaryOp(left, bop, right);
		throw new RuntimeException("BOp enum is inconsistent with grammar");
	}

	@Override
	public Ast visitExprCast(ExprCastContext ctx) {
        Expr arg = (Expr) visit(ctx.expr());
        String typeName = ctx.referenceType().getText();
        return new Cast(arg, typeName);
	}

	@Override
	public Ast visitExprParentheses(ExprParenthesesContext ctx) {
		return visit(ctx.expr());
	}

	@Override
	public Ast visitExprUnary(ExprUnaryContext ctx) {
		Expr expr = (Expr) visit(ctx.expr());
		String op = ctx.getChild(0).getText();
		for (UOp uop : UOp.values())
			if (uop.repr.equals(op))
				return new UnaryOp(uop, expr);
		throw new RuntimeException("UOp enum is inconsistent with grammar");
	}

	@Override
	public Ast visitExprConstant(ExprConstantContext ctx) {
		return visit(ctx.literal());
	}

	@Override
	public Ast visitNullLiteral(NullLiteralContext ctx) {
		return new NullConst();
	}

	@Override
	public Ast visitBoolLiteral(BoolLiteralContext ctx) {
		return new BooleanConst(ctx.Boolean().getText().equals("true"));
	}

	@Override
	public Ast visitIntLiteral(IntLiteralContext ctx) {
		try {
			return new IntConst(Integer.decode(ctx.Integer().getText()));
		} catch (NumberFormatException exception) {
			throw new ParseFailure(ctx.start.getLine(), "Value not in 32bit signed int range");
		}
	}

	@Override
	public Ast visitExprIdentAccess(ExprIdentAccessContext ctx) {
		return visit(ctx.identAccess());
	}

	/**
	 * Obtain the parameters for a method call
	 * @param paramListContext Formal parameter list from the ANTRL tree
	 * @return List of parameters for the method call
	 */
	private List<Expr> getParams(ActualParamListContext paramListContext) {
		List<Expr> paramList = new ArrayList<>();
		if (paramListContext != null) {
			Seq paramSeq = (Seq) visit(paramListContext);
			for (Ast ast : paramSeq.children())
				paramList.add((Expr) ast);
		}
		return paramList;
	}
}