#
# JobLauncher.R
#
# Copyright (C) 2022 by RStudio, PBC
#
# Unless you have received this program directly from RStudio pursuant
# to the terms of a commercial license agreement with RStudio, then
# this program is licensed to you under the terms of version 3 of the
# GNU Affero General Public License. This program is distributed WITHOUT
# ANY EXPRESS OR IMPLIED WARRANTY, INCLUDING THOSE OF NON-INFRINGEMENT,
# MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. Please refer to the
# AGPL (http://www.gnu.org/licenses/agpl-3.0.txt) for more details.
#
#

# -------------------------------------------------------------
# Helpers
# -------------------------------------------------------------

.rs.addFunction("launcherhelper.fieldnames", function() {
  c("args", "cluster", "command", "config", "container", "environment",
    "exe", "exitCode", "exposedPorts", "host", "id", "lastUpdateTime",
    "mounts", "name", "placementConstraints", "queues", "resourceLimits",
    "status", "statusMessage", "stderrFile", "stdin", "stdoutFile", "submissionTime",
    "tags", "pid", "user", "workingDirectory")
})

.rs.addFunction("launcherhelper.validatefields", function(fields) {
  if (!is.null(fields) && !is.character(fields)) {
    stop("'fields' should be a character vector or NULL")
  }
  if (!is.null(fields) &&
      (sum(fields %in% .rs.launcherhelper.fieldnames()) != length(fields))) {
    stop(paste0("Supported 'fields' are: ",
           paste0(.rs.launcherhelper.fieldnames(), sep = "", collapse = ", "),
         sep = "", collapse = ""))
  }
})

.rs.addFunction("launcherhelper.statusnames", function() {
  c("Pending", "Running", "Suspended", "Failed", "Finished", "Killed", "Canceled")
})

.rs.addFunction("launcherhelper.validatestatuses", function(statuses) {
  if (!is.null(statuses) && !is.character(statuses)) {
    stop("'statuses should be a character vector or NULL")
  }
  if (!is.null(statuses) &&
      (sum(statuses %in% .rs.launcherhelper.statusnames()) != length(statuses))) {
    stop(paste0("Supported 'statuses' are: ",
           paste0(.rs.launcherhelper.statusnames(), sep = "", collapse = ", "),
         sep = "", collapse = ""))
  }
})

.rs.addFunction("launcherhelper.operations", function() {
  c("suspend", "resume", "stop", "kill", "cancel")
})

.rs.addFunction("launcherhelper.validateoperation", function(operation) {
  if (is.null(operation) || !is.character(operation) || length(operation) != 1 ||
      (!(operation %in% .rs.launcherhelper.operations()))) {
    stop("'operation' must be one of 'suspend', 'resume', 'stop', 'kill', or 'cancel'")
  }
})

.rs.addFunction("launcherhelper.resourcelimits", function() {
  c("cpuCount", "cpuFrequency", "cpuSet", "cpuTime", "memory", "memorySwap")
})

.rs.addFunction("launcherhelper.validateresourcelimit", function(limit) {
  if (is.null(limit) || !is.character(limit) || length(limit) != 1 ||
      (!(limit %in% .rs.launcherhelper.resourcelimits()))) {
    stop("'resourcelimit' must be one of 'cpuCount', 'cpuFrequency', 'cpuSet', 'cpuTime', 'memory', or 'memorySwap'")
  }
})

.rs.addFunction("launcherhelper.validatejobid", function(jobId) {
  if (is.null(jobId) || !is.character(jobId) || length(jobId) != 1)
    stop("'jobId' must be a single element character vector")
})

# ensure that all entries in list 'theList' inherit from 'expectedClass'
.rs.addFunction("checkListMemberClass", function(theList, expectedClass) {
  all(vapply(theList, inherits, expectedClass, FUN.VALUE = logical(1)))
})

# -------------------------------------------------------------
# invoke launcher job APIs from R
# -------------------------------------------------------------

.rs.addApiFunction("launcher.getInfo", function() {
  result <- .Call("rs_invokeServerRpc", "/job_launcher_rpc/get_info", NULL)
  if (result[["result"]] == FALSE) {
    return(NULL)
  }
  result
})

.rs.addApiFunction("launcher.jobsFeatureAvailable", function() {
   .Call("rs_launcherJobsFeatureAvailable")
})

