package mclab.ide.callgraph;

import mclint.MatlabProgram;
import mclint.Project;
import mclint.util.AstUtil;
import natlab.LookupFile;
import natlab.toolkits.analysis.varorfun.VFAnalysis;
import natlab.toolkits.analysis.varorfun.VFPreorderAnalysis;
import natlab.toolkits.filehandling.FunctionOrScriptQuery;
import natlab.toolkits.path.FileEnvironment;
import nodecases.AbstractNodeCaseHandler;
import ast.ASTNode;
import ast.Expr;
import ast.ExprStmt;
import ast.Function;
import ast.FunctionHandleExpr;
import ast.Name;
import ast.NameExpr;
import ast.ParameterizedExpr;
import ast.Program;
import ast.Stmt;
import ast.StringLiteralExpr;

public class CallgraphTransformer extends AbstractNodeCaseHandler {
  private static FunctionOrScriptQuery BASELINE_QUERY = LookupFile.getFunctionOrScriptQueryObject();
  private static FunctionOrScriptQuery getKindAnalysisEnvironment(Program node) {
    final FunctionOrScriptQuery env = new FileEnvironment(node.getFile())
        .getFunctionOrScriptQuery(node.getFile());
    return new FunctionOrScriptQuery() {
      @Override public boolean isFunctionOrScript(String name) {
        return env.isFunctionOrScript(name) || BASELINE_QUERY.isFunctionOrScript(name);
      }

      @Override public boolean isPackage(String name) {
        return env.isPackage(name) || BASELINE_QUERY.isPackage(name);
      }
    };
  }

  public static void instrument(Project project) {
    project.getMatlabPrograms().forEach(CallgraphTransformer::instrument);
  }

  public static void instrument(MatlabProgram program) {
    Program node = program.parse();
    VFAnalysis kindAnalysis = new VFPreorderAnalysis(node, getKindAnalysisEnvironment(node));
    kindAnalysis.analyze();
    node.analyze(new CallgraphTransformer(kindAnalysis, program.getPath().toString()));
  }

  private VFAnalysis kindAnalysis;
  private String relativePath;

  private CallgraphTransformer(VFAnalysis kindAnalysis, String relativePath) {
    this.kindAnalysis = kindAnalysis;
    this.relativePath = relativePath;
  }

  private String identifier(ASTNode node, String id) {
    return String.format("%s@%s:%d,%d", id, relativePath, node.getStartLine(), node.getStartColumn());
  }

  private String identifier(Name name) {
    return identifier(name, name.getID());
  }

  private String identifier(Function function) {
    return identifier(function, function.getName());
  }

  private static Name getTarget(ParameterizedExpr call) {
    return ((NameExpr) call.getTarget()).getName();
  }

  private boolean isCall(ParameterizedExpr call) {
    if (!(call.getTarget() instanceof NameExpr)) {
      return false;
    }
    return kindAnalysis.getResult(getTarget(call)).isFunction();
  }

  private static <T extends ASTNode> ast.List<T> concat(ast.List<T> list, ast.List<T> other) {
    for (T t : other) {
      list = list.add(t);
    }
    return list;
  }

  private Expr wrapWithTraceCall(ParameterizedExpr call) {
    Name target = getTarget(call);
    return new ParameterizedExpr(
        new NameExpr(new Name("mclab_callgraph_log_then_run")),
        concat(
          new ast.List<Expr>()
            .add(new StringLiteralExpr("call " + identifier(target) + "\\n"))
            .add(new FunctionHandleExpr(new Name(target.getID()))),
          call.getArgs()
        )
    );

  }

  private Stmt functionEntryLogStmt(Function f) {
    return new ExprStmt(
        new ParameterizedExpr(
            new NameExpr(new Name("mclab_callgraph_log")),
            new ast.List<Expr>()
              .add(new StringLiteralExpr("enter " + identifier(f) + "\\n"))
        )
    );
  }

  @Override public void caseASTNode(ASTNode node) {
    for (int i = 0; i < node.getNumChild(); ++i) {
      node.getChild(i).analyze(this);
    }
  }

  @Override public void caseFunction(Function f) {
    caseASTNode(f);
    f.getStmts().insertChild(functionEntryLogStmt(f), 0);
  }

  @Override public void caseParameterizedExpr(ParameterizedExpr e) {
    if (!isCall(e)) {
      return;
    }
    e.getArgs().analyze(this);
    AstUtil.replace(e, wrapWithTraceCall(e));
  }
}

@@ -2399,29 +2399,22 @@
 ate 
-static Name getTarget
+boolean isCall
 (Par
@@ -2447,16 +2447,62 @@
  return 
+call.getTarget() instanceof NameExpr && isCall
 ((NameEx
@@ -2522,26 +2522,16 @@
 arget())
-.getName()
 ;%0A  %7D%0A%0A 
@@ -2550,37 +2550,28 @@
 lean isCall(
-Parameterized
+Name
 Expr call) %7B
@@ -2575,85 +2575,8 @@
 ) %7B%0A
-    if (!(call.getTarget() instanceof NameExpr)) %7B%0A      return false;%0A    %7D%0A
     
@@ -2605,30 +2605,29 @@
 tResult(
-getTarget(call
+call.getName(
 )).isFun
@@ -2743,59 +2743,38 @@
     
-for (T t : other) %7B%0A      list = list.add(t);%0A    %7D
+other.forEach(list::addChild);
 %0A   
@@ -2873,22 +2873,46 @@
 t = 
-getTarget(call
+((NameExpr) call.getTarget()).getName(
 );%0A 
@@ -3904,16 +3904,47 @@
 ) %7B%0A
+    e.getArgs().analyze(this);%0A
     if (
 !isC
@@ -3943,9 +3943,8 @@
 if (
-!
 isCa
@@ -3962,98 +3962,238 @@
     
-return;%0A    %7D%0A    e.getArgs().analyze(this);%0A    AstUtil.replace(e, wrapWithTraceCall(e));
+AstUtil.replace(e, wrapWithTraceCall(e));%0A    %7D%0A  %7D%0A%0A  @Override public void caseNameExpr(NameExpr e) %7B%0A    if (isCall(e)) %7B%0A      AstUtil.replace(e, wrapWithTraceCall(new ParameterizedExpr(e.fullCopy(), new ast.List%3C%3E())));%0A    %7D
 %0A  %7D

add support for calls with no parens.