package mc.astgen
import org.antlr.v4.runtime.Token
import org.antlr.v4.runtime.CommonTokenStream
import org.antlr.v4.runtime.ParserRuleContext
import java.io.{PrintWriter,File}
import org.antlr.v4.runtime.ANTLRFileStream
import mc.utils._
import scala.collection.JavaConverters._
import org.antlr.v4.runtime.tree._
import mc.parser._
import mc.parser.MCParser._

class ASTGeneration extends MCBaseVisitor[Any] {
	def higAppend (a: List[AST] , b: List[AST]) = higReverse(a).foldLeft(b) ((x,y) => y::x)
	def higReverse (a : List[AST]) = a.foldLeft(List[AST]())((x,y)=> y::x)
//program : declaration+ EOF ;
	override def visitProgram(ctx:ProgramContext) =
		Program(ctx.declaration.asScala.toList.foldLeft(List[Decl]())(
			(a,b) => higAppend(a,b.accept(this).asInstanceOf[List[Decl]]).asInstanceOf[List[Decl]]
		))

// declaration: var_declr | fun_declr;
	override def visitDeclaration(ctx : DeclarationContext) =
		if(ctx.var_declr != null)ctx.var_declr.accept(this)
		else List(ctx.fun_declr.accept(this))



// var_declr: pri_type variable (COM variable)* SEMI;
	override def visitVar_declr(ctx: Var_declrContext) =
		// int a,b,c; => int a, int b và int c
		ctx.variable.asScala.toList.map(x =>
			if (x.getChildCount() == 1) VarDecl(x.accept(this).asInstanceOf[Id], ctx.pri_type.accept(this).asInstanceOf[Type])
			else VarDecl(x.accept(this).asInstanceOf[Id], ArrayType(IntLiteral(x.INTLIT.getText.toInt) ,ctx.pri_type.accept(this).asInstanceOf[Type]))
		)


// pri_type: INTTYPE | STRINGTYPE | BOOLEANTYPE | FLOATTYPE ;
	override def visitPri_type(ctx : Pri_typeContext) =
		if(ctx.INTTYPE != null ) IntType
		else if(ctx.STRINGTYPE != null) StringType
		else if (ctx.BOOLEANTYPE != null) BoolType
		else FloatType


// variable : ID | ID LSB INTLIT RSB ;
	override def visitVariable(ctx: VariableContext) =  Id(ctx.ID.getText)


// fun_declr : func_type ID LB (para_declr (COM para_declr)* )? RB block_stmt ;
	override def visitFun_declr(ctx: Fun_declrContext) =
		FuncDecl(Id(ctx.ID.getText),
			ctx.para_declr.asScala.toList.foldLeft(List[VarDecl]())(
				 (a,b) => higAppend(a,b.accept(this).asInstanceOf[List[VarDecl]]).asInstanceOf[List[VarDecl]]),
			ctx.func_type.accept(this).asInstanceOf[Type],
			ctx.block_stmt.accept(this).asInstanceOf[Stmt]
		)

// func_type : pri_type | arr_pointer_type | VOIDTYPE ;
	override def visitFunc_type (ctx: Func_typeContext) =
		if(ctx.pri_type !=null) ctx.pri_type.accept(this)
		else if (ctx.VOIDTYPE != null) VoidType
		else ctx.arr_pointer_type.accept(this)

// arr_pointer_type : pri_type LSB RSB;
	override def visitArr_pointer_type( ctx:Arr_pointer_typeContext) =
		ArrayPointerType(ctx.pri_type.accept(this).asInstanceOf[Type])

// para_declr: (pri_type ID) | (pri_type ID LSB RSB) ;;
	override def visitPara_declr (ctx: Para_declrContext) =
		if(ctx.getChildCount == 2) List(VarDecl(Id(ctx.ID.getText), ctx.pri_type.accept(this).asInstanceOf[Type]))
		else List(VarDecl(Id(ctx.ID.getText) ,ArrayPointerType(ctx.pri_type.accept(this).asInstanceOf[Type])))

// block_stmt : LP var_declr* stmt* RP ;
	override def visitBlock_stmt( ctx: Block_stmtContext) =
		Block(
			ctx.var_declr.asScala.toList.foldLeft(List[Decl]())(
				(a,b) => higAppend(a,b.accept(this).asInstanceOf[List[Decl]]).asInstanceOf[List[Decl]]
			),
			ctx.stmt.asScala.toList.map(_.accept(this).asInstanceOf[Stmt])
		)

// stmt : if_stmt | do_stmt | for_stmt | break_stmt | continue_stmt | return_stmt | exp_stmt | block_stmt;
	override def visitStmt (ctx: StmtContext) =
		ctx.getChild(0).accept(this)
// if_stmt : IF LB exp RB stmt ELSE stmt | IF LB exp RB stmt;
	override def visitIf_stmt (ctx:If_stmtContext) =
		if(ctx.getChildCount == 7)
			If(
				ctx.exp.accept(this).asInstanceOf[Expr],
				ctx.stmt(0).accept(this).asInstanceOf[Stmt],
				Some(ctx.stmt(1).accept(this).asInstanceOf[Stmt])
			)
		else
			If(
				ctx.exp.accept(this).asInstanceOf[Expr],
				ctx.stmt(0).accept(this).asInstanceOf[Stmt],
				None
			)

// do_stmt : DO stmt* WHILE  exp SEMI ;
	override def visitDo_stmt (ctx: Do_stmtContext) =
		Dowhile(
			ctx.stmt.asScala.toList.foldLeft(List[Stmt]())(
				(a,b) => higAppend(a,List(b.accept(this)).asInstanceOf[List[Stmt]]).asInstanceOf[List[Stmt]]
			)
		,
			ctx.exp.accept(this).asInstanceOf[Expr]
		)
// for_stmt : FOR LB exp SEMI exp SEMI exp RB stmt ;
	override def visitFor_stmt(ctx: For_stmtContext) =
		For(
			ctx.exp(0).accept(this).asInstanceOf[Expr],
			ctx.exp(1).accept(this).asInstanceOf[Expr],
			ctx.exp(2).accept(this).asInstanceOf[Expr],
			ctx.stmt.accept(this).asInstanceOf[Stmt]
		)
// break_stmt : BREAK SEMI ;
	override def visitBreak_stmt (ctx: Break_stmtContext) =
		Break

// continue_stmt : CONTINUE SEMI ;
	override def visitContinue_stmt(ctx: Continue_stmtContext) =
		Continue

// return_stmt : RETURN SEMI | RETURN exp SEMI ;
	override def visitReturn_stmt(ctx: Return_stmtContext) =
		if(ctx.getChildCount() == 2) Return(None)
		else Return(Some(ctx.exp.accept(this).asInstanceOf[Expr]))

// exp_stmt : exp SEMI;
	override def visitExp_stmt (ctx: Exp_stmtContext) =
		ctx.exp.accept(this)

// exp : term_or (ASSIGN term_or)* ;
// exp : term_or ASSIGN exp | term_or ;
	override def visitExp(ctx: ExpContext) =
		// ctx.term_or.asScala.toList.tail.zip(ctx.ASSIGN.asScala.toList).foldRight(ctx.term_or(0).accept(this).asInstanceOf[Expr])(
  //  			 (x,y) => BinaryOp(x._2.getText,y,x._1.accept(this).asInstanceOf[Expr])
  //  			 )
  		if(ctx.getChildCount()==1) ctx.term_or.accept(this)
  		else BinaryOp(ctx.ASSIGN.getText,ctx.term_or.accept(this).asInstanceOf[Expr],ctx.exp.accept(this).asInstanceOf[Expr])

// term_or : (term_and OR)* term_and;
	override def visitTerm_or(ctx: Term_orContext) =
		ctx.term_and.asScala.toList.tail.zip(ctx.OR.asScala.toList).foldLeft(ctx.term_and(0).accept(this).asInstanceOf[Expr])(
   			 (x,y) => BinaryOp(y._2.getText,x,y._1.accept(this).asInstanceOf[Expr])
   			 )
// term_and: (term_equal AND)* term_equal  ;
	override def visitTerm_and (ctx: Term_andContext) =
		ctx.term_equal.asScala.toList.tail.zip(ctx.AND.asScala.toList).foldLeft(ctx.term_equal(0).accept(this).asInstanceOf[Expr])(
   			 (x,y) => BinaryOp(y._2.getText,x,y._1.accept(this).asInstanceOf[Expr])
   			 )
// term_equal : term_less EQUAL term_less | term_less NOTEQUAL term_less | term_less ;
	override def visitTerm_equal (ctx: Term_equalContext) =
		if (ctx.getChildCount() == 1) ctx.term_less(0).accept(this)
		else {
			BinaryOp(ctx.getChild(1).getText,ctx.term_less(0).accept(this).asInstanceOf[Expr],ctx.term_less(1).accept(this).asInstanceOf[Expr])
		}
// term_less : term_add LT term_add | term_add GT term_add | term_add LEQ term_add | term_add GEQ term_add | term_add ;
	override def visitTerm_less (ctx: Term_lessContext) =
		if (ctx.getChildCount() == 1) ctx.term_add(0).accept(this)
		else {
			BinaryOp(ctx.getChild(1).getText,ctx.term_add(0).accept(this).asInstanceOf[Expr],ctx.term_add(1).accept(this).asInstanceOf[Expr])
		}
// term_add : (term_div (ADD|SUB))* term_div;
// term_add : term_add ADD term_div | term_add SUB term_div | term_div ;
	override def visitTerm_add (ctx: Term_addContext) =
		if(ctx.getChildCount() == 1 ) ctx.term_div.accept(this)
		else {
			BinaryOp(ctx.getChild(1).getText,ctx.term_add.accept(this).asInstanceOf[Expr],ctx.term_div.accept(this).asInstanceOf[Expr])
		}

// term_div : (term_not (DIV|MUL|MOD))* term_not ;
// term_div : term_div DIV term_not | term_div MUL term_not | term_div MOD term_not | term_not
	override def visitTerm_div (ctx: Term_divContext) =
		if(ctx.getChildCount() == 1 ) ctx.term_not.accept(this)
		else {
			BinaryOp(ctx.getChild(1).getText,ctx.term_div.accept(this).asInstanceOf[Expr],ctx.term_not.accept(this).asInstanceOf[Expr])
		}
// term_not : (SUB|NOT)* term_square ;
// term_not : SUB term_not | NOT term_not | term_square ;
	override def visitTerm_not (ctx: Term_notContext) =
		if(ctx.getChildCount() == 1 ) ctx.term_square.accept(this)
		else {
			UnaryOp(ctx.getChild(0).getText,ctx.term_not.accept(this).asInstanceOf[Expr])
		}
// term_square : factor LSB exp RSB | factor ;
	override def visitTerm_square (ctx: Term_squareContext) =
		if(ctx.getChildCount() == 1) ctx.factor.accept(this)
		else ArrayCell(ctx.factor.accept(this).asInstanceOf[Expr],
			ctx.exp.accept(this).asInstanceOf[Expr])

// factor : LB exp RB | INTLIT | FLOATLIT | STRINGLIT | BOOLEANLIT | ID | funcall ;
	override def visitFactor( ctx : FactorContext) =
		if(ctx.funcall != null) ctx.funcall.accept(this)
		else if (ctx.getChildCount() == 3) ctx.exp.accept(this)
		else if (ctx.INTLIT != null) IntLiteral(ctx.INTLIT.getText.toInt)
		else if (ctx.FLOATLIT != null) FloatLiteral(ctx.FLOATLIT.getText.toFloat)
		else if (ctx.STRINGLIT != null) StringLiteral(ctx.STRINGLIT.getText)
		else if (ctx.BOOLEANLIT != null) BooleanLiteral(ctx.BOOLEANLIT.getText.toBoolean)
		else if (ctx.ID != null) Id(ctx.ID.getText)
// funcall : ID LB (exp (COM exp)*)? RB ;
	override def visitFuncall (ctx: FuncallContext) =
		CallExpr(Id(ctx.ID.getText),ctx.exp.asScala.toList.map(_.accept(this).asInstanceOf[Expr]))



}

