package io.ray.test;

import static io.ray.api.runtimeenv.types.RuntimeEnvName.JARS;

import com.google.common.collect.ImmutableList;
import io.ray.api.ActorHandle;
import io.ray.api.Ray;
import io.ray.api.runtimeenv.RuntimeEnv;
import io.ray.api.runtimeenv.RuntimeEnvConfig;
import io.ray.api.runtimeenv.types.RuntimeEnvName;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.testng.Assert;
import org.testng.annotations.Test;

@Test(groups = "cluster")
public class RuntimeEnvTest {

  private static final String FOO_JAR_URL =
      "https://github.com/ray-project/test_packages/raw/main/raw_resources/foo.jar";
  private static final String BAR_JAR_URL =
      "https://github.com/ray-project/test_packages/raw/main/raw_resources/bar.jar";
  private static final String FOO_ZIP_URL =
      "https://github.com/ray-project/test_packages/raw/main/raw_resources/foo.zip";
  private static final String BAR_ZIP_URL =
      "https://github.com/ray-project/test_packages/raw/main/raw_resources/bar.zip";

  private static final String FOO_CLASS_NAME = "io.testpackages.Foo";
  private static final String BAR_CLASS_NAME = "io.testpackages.Bar";

  private static class A {

    public String getEnv(String key) {
      return System.getenv(key);
    }

    public boolean findClass(String className) {
      try {
        Class.forName(className);
      } catch (ClassNotFoundException e) {
        return false;
      }
      return true;
    }
  }

  public void testPerJobEnvVars() {
    System.setProperty("ray.job.runtime-env.env-vars.KEY1", "A");
    System.setProperty("ray.job.runtime-env.env-vars.KEY2", "B");

    try {
      Ray.init();
      ActorHandle<A> actor = Ray.actor(A::new).remote();
      String val = actor.task(A::getEnv, "KEY1").remote().get();
      Assert.assertEquals(val, "A");
      val = actor.task(A::getEnv, "KEY2").remote().get();
      Assert.assertEquals(val, "B");
    } finally {
      System.clearProperty("ray.job.runtime-env.env-vars.KEY1");
      System.clearProperty("ray.job.runtime-env.env-vars.KEY2");
      Ray.shutdown();
    }
  }

  private static String getEnvVar(String key) {
    return System.getenv(key);
  }

  public void testEnvVarsForNormalTask() {
    try {
      Ray.init();
      RuntimeEnv runtimeEnv = new RuntimeEnv.Builder().build();
      Map<String, String> envMap =
          new HashMap<String, String>() {
            {
              put("KEY1", "A");
              put("KEY2", "B");
              put("KEY1", "C");
            }
          };
      runtimeEnv.set(RuntimeEnvName.ENV_VARS, envMap);

      String val =
          Ray.task(RuntimeEnvTest::getEnvVar, "KEY1").setRuntimeEnv(runtimeEnv).remote().get();
      Assert.assertEquals(val, "C");
      val = Ray.task(RuntimeEnvTest::getEnvVar, "KEY2").setRuntimeEnv(runtimeEnv).remote().get();
      Assert.assertEquals(val, "B");
    } finally {
      Ray.shutdown();
    }
  }

  /// overwrite the runtime env from job config.
  public void testPerTaskEnvVarsOverwritePerJobEnvVars() {
    System.setProperty("ray.job.runtime-env.env-vars.KEY1", "A");
    System.setProperty("ray.job.runtime-env.env-vars.KEY2", "B");
    try {
      Ray.init();
      Map<String, String> envMap =
          new HashMap<String, String>() {
            {
              put("KEY1", "C");
            }
          };
      RuntimeEnv runtimeEnv = new RuntimeEnv.Builder().build();
      runtimeEnv.set(RuntimeEnvName.ENV_VARS, envMap);

      /// value of KEY1 is overwritten to `C` and KEY2s is extended from job config.
      String val =
          Ray.task(RuntimeEnvTest::getEnvVar, "KEY1").setRuntimeEnv(runtimeEnv).remote().get();
      Assert.assertEquals(val, "C");
      val = Ray.task(RuntimeEnvTest::getEnvVar, "KEY2").setRuntimeEnv(runtimeEnv).remote().get();
      Assert.assertEquals(val, "B");
    } finally {
      System.clearProperty("ray.job.runtime-env.env-vars.KEY1");
      System.clearProperty("ray.job.runtime-env.env-vars.KEY2");
      Ray.shutdown();
    }
  }

  private static void testDownloadAndLoadPackage(String url) {
    try {
      Ray.init();
      RuntimeEnv runtimeEnv = new RuntimeEnv.Builder().build();
      runtimeEnv.set(JARS, ImmutableList.of(url));
      ActorHandle<A> actor1 = Ray.actor(A::new).setRuntimeEnv(runtimeEnv).remote();
      boolean ret = actor1.task(A::findClass, FOO_CLASS_NAME).remote().get();
      Assert.assertTrue(ret);
    } finally {
      Ray.shutdown();
    }
  }

  public void testJarPackageInActor() {
    testDownloadAndLoadPackage(FOO_JAR_URL);
  }

  public void testZipPackageInActor() {
    testDownloadAndLoadPackage(FOO_ZIP_URL);
  }

  private static boolean findClasses(List<String> classNames) {
    try {
      for (String name : classNames) {
        Class.forName(name);
      }
    } catch (ClassNotFoundException e) {
      return false;
    }
    return true;
  }

  private static void testDownloadAndLoadPackagesForTask(
      List<String> urls, List<String> classNames) {
    try {
      Ray.init();
      RuntimeEnv runtimeEnv = new RuntimeEnv.Builder().build();
      runtimeEnv.set(JARS, urls);
      boolean ret =
          Ray.task(RuntimeEnvTest::findClasses, classNames)
              .setRuntimeEnv(runtimeEnv)
              .remote()
              .get();
      Assert.assertTrue(ret);
    } finally {
      Ray.shutdown();
    }
  }

  private static void testDownloadAndLoadPackagesForTask(String url, String className) {
    testDownloadAndLoadPackagesForTask(ImmutableList.of(url), ImmutableList.of(className));
  }

  public void testJarPackageForTask() {
    testDownloadAndLoadPackagesForTask(BAR_JAR_URL, BAR_CLASS_NAME);
  }

  public void testZipPackageForTask() {
    testDownloadAndLoadPackagesForTask(FOO_ZIP_URL, FOO_CLASS_NAME);
  }

  /// This case tests that a task needs 2 jars for load different classes.
  public void testMultipleJars() {
    testDownloadAndLoadPackagesForTask(
        ImmutableList.of(FOO_JAR_URL, BAR_JAR_URL),
        ImmutableList.of(BAR_CLASS_NAME, FOO_CLASS_NAME));
  }

  public void testRuntimeEnvJarsForJob() {
    System.setProperty("ray.job.runtime-env.jars.0", FOO_JAR_URL);
    System.setProperty("ray.job.runtime-env.jars.1", BAR_JAR_URL);
    try {
      Ray.init();
      boolean ret =
          Ray.task(RuntimeEnvTest::findClasses, ImmutableList.of(BAR_CLASS_NAME, FOO_CLASS_NAME))
              .remote()
              .get();
      Assert.assertTrue(ret);
    } finally {
      System.clearProperty("ray.job.runtime-env.jars.0");
      System.clearProperty("ray.job.runtime-env.jars.1");
      Ray.shutdown();
    }
  }

  private static class Pip {
    private String[] packages;
    private Boolean pipCheck;

    public String[] getPackages() {
      return packages;
    }

    public void setPackages(String[] packages) {
      this.packages = packages;
    }

    public Boolean getPipCheck() {
      return pipCheck;
    }

    public void setPipCheck(Boolean pipCheck) {
      this.pipCheck = pipCheck;
    }
  }

