package org.karaffe.compiler.visitor;

import org.karaffe.compiler.frontend.karaffe.antlr.KaraffeBaseVisitor;
import org.karaffe.compiler.frontend.karaffe.antlr.KaraffeParser;
import org.karaffe.compiler.tree.NodeType;
import org.karaffe.compiler.tree.Tree;
import org.karaffe.compiler.tree.attr.InferredType;
import org.karaffe.compiler.tree.attr.JavaModifier;
import org.karaffe.compiler.tree.attr.MethodSignature;
import org.karaffe.compiler.tree.attr.ModifierAttribute;
import org.karaffe.compiler.tree.attr.SuperClass;
import org.karaffe.compiler.util.CompilerContext;
import org.karaffe.compiler.util.MethodParameterEntry;
import org.karaffe.compiler.util.Position;

import java.util.Collections;
import java.util.List;

public class KaraffeASTCreateVisitor extends KaraffeBaseVisitor<Tree> {

    private final CompilerContext context;
    private final Tree compilationUnit;

    public KaraffeASTCreateVisitor(CompilerContext context) {
        this.context = context;
        this.compilationUnit = new Tree(NodeType.CompilationUnit, "", null);
    }

    public Tree getCompilationUnit() {
        return compilationUnit;
    }

    @Override
    public Tree visitSourceFile(KaraffeParser.SourceFileContext ctx) {
        Tree tree = new Tree(NodeType.SourceFile, ctx.EOF().getSymbol().getInputStream().getSourceName(), new Position(ctx));
        for (KaraffeParser.ClassDefContext classDefContext : ctx.classDef()) {
            tree.addChild(classDefContext.accept(this));
        }
        this.compilationUnit.addChild(tree);
        return tree;
    }

    @Override
    public Tree visitClassDef(KaraffeParser.ClassDefContext ctx) {
        Tree tree = new Tree(NodeType.DefClass, ctx.Identifier().getText(), new Position(ctx));
        tree.addAttribute(new SuperClass() /*java.lang.Object*/);
        tree.addAttribute(new ModifierAttribute(JavaModifier.PUBLIC));
        if (ctx.typeDefBody() != null) {
            List<KaraffeParser.StatementContext> statements = ctx.typeDefBody().statement();
            for (KaraffeParser.StatementContext statement : statements) {
                tree.addChild(statement.accept(this));
            }
        }
        super.visitClassDef(ctx);
        return tree;
    }

    @Override
    public Tree visitEntryPointBlock(KaraffeParser.EntryPointBlockContext ctx) {
        Tree tree = new Tree(NodeType.DefMethod, "main", new Position(ctx));
        ModifierAttribute modifierAttribute = new ModifierAttribute(JavaModifier.PUBLIC, JavaModifier.STATIC);
        tree.addAttribute(modifierAttribute);
        tree.addAttribute(new MethodSignature(void.class, Collections.singletonList(new MethodParameterEntry("args", String[].class))));
        List<KaraffeParser.StatementContext> statements = ctx.statement();
        for (KaraffeParser.StatementContext statement : statements) {
            tree.addChild(statement.accept(this));
        }
        return tree;
    }

    @Override
    public Tree visitVarDef(KaraffeParser.VarDefContext ctx) {
        Tree tree = new Tree(NodeType.DefVar, ctx.Identifier().getText(), new Position(ctx));
        return tree;
    }

    @Override
    public Tree visitPrintFunction(KaraffeParser.PrintFunctionContext ctx) {
        Tree tree = new Tree(NodeType.Apply, "()", new Position(ctx));
        tree.addChild(new Tree(NodeType.Select, "print", new Position(ctx.PRINT().getSymbol())));
        if (ctx.body != null) {
            tree.addChild(ctx.body.accept(this));
        }
        return tree;
    }

    public Tree visitLiteral(KaraffeParser.LiteralContext ctx) {
        Tree tree;
        if (ctx.IntegerLiteral() != null) {
            tree = new Tree(NodeType.IntLiteral, ctx.IntegerLiteral().getText(), new Position(ctx));
            tree.addAttribute(new InferredType(int.class));
        } else if (ctx.StringLiteral() != null) {
            tree = new Tree(NodeType.StringLiteral, ctx.StringLiteral().getText(), new Position(ctx));
            tree.addAttribute(new InferredType(String.class));
        } else {
            throw new IllegalArgumentException();
        }
        return tree;
    }
}

@@ -255,60 +255,8 @@
 ee;%0A
-import org.karaffe.compiler.tree.attr.InferredType;%0A
 impo
@@ -419,58 +419,8 @@
 te;%0A
-import org.karaffe.compiler.tree.attr.SuperClass;%0A
 impo
@@ -1630,32 +1630,34 @@
 Position(ctx));%0A
+//
         tree.add
@@ -1706,16 +1706,18 @@
 ect*/);%0A
+//
         
@@ -2120,32 +2120,773 @@
 rn tree;%0A    %7D%0A%0A
+    @Override%0A    public Tree visitExpr(KaraffeParser.ExprContext ctx) %7B%0A        if (ctx.lit != null) %7B%0A            //literal%0A            return ctx.literal().accept(this);%0A        %7D else if (ctx.t != null) %7B%0A            return new Tree(NodeType.This, %22this%22, new Position(ctx.t));%0A        %7D else if (ctx.op != null) %7B%0A            Tree tree = new Tree(NodeType.Apply, %22()%22, new Position(ctx));%0A            Tree operator = new Tree(NodeType.Select, ctx.op.getText(), new Position(ctx.op));%0A            tree.addChild(operator);%0A            operator.addChild(ctx.left.accept(this));%0A            operator.addChild(ctx.right.accept(this));%0A            return tree;%0A        %7D else %7B%0A            throw new IllegalStateException();%0A        %7D%0A    %7D%0A%0A
     @Override%0A  
@@ -3938,30 +3938,30 @@
         
-tree.addChild(
+Tree select = 
 new Tree
@@ -4025,16 +4025,46 @@
 mbol()))
+;%0A        tree.addChild(select
 );%0A     
@@ -4094,36 +4094,38 @@
 ) %7B%0A            
-tree
+select
 .addChild(ctx.bo
@@ -4401,32 +4401,34 @@
 Position(ctx));%0A
+//
             tree
@@ -4616,32 +4616,34 @@
 Position(ctx));%0A
+//
             tree

Disable set attributes