package com.jetbrains.python.psi.impl;

import com.intellij.openapi.util.io.FileUtil;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.psi.PsiElement;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.resolve.PyResolveContext;
import org.jetbrains.annotations.Nullable;

import java.io.File;

/**
 * @author yole
 */
public class PyPathEvaluator {

  private PyPathEvaluator() {
  }

  @Nullable
  public static String evaluate(PyExpression expr) {
    if (expr == null) {
      return null;
    }
    VirtualFile vFile = expr.getContainingFile().getVirtualFile();
    return evaluate(expr, vFile == null ? null : vFile.getPath());
  }

  @Nullable
  public static String evaluate(PyExpression expr, String path) {
    if (expr == null) {
      return null;
    }
    if (expr instanceof PyCallExpression) {
      final PyCallExpression call = (PyCallExpression)expr;
      final PyExpression[] args = call.getArguments();
      if (call.isCalleeText(PyNames.DIRNAME) && args.length == 1) {
        String argValue = evaluate(args[0], path);
        return argValue == null ? null : new File(argValue).getParent();
      }
      else if (call.isCalleeText(PyNames.JOIN) && args.length >= 1) {
        String result = null;
        for (int i = 0; i<args.length; i++) {
          String arg = evaluate(args[i], path);
          if (arg == null) {
            return null;
          }
          if (result == null) {
            result = arg;
          } else {
            result = new File(result, arg).getPath();
          }
        }
        return result;
      }
      else if (call.isCalleeText(PyNames.ABSPATH) && args.length == 1) {
        String argValue = evaluate(args[0], path);
        // relative to directory of 'path', not file
        if (argValue == null) {
          return null;
        }
        if (FileUtil.isAbsolute(argValue)) {
          return argValue;
        }
        else {
          return new File(new File(path).getParent(), argValue).getPath();
        }
      }
      else if (call.isCalleeText(PyNames.REPLACE) && args.length == 2) {
        final PyExpression callee = call.getCallee();
        if (!(callee instanceof PyQualifiedExpression)) return null;
        final PyExpression qualifier = ((PyQualifiedExpression)callee).getQualifier();
        String result = evaluate(qualifier, path);
        if (result == null) return null;
        String arg1 = evaluate(args[0], path);
        String arg2 = evaluate(args[1], path);
        if (arg1 == null || arg2 == null) return null;
        return result.replace(arg1, arg2);
      }
    }
    else if (expr instanceof PyReferenceExpression) {
      if (((PyReferenceExpression)expr).getQualifier() == null) {
        final String refName = ((PyReferenceExpression)expr).getReferencedName();
        if (PyNames.FILE.equals(refName)) {
          return path;
        }
        PsiElement result = ((PyReferenceExpression)expr).getReference(PyResolveContext.noImplicits()).resolve();
        if (result instanceof PyTargetExpression) {
          result = ((PyTargetExpression)result).findAssignedValue();
        }
        if (result instanceof PyExpression) {
          return evaluate((PyExpression)result, path);
        }
      }
    }
    else if (expr instanceof PyStringLiteralExpression) {
      return ((PyStringLiteralExpression)expr).getStringValue();
    }
    return null;
  }
}

@@ -770,17 +770,31 @@
  String 
-p
+containingFileP
 ath) %7B%0A 
@@ -1104,33 +1104,47 @@
 aluate(args%5B0%5D, 
-p
+containingFileP
 ath);%0A        re
@@ -1296,364 +1296,72 @@
     
-String result = null;%0A        for (int i = 0; i%3Cargs.length; i++) %7B%0A          String arg = evaluate(args%5Bi%5D, path);%0A          if (arg == null) %7B%0A            return null;%0A          %7D%0A          if (result == null) %7B%0A            result = arg;%0A          %7D else %7B%0A            result = new File(result, arg).getPath();%0A          %7D%0A        %7D%0A        return result
+return evaluatePathInJoin(containingFilePath, args, args.length)
 ;%0A  
@@ -1475,33 +1475,47 @@
 aluate(args%5B0%5D, 
-p
+containingFileP
 ath);%0A        //
@@ -1541,17 +1541,31 @@
 ory of '
-p
+containingFileP
 ath', no
@@ -1768,17 +1768,31 @@
 ew File(
-p
+containingFileP
 ath).get
@@ -2167,17 +2167,31 @@
 lifier, 
-p
+containingFileP
 ath);%0A  
@@ -2269,17 +2269,31 @@
 rgs%5B0%5D, 
-p
+containingFileP
 ath);%0A  
@@ -2330,17 +2330,31 @@
 rgs%5B1%5D, 
-p
+containingFileP
 ath);%0A  
@@ -2726,17 +2726,31 @@
  return 
-p
+containingFileP
 ath;%0A   
@@ -3095,17 +3095,31 @@
 result, 
-p
+containingFileP
 ath);%0A  
@@ -3289,11 +3289,452 @@
 ull;%0A  %7D
+%0A%0A  public static String evaluatePathInJoin(String containingFilePath, PyExpression%5B%5D args, int endElement) %7B%0A    String result = null;%0A    for (int i = 0; i%3C endElement; i++) %7B%0A      String arg = evaluate(args%5Bi%5D, containingFilePath);%0A      if (arg == null) %7B%0A        return null;%0A      %7D%0A      if (result == null) %7B%0A        result = arg;%0A      %7D else %7B%0A        result = new File(result, arg).getPath();%0A      %7D%0A    %7D%0A    return result;%0A  %7D
 %0A%7D%0A

Add path completion to os.path.join elements.