package io.ray.test;

import com.google.common.collect.ImmutableList;
import io.ray.runtime.runner.RunManager;
import java.io.File;
import java.time.LocalDateTime;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.SystemUtils;
import org.testng.IInvokedMethod;
import org.testng.IInvokedMethodListener;
import org.testng.ITestContext;
import org.testng.ITestListener;
import org.testng.ITestResult;
import org.testng.SkipException;

public class TestProgressListener implements IInvokedMethodListener, ITestListener {

  // Travis aborts CI if no outputs for 10 minutes. So threshold needs to be smaller than 10m.
  private static final long hangDetectionThresholdMillis = 5 * 60 * 1000;
  private static final int TAIL_NO_OF_LINES = 500;
  private Thread testMainThread;
  private long testStartTimeMillis;

  private String getFullTestName(ITestResult testResult) {
    return testResult.getTestClass().getName() + "." + testResult.getMethod().getMethodName();
  }

  private void printSection(String sectionName) {
    System.out.println(
        "============ [" + LocalDateTime.now().toString() + "] " + sectionName + " ============");
  }

  private void printTestStage(String tag, String content) {
    printSection("[" + tag + "] " + content);
  }

  @Override
  public void beforeInvocation(IInvokedMethod method, ITestResult testResult) {}

  @Override
  public void afterInvocation(IInvokedMethod method, ITestResult testResult) {}

  @Override
  public void onTestStart(ITestResult result) {
    printTestStage("TEST START", getFullTestName(result));
    testStartTimeMillis = System.currentTimeMillis();
    // TODO(kfstorm): Add a timer to detect hang
    if (testMainThread == null) {
      testMainThread = Thread.currentThread();
      Thread hangDetectionThread =
          new Thread(
              () -> {
                try {
                  // If current task case has ran for more than 5 minutes.
                  while (System.currentTimeMillis() - testStartTimeMillis
                      < hangDetectionThresholdMillis) {
                    Thread.sleep(1000);
                  }
                  printDebugInfo(null, /*testHanged=*/ true);
                } catch (InterruptedException e) {
                  // ignored
                }
              });
      hangDetectionThread.setDaemon(true);
      hangDetectionThread.start();
    }
  }

  @Override
  public void onTestSuccess(ITestResult result) {
    printTestStage("TEST SUCCESS", getFullTestName(result));
  }

  @Override
  public void onTestFailure(ITestResult result) {
    printTestStage("TEST FAILURE", getFullTestName(result));
    printDebugInfo(result, /*testHanged=*/ false);
  }

  @Override
  public void onTestSkipped(ITestResult result) {
    printTestStage("TEST SKIPPED", getFullTestName(result));
    printDebugInfo(result, /*testHanged=*/ false);
  }

  @Override
  public void onTestFailedButWithinSuccessPercentage(ITestResult result) {
    printTestStage("TEST FAILED BUT WITHIN SUCCESS PERCENTAGE", getFullTestName(result));
  }

  @Override
  public void onStart(ITestContext context) {}

  @Override
  public void onFinish(ITestContext context) {}

  private void printDebugInfo(ITestResult result, boolean testHanged) {
    boolean testFailed = false;
    if (result != null) {
      Throwable throwable = result.getThrowable();
      if (throwable != null && !(throwable instanceof SkipException)) {
        testFailed = true;
        throwable.printStackTrace();
      }
    }
    if (!testFailed && !testHanged) {
      return;
    }

    if (testHanged) {
      printSection("TEST CASE HANGED");
      printSection("STACK TRACE OF TEST THREAD");
      for (StackTraceElement element : testMainThread.getStackTrace()) {
        System.out.println(element.toString());
      }
      Set<Integer> javaPids = getJavaPids();
      for (Integer pid : javaPids) {
        runCommandSafely(ImmutableList.of("jstack", pid.toString()));
        // TODO(kfstorm): Check lldb or gdb exists rather than detecting OS type.
        if (SystemUtils.IS_OS_MAC) {
          runCommandSafely(
              ImmutableList.of("lldb", "--batch", "-o", "bt all", "-p", pid.toString()));
        } else {
          runCommandSafely(
              ImmutableList.of(
                  "sudo", "gdb", "-batch", "-ex", "thread apply all bt", "-p", pid.toString()));
        }
      }
    }

    printLogFiles();

    if (testHanged) {
      printSection("ABORT TEST");
      System.exit(1);
    }
  }

  private String runCommandSafely(List<String> command) {
    String output;
    String commandString = String.join(" ", command);
    printSection(commandString);
    try {
      output = RunManager.runCommand(command);
      System.out.println(output);
    } catch (Exception e) {
      System.out.println("Failed to execute command: " + commandString);
      e.printStackTrace();
      output = "";
    }
    return output;
  }

  private Set<Integer> getJavaPids() {
    Set<Integer> javaPids = new HashSet<>();
    String jpsOutput = runCommandSafely(ImmutableList.of("jps", "-v"));
    try {
      for (String line : StringUtils.split(jpsOutput, "\n")) {
        String[] parts = StringUtils.split(line);
        if (parts.length > 1 && parts[1].toLowerCase().equals("jps")) {
          // Skip jps.
          continue;
        }
        Integer pid = Integer.valueOf(parts[0]);
        javaPids.add(pid);
      }
    } catch (Exception e) {
      System.out.println("Failed to parse jps output.");
      e.printStackTrace();
    }

    String pgrepJavaResult = runCommandSafely(ImmutableList.of("pgrep", "java"));
    try {
      for (String line : StringUtils.split(pgrepJavaResult, "\n")) {
        Integer pid = Integer.valueOf(line);
        javaPids.add(pid);
      }
    } catch (Exception e) {
      System.out.println("Failed to parse pgrep java output.");
      e.printStackTrace();
    }

    return javaPids;
  }

  private void printLogFiles() {
    Collection<File> logFiles =
        FileUtils.listFiles(new File("/tmp/ray/session_latest/logs"), null, false);
    for (File file : logFiles) {
      runCommandSafely(
          ImmutableList.of("tail", "-n", String.valueOf(TAIL_NO_OF_LINES), file.getAbsolutePath()));
    }
  }
}

@@ -3733,32 +3733,195 @@
  (testHanged) %7B%0A
+      if (result != null) %7B%0A        printSection(%22TEST CASE HANGED: %22 + getFullTestName(result));%0A      %7D else %7B%0A        printSection(%22TEST CASE HANGED%22);%0A      %7D%0A
       printSecti

print hang test case name (#21804)