  public void testRuntimeEnvAPI() {
    try {
      Ray.init();
      RuntimeEnv runtimeEnv = new RuntimeEnv.Builder().build();
      String workingDir = "https://path/to/working_dir.zip";
      runtimeEnv.set("working_dir", workingDir);
      String[] pyModules =
          new String[] {"https://path/to/py_modules1.zip", "https://path/to/py_modules2.zip"};
      runtimeEnv.set("py_modules", pyModules);
      Pip pip = new Pip();
      pip.setPackages(new String[] {"requests", "tensorflow"});
      pip.setPipCheck(true);
      runtimeEnv.set("pip", pip);
      String serializedRuntimeEnv = runtimeEnv.serialize();

      RuntimeEnv runtimeEnv2 = RuntimeEnv.deserialize(serializedRuntimeEnv);
      Assert.assertEquals(runtimeEnv2.get("working_dir", String.class), workingDir);
      Assert.assertEquals(runtimeEnv2.get("py_modules", String[].class), pyModules);
      Pip pip2 = runtimeEnv2.get("pip", Pip.class);
      Assert.assertEquals(pip2.getPackages(), pip.getPackages());
      Assert.assertEquals(pip2.getPipCheck(), pip.getPipCheck());

      Assert.assertEquals(runtimeEnv2.remove("working_dir"), true);
      Assert.assertEquals(runtimeEnv2.remove("py_modules"), true);
      Assert.assertEquals(runtimeEnv2.remove("pip"), true);
      Assert.assertEquals(runtimeEnv2.remove("conda"), false);
      Assert.assertEquals(runtimeEnv2.get("working_dir", String.class), null);
      Assert.assertEquals(runtimeEnv2.get("py_modules", String[].class), null);
      Assert.assertEquals(runtimeEnv2.get("pip", Pip.class), null);
    } finally {
      Ray.shutdown();
    }
  }

  public void testRuntimeEnvJsonStringAPI() {
    try {
      Ray.init();
      RuntimeEnv runtimeEnv = new RuntimeEnv.Builder().build();
      String pipString = "{\"packages\":[\"requests\",\"tensorflow\"],\"pip_check\":false}";
      runtimeEnv.setJsonStr("pip", pipString);
      String serializedRuntimeEnv = runtimeEnv.serialize();

      RuntimeEnv runtimeEnv2 = RuntimeEnv.deserialize(serializedRuntimeEnv);
      Assert.assertEquals(runtimeEnv2.getJsonStr("pip"), pipString);
    } finally {
      Ray.shutdown();
    }
  }

  public void testRuntimeEnvContextForJob() {
    System.setProperty("ray.job.runtime-env.jars.0", FOO_JAR_URL);
    System.setProperty("ray.job.runtime-env.jars.1", BAR_JAR_URL);
    System.setProperty("ray.job.runtime-env.config.setup-timeout-seconds", "1");
    try {
      Ray.init();
      RuntimeEnv runtimeEnv = Ray.getRuntimeContext().getCurrentRuntimeEnv();
      Assert.assertNotNull(runtimeEnv);
      List<String> jars = runtimeEnv.get(JARS, List.class);
      Assert.assertNotNull(jars);
      Assert.assertEquals(jars.size(), 2);
      Assert.assertEquals(jars.get(0), FOO_JAR_URL);
      Assert.assertEquals(jars.get(1), BAR_JAR_URL);
      RuntimeEnvConfig runtimeEnvConfig = runtimeEnv.getConfig();
      Assert.assertNotNull(runtimeEnvConfig);
      Assert.assertEquals((int) runtimeEnvConfig.getSetupTimeoutSeconds(), 1);

    } finally {
      System.clearProperty("ray.job.runtime-env.jars.0");
      System.clearProperty("ray.job.runtime-env.jars.1");
      System.clearProperty("ray.job.runtime-env.config.setup-timeout-seconds");
      Ray.shutdown();
    }
  }

  private static Integer getRuntimeEnvTimeout() {
    RuntimeEnv runtimeEnv = Ray.getRuntimeContext().getCurrentRuntimeEnv();
    if (runtimeEnv != null) {
      return runtimeEnv.getConfig().getSetupTimeoutSeconds();
    }
    return null;
  }

  public void testRuntimeEnvContextForTask() {
    try {
      Ray.init();
      RuntimeEnv currentRuntimeEnv = Ray.getRuntimeContext().getCurrentRuntimeEnv();
      Assert.assertTrue(currentRuntimeEnv.isEmpty());
      RuntimeEnv runtimeEnv = new RuntimeEnv.Builder().build();
      RuntimeEnvConfig runtimeEnvConfig = new RuntimeEnvConfig(1, false);
      runtimeEnv.setConfig(runtimeEnvConfig);
      Integer result =
          Ray.task(RuntimeEnvTest::getRuntimeEnvTimeout).setRuntimeEnv(runtimeEnv).remote().get();
      Assert.assertNotNull(result);
      Assert.assertEquals((int) result, 1);
    } finally {
      Ray.shutdown();
    }
  }
}