==========================================================MC.G4==============================================
/**
 * Student name: Le Huu Viet Anh
 * Student ID: 1510054
 */
grammar MC;

@lexer::header{
	package mc.parser;
}

@lexer::members{
@Override
public Token emit() {
    switch (getType()) {
    case UNCLOSE_STRING:
        Token result = super.emit();
        // you'll need to define this method
        throw new UncloseString(result.getText().substring(1, getText().length()));

    case ILLEGAL_ESCAPE:
    	result = super.emit();
    	throw new IllegalEscape(result.getText().substring(1, getText().length()));

    case ERROR_CHAR:
    	result = super.emit();
    	throw new ErrorToken(result.getText());

    default:
        return super.emit();
    }
}
}

@parser::header{
	package mc.parser;
}

options{
	language=Java;
}

program : declaration+ EOF ;

declaration: var_declr | fun_declr;

var_declr: pri_type variable (COM variable)* SEMI;

pri_type: INTTYPE | STRINGTYPE | BOOLEANTYPE | FLOATTYPE ;

variable : ID | ID LSB INTLIT RSB ;

fun_declr : func_type ID LB (para_declr (COM para_declr)* )? RB block_stmt ;

func_type : pri_type | arr_pointer_type | VOIDTYPE ;

para_declr: (pri_type ID) | (pri_type ID LSB RSB) ;