.rs.addApiFunction("launcher.getJobs", function(statuses = NULL,
                                                fields = NULL,
                                                tags = NULL,
                                                includeSessions = FALSE) {
  .rs.launcherhelper.validatefields(fields)
  .rs.launcherhelper.validatestatuses(statuses)

  if (!is.null(tags) && !is.character(tags)) {
    stop("'tags' should be a character vector or NULL")
  }

  if (is.null(statuses))
    statuses <- list()
  if (is.null(fields)) {
    fields <- list()
  } else if (!includeSessions) {
    # must have "tags" field to exclude session jobs
    tagfield <- c("tags")
    fields <- union(fields, tagfield)
  }
  if (is.null(tags))
    tags <- list()

  args <- list(fields = fields, statuses = statuses, tags = tags)
  result <- .Call("rs_invokeServerRpc", "/job_launcher_rpc/get_jobs", args)
  if (result$result == FALSE) {
    return(NULL)
  }
  jobs <- result[["jobs"]]
  if (includeSessions) {
    jobs
  } else if (length(jobs) > 0) {
    jobs[sapply(jobs, function(x) { !("rstudio-r-session" %in% x$tags)})]
  } else {
    return(list())
  }
})

.rs.addApiFunction("launcher.getJob", function(jobId) {
  .rs.launcherhelper.validatejobid(jobId)
  result <- .Call("rs_invokeServerRpc", "/job_launcher_rpc/get_job", jobId)
  if (result[["result"]] == FALSE) {
    return(NULL)
  }
  result$job
})

.rs.addApiFunction("launcher.newConfig",
                function(name, value = NULL, valueType = NULL) {
  structure(list(name = .rs.scalar(name), value = .rs.scalar(value), valueType = .rs.scalar(valueType)),
            class = "rs_launcher_config")
})

.rs.addApiFunction("launcher.newContainer",
                function(image, runAsUserId = NULL, runAsGroupId = NULL) {
  structure(list(image = .rs.scalar(image), runAsUserId = .rs.scalar(runAsUserId), runAsGroupId = .rs.scalar(runAsGroupId)),
            class = "rs_launcher_container")
})

.rs.addApiFunction("launcher.newHostMount",
                function(path, mountPath, readOnly = TRUE) {

  hostMount <- list(path = .rs.scalar(path))
  hostMountSource <- list(type = .rs.scalar("host"), source = .rs.scalar(hostMount))
  structure(list(mountSource = hostMountSource, mountPath = .rs.scalar(mountPath), readOnly = .rs.scalar(readOnly)),
            class = c("rs_launcher_hostMount", "rs_launcher_mount"))
})

.rs.addApiFunction("launcher.newNfsMount",
                function(host, path, mountPath, readOnly = TRUE) {

  nfsMount <- list(host = .rs.scalar(host), path = .rs.scalar(path))
  nfsMountSource <- list(type = .rs.scalar("nfs"), source = .rs.scalar(nfsMount))
  structure(list(mountSource = nfsMountSource, mountPath = .rs.scalar(mountPath), readOnly = .rs.scalar(readOnly)),
            class = c("rs_launcher_nfsMount", "rs_launcher_mount"))
})

.rs.addApiFunction("launcher.newPlacementConstraint",
                function(name, value = NULL) {
  structure(list(name = .rs.scalar(name), value = .rs.scalar(value)),
            class = "rs_launcher_placement")
})

.rs.addApiFunction("launcher.newResourceLimit",
                function(type, value) {
  structure(list(type = .rs.scalar(type), value = .rs.scalar(value)),
            class = "rs_launcher_resourcelimit")
})