arr_pointer_type : pri_type LSB RSB;

block_stmt : LP var_declr* stmt* RP ;

stmt : if_stmt | do_stmt | for_stmt | break_stmt | continue_stmt | return_stmt | exp_stmt | block_stmt;

if_stmt : IF LB exp RB stmt ELSE stmt | IF LB exp RB stmt;

do_stmt : DO stmt+ WHILE  exp SEMI ;

for_stmt : FOR LB exp SEMI exp SEMI exp RB stmt ;

break_stmt : BREAK SEMI ;

continue_stmt : CONTINUE SEMI ;

return_stmt : RETURN SEMI | RETURN exp SEMI ;

exp_stmt : exp SEMI;

//exp : term_or (ASSIGN term_or)* ;
exp : term_or ASSIGN exp | term_or ;

term_or : (term_and OR)* term_and;

term_and: (term_equal AND)* term_equal  ;

term_equal : term_less EQUAL term_less | term_less NOTEQUAL term_less | term_less ;

term_less : term_add LT term_add | term_add GT term_add | term_add LEQ term_add | term_add GEQ term_add | term_add ;

//term_add : (term_div (ADD|SUB))* term_div;

//term_div : (term_not (DIV|MUL|MOD))* term_not ;

//term_not : SUB term_not | NOT term_not | term_square ;

term_add : term_add ADD term_div |term_add SUB term_div | term_div ;

term_div : term_div DIV term_not | term_div MUL term_not | term_div MOD term_not | term_not ;

term_not : SUB term_not | NOT term_not | term_square ;

term_square : factor LSB exp RSB | factor ;

factor : LB exp RB | INTLIT | FLOATLIT | STRINGLIT | BOOLEANLIT | ID | funcall ;

funcall : ID LB (exp (COM exp)*)? RB ;

INTTYPE: 'int' ;

VOIDTYPE: 'void' ;

FLOATTYPE: 'float' ;

STRINGTYPE: 'string' ;

BOOLEANTYPE: 'boolean' ;

LINECOMMENT: '//'~[\n]* -> skip;

BLOCKCOMMENT:'/*' .*? '*/' -> skip;

IF : 'if' ;

ELSE : 'else' ;

DO : 'do' ;

WHILE : 'while' ;

FOR : 'for' ;

BREAK : 'break' ;

CONTINUE : 'continue' ;

RETURN : 'return' ;

INTLIT: [0-9]+;

fragment EXPONENT: [eE]'-'?[0-9]+ ;

FLOATLIT: ([0-9]+(EXPONENT))|(([0-9]+'.'[0-9]*|'.'[0-9]+)(EXPONENT)?) ;

fragment ESCAPE: '\\'[bnfrt\\\'"] ;
STRINGLIT:'"' (~["\\\n]|ESCAPE)* '"'{setText(getText().substring(1, getText().length()-1));};

BOOLEANLIT: 'true'|'false';

KEYWORD: ('boolean'|'break'|'continue'|'else'|'for'|'float'|'if'|'int'|'return'|'void'|'do'|'while'|'true'|'false'|'string');

ID: [_a-zA-Z][_a-zA-Z0-9]* ;

ADD: '+';

SUB: '-';

MUL: '*';

DIV: '/';

MOD: '%';

NOT: '!';

OR: '||';

AND: '&&';

ASSIGN: '=';

NOTEQUAL: '!=';

EQUAL: '==';

LT: '<';

GT: '>';

LEQ: '<=';

GEQ: '>=';

LB: '(' ;

RB: ')' ;

LP: '{';

RP: '}';

SEMI: ';' ;

LSB: '[';

RSB: ']';

COM: ',';

WS : [ \f\t\r\n]+ -> skip ; // skip spaces, tabs, newlines

ERROR_CHAR: .;

UNCLOSE_STRING: '"' (~["\\\n]|ESCAPE)* EOF? ;

fragment IES: '\\'(~[bnfrt\\\'"]) ; // Illegal escape
ILLEGAL_ESCAPE: '"' (~["\\\n]|ESCAPE)*IES ;