.rs.addApiFunction("launcher.submitJob",
                function(name,
                         args = NULL,
                         cluster = "Local",
                         command = NULL,
                         config = NULL,
                         container = NULL,
                         environment = NULL,
                         exe = NULL,
                         exposedPorts = NULL,
                         host = NULL,
                         mounts = NULL,
                         placementConstraints = NULL,
                         queues = NULL,
                         resourceLimits = NULL,
                         stderrFile = NULL,
                         stdin = NULL,
                         stdoutFile = NULL,
                         tags = NULL,
                         user = Sys.getenv("USER"),
                         workingDirectory = NULL,
                         applyConfigSettings = TRUE) {

  if (!is.null(environment) &&
      (!is.vector(environment, mode = "character") || is.null(names(environment))))
      stop("'environment' must be a named character vector")

  if (!is.null(config)) {
    if (!.rs.checkListMemberClass(config, "rs_launcher_config")) {
      stop("'config' must be of class 'rs_launcher_config'")
    }
  }

  if (!is.null(container)) {
    if (!inherits(container, "rs_launcher_container")) {
      stop("'container' must be of type 'rs_launcher_container'")
    }
    # allow container to be started without a command, or only one of 'command' or 'exe'; implies a container
    # that knows what work to do without being given a command
    if (length(c(command, exe)) > 1) {
      stop("cannot specify both 'command' and 'exe'")
    }
  } else if (length(c(command, exe)) != 1) {
    stop("'command' or 'exe' must be supplied, not both")
  }

  # the launcher supports specifying protocol and publishedPort in addition to
  # the targetPort we are exposing here; this limitation matches the current
  # limitation for customizing ports via the launcher-ports and launcher-adhoc-ports
  # config files
  if (!is.null(exposedPorts) && !is.vector(exposedPorts, mode = "integer")) {
    stop("'exposedPorts' must be a integer vector")
  }

  if (!is.null(mounts)) {
    if (!.rs.checkListMemberClass(mounts, "rs_launcher_mount")) {
      stop("'mounts' must be of class 'rs_launcher_mounts'")
    }
  }

  if (!is.null(placementConstraints)) {
    if (!.rs.checkListMemberClass(placementConstraints, "rs_launcher_placement")) {
      stop("'placementConstraints' must be of class 'rs_launcher_placement'")
    }
  }

  if (!is.null(resourceLimits)) {
    if (!.rs.checkListMemberClass(resourceLimits, "rs_launcher_resourcelimit")) {
      stop("'resourceLimits' must be of class 'rs_launcher_resourcelimit'")
    }
    lapply(resourceLimits, function(x) { .rs.launcherhelper.validateresourcelimit(x$type) })
  }

  if (!is.null(container) && !is.null(container$image) && is.null(container$runAsUserId)) {
    userInfo <- .Call("rs_launcherJobContainerUser")
    if (!is.null(userInfo)) {
      container$runAsUserId <- .rs.scalar(userInfo$uid)
      container$runAsGroupId <- .rs.scalar(userInfo$gid)

      if (!is.null(environment)) {
        if (is.na(environment["HOME"])) {
          environment <- c(HOME = userInfo$home, environment)
        }
      } else {
        environment <- c(HOME = userInfo$home)
      }
    }
  }

  env <- NULL
  if (!is.null(environment)) {
    for (i in 1:length(environment)) {
      env[[i]] <- .rs.scalarListFromList(list(name = names(environment)[[i]], value = environment[[i]]))
    }
  }

  ports <- lapply(exposedPorts, function(x) .rs.scalarListFromList(list(targetPort = x)))

  job <- list(args = args,
              cluster = .rs.scalar(cluster),
              command = .rs.scalar(command),
              config = config,
              container = container,
              environment = env,
              exe = .rs.scalar(exe),
              exposedPorts = ports,
              host = .rs.scalar(host),
              mounts = mounts,
              name = .rs.scalar(name),
              placementConstraints = placementConstraints,
              queues = queues,
              resourceLimits = resourceLimits,
              stderrFile = .rs.scalar(stderrFile),
              stdin = .rs.scalar(stdin),
              stdoutFile = .rs.scalar(stdoutFile),
              tags = tags,
              user = .rs.scalar(user),
              workingDirectory = .rs.scalar(workingDirectory))

  args <- list(job = .rs.scalar(job), applyConfigSettings = .rs.scalar(applyConfigSettings))

  result <- .Call("rs_invokeServerRpc", "/job_launcher_rpc/api_submit_job", args)
  if (result$result == FALSE) {
    return(NULL)
  }
  result$job$id
})

.rs.addApiFunction("launcher.controlJob", function(jobId, operation) {
  .rs.launcherhelper.validatejobid(jobId)
  .rs.launcherhelper.validateoperation(operation)
  .Call("rs_invokeServerRpc", "/job_launcher_rpc/control_job", c(jobId, operation))
})

.rs.addApiFunction("launcher.startStatusStream", function(jobId = "*") {
  invisible(.Call("rs_launcherStartStatusStream", jobId))
})

.rs.addApiFunction("launcher.stopStatusStream", function(jobId = "*") {
  invisible(.Call("rs_launcherStopStatusStream", jobId))
})

.rs.addApiFunction("launcher.streamOutput", function(jobId, listening) {
  invisible(.Call("rs_launcherStreamOutput", jobId, listening))
})

