# Copyright (C) 2023-2024 Hibiki AI Limited <info@hibiki-ai.com>
#
# This file is part of mirai.
#
# mirai is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# mirai is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# mirai. If not, see <https://www.gnu.org/licenses/>.

# mirai ------------------------------------------------------------------------

#' Launch Daemon
#'
#' \code{launch_local} spawns a new background \code{Rscript} process calling
#' \code{\link{daemon}} with the specified arguments.
#'
#' These functions may be used to re-launch daemons that have exited after
#' reaching time or task limits.
#'
#' Daemons must already be set for launchers to work.
#'
#' The generated command contains the argument \sQuote{rs} specifying the length
#' 7 L'Ecuyer-CMRG random seed supplied to the daemon. The values will be
#' different each time the function is called.
#'
#' @inheritParams mirai
#' @param url the character host URL or vector of host URLs, including the port
#'   to connect to (and optionally for websockets, a path), e.g.
#'   'tcp://hostname:5555' or 'ws://10.75.32.70:5555/path'
#'
#'   \strong{or} integer index value, or vector of index values, of the
#'   dispatcher URLs, or 1L for the host URL (when not using dispatcher).
#'
#'   \strong{or} for \code{launch_remote} only, a \sQuote{miraiCluster} or
#'   \sQuote{miraiNode}.
#' @param ... (optional) additional arguments passed through to
#'   \code{\link{daemon}}. These include \sQuote{autoexit}, \sQuote{cleanup},
#'   \sQuote{output}, \sQuote{maxtasks}, \sQuote{idletime}, \sQuote{walltime}
#'   and \sQuote{timerstart}.
#' @param tls [default NULL] required for secure TLS connections over tls+tcp or
#'   wss. Zero-configuration TLS certificates generated by \code{\link{daemons}}
#'   are automatically passed to the daemon, without requiring to be specified
#'   here. Otherwise, supply \strong{either} the character path to a file
#'   containing X.509 certificate(s) in PEM format, comprising the certificate
#'   authority certificate chain, \strong{or} a length 2 character vector
#'   comprising [i] the certificate authority certificate chain and [ii] the
#'   empty character ''.
#'
#' @return For \strong{launch_local}: Invisible NULL.
#'
#' @examples
#' if (interactive()) {
#' # Only run examples in interactive R sessions
#'
#' daemons(url = host_url(ws = TRUE), dispatcher = "none")
#' status()
#' launch_local(status()$daemons, maxtasks = 10L)
#' launch_remote(1L, maxtasks = 10L)
#' Sys.sleep(1)
#' status()
#' daemons(0)
#'
#' daemons(n = 2L, url = host_url(tls = TRUE))
#' status()
#' launch_local(1:2, idletime = 60000L, timerstart = 1L)
#' launch_remote(1:2, idletime = 60000L, timerstart = 1L)
#' Sys.sleep(1)
#' status()
#' daemons(0)
#'
#' }
#'
#' @export
#'
launch_local <- function(url, ..., tls = NULL, .compute = "default") {

  envir <- ..[[.compute]]
  is.null(envir) && stop(._[["daemons_unset"]])
  dots <- parse_dots(...)
  output <- attr(dots, "output")
  if (is.null(tls)) tls <- envir[["tls"]]
  url <- process_url(url = url, envir = envir)
  is.character(url) || stop(._[["url_spec"]])
  for (u in url)
    launch_daemon(wa3(u, dots, next_stream(envir), tls), output)

}

#' Launch Daemon
#'
#' \code{launch_remote} returns the shell command for deploying daemons as a
#' character vector. If a configuration generated by \code{\link{remote_config}}
#' or \code{\link{ssh_config}} is supplied then this is used to launch the
#' daemon on the remote machine.
#'
#' @param remote required only for launching remote daemons, a configuration
#'   generated by \code{\link{remote_config}} or \code{\link{ssh_config}}. An
#'   empty \code{\link{remote_config}} does not effect any daemon launches but
#'   returns the shell commands for deploying manually on remote machines.
#'
#' @return For \strong{launch_remote}: A character vector of daemon launch
#'   commands, classed as \sQuote{miraiLaunchCmd}. The printed output may be
#'   copy / pasted directly to the remote machine.
#'
#' @rdname launch_local
#' @export
#'
launch_remote <- function(url, remote = remote_config(), ..., tls = NULL, .compute = "default") {

  if (!is.character(url) && inherits(url, c("miraiCluster", "miraiNode"))) {
    .compute <- attr(url, "id")
    url <- rep(..[[.compute]][["urls"]], max(length(url), 1L))
  }
  envir <- ..[[.compute]]
  is.null(envir) && stop(._[["daemons_unset"]])
  dots <- parse_dots(...)
  if (is.null(tls)) tls <- envir[["tls"]]
  url <- process_url(url = url, envir = envir)
  is.character(url) || stop(._[["url_spec"]])

  ulen <- length(url)
  command <- remote[["command"]]
  rscript <- remote[["rscript"]]
  quote <- remote[["quote"]]

  if (length(command)) {

    args <- remote[["args"]]

    if (is.list(args)) {

      if (length(args) == 1L) {
        args <- args[[1L]]

      } else if (ulen == 1L || ulen == length(args)) {

        arglen <- length(args)
        cmds <- character(arglen)
        for (i in seq_along(args))
          cmds[i] <- sprintf("%s -e %s", rscript, wa3(url[min(i, ulen)], dots, next_stream(envir), tls))

        for (i in seq_along(args))
          system2(command = command, args = `[<-`(args[[i]], find_dot(args[[i]]), if (quote) shQuote(cmds[i]) else cmds[i]), wait = FALSE)

        return(`class<-`(cmds, "miraiLaunchCmd"))

      } else {
        stop(._[["arglen"]])
      }

    }
  }

  cmds <- character(ulen)
  for (i in seq_along(url))
    cmds[i] <- sprintf("%s -e %s", rscript, wa3(url[i], dots, next_stream(envir), tls))

  if (length(command))
    for (cmd in cmds)
      system2(command = command, args = `[<-`(args, find_dot(args), if (quote) shQuote(cmd) else cmd), wait = FALSE)

  `class<-`(cmds, "miraiLaunchCmd")

}

#' Generic and SSH Remote Launch Configuration
#'
#' \code{remote_config} provides a flexible generic framework for generating the
#' shell commands to deploy daemons remotely.
#'
#' @param command the command used to effect the daemon launch on the remote
#'   machine as a character string (e.g. \code{'ssh'}). Defaults to \sQuote{ssh}
#'   for \code{ssh_config}, although may be substituted for the full path to a
#'   specific SSH application. The default NULL for \code{remote_config} does
#'   not effect any launches, but causes \code{\link{launch_remote}} to return
#'   the shell commands for manual deployment on remote machines.
#' @param args (optional) arguments passed to \sQuote{command}, as a character
#'   vector that must include \code{"."} as an element, which will be
#'   substituted for the daemon launch command. Alternatively, a list of such
#'   character vectors to effect multiple launches (one for each list element).
#' @param rscript (optional) name / path of the Rscript executable on the remote
#'   machine. The default assumes \sQuote{Rscript} is on the executable search
#'   path. Prepend the full path if necessary. If launching on Windows,
#'   \sQuote{Rscript} should be replaced with \sQuote{Rscript.exe}.
#' @param quote [default FALSE] logical value whether or not to quote the daemon
#'   launch command (not required for Slurm \sQuote{srun} for example, but
#'   required for \sQuote{ssh}).
#'
#' @return A list in the required format to be supplied to the \sQuote{remote}
#'   argument of \code{\link{launch_remote}}, \code{\link{daemons}}, or
#'   \code{\link{make_cluster}}.
#'
#' @examples
#' # example for Slurm
#' remote_config(
#'   command = "srun",
#'   args = c("--mem 512", "-n 1", "."),
#'   rscript = file.path(R.home("bin"), "Rscript")
#' )
#'
#' # commands like SSH require quoting of the daemon launch command
#' remote_config(
#'   command = "/usr/bin/ssh",
#'   args = c("-fTp 22 10.75.32.90", "."),
#'   quote = TRUE
#' )
#'
#' # can be used to start local dameons with special configurations
#' remote_config(
#'   command = "Rscript",
#'   rscript = "--default-packages=NULL --vanilla"
#' )
#'
#' @export
#'
remote_config <- function(command = NULL, args = c("", "."), rscript = "Rscript", quote = FALSE) {

  if (is.list(args)) lapply(args, find_dot) else find_dot(args)
  list(command = command, args = args, rscript = rscript, quote = quote)

}

#' SSH Remote Launch Configuration
#'
#' \code{ssh_config} generates a remote configuration for launching daemons over
#' SSH, with the option of SSH tunnelling.
#'
#' @param remotes the character URL or vector of URLs to SSH into, using the
#'   'ssh://' scheme and including the port open for SSH connections (defaults
#'   to 22 if not specified), e.g. 'ssh://10.75.32.90:22' or 'ssh://nodename'.
#' @param tunnel [default FALSE] logical value whether to use SSH reverse
#'   tunnelling. If TRUE, a tunnel is created between the same ports on the
#'   local and remote machines. See the \sQuote{SSH Tunnelling} section below
#'   for how to correctly specify required settings.
#' @param timeout [default 10] maximum time allowed for connection setup in
#'   seconds.
#' @param host (optional) only applicable for reverse tunnelling. Should be
#'   specified if creating a standalone configuration object. If calling this
#'   function directly as an argument to \code{\link{daemons}}, this is not
#'   required and can be inferred from the \sQuote{url} supplied (see
#'   \sQuote{SSH Tunnelling} section below).
#'
#' @section SSH Direct Connections:
#'
#' The simplest use of SSH is to execute the daemon launch command on a remote
#' machine, for it to dial back to the host / dispatcher URL.
#'
#' It is assumed that SSH key-based authentication is already in place. The
#' relevant port on the host must also be open to inbound connections from the
#' remote machine.
#'
#' @section SSH Tunnelling:
#'
#' Use of SSH tunnelling provides a convenient way to launch remote daemons
#' without requiring the remote machine to be able to access the host. Often
#' firewall configurations or security policies may prevent opening a port to
#' accept outside connections.
#'
#' In these cases SSH tunnelling offers a solution by creating a tunnel once the
#' initial SSH connection is made. For simplicity, this SSH tunnelling
#' implementation uses the same port on both the side of the host and that of
#' the daemon. SSH key-based authentication must also already be in place.
#'
#' Tunnelling requires the hostname for the \sQuote{host} argument (or the
#' \sQuote{url} argument to \code{\link{daemons}} if called directly in that
#' context) to be either \sQuote{127.0.0.1} or \sQuote{localhost}. This is as
#' the tunnel is created between \code{127.0.0.1:port} or equivalently
#' \code{localhost:port} on each machine. The host listens to \code{port} on its
#' machine and the remotes each dial into \code{port} on their own respective
#' machines.
#'
#' @examples
#' # simple SSH example
#' ssh_config(
#'   remotes = c("ssh://10.75.32.90:222", "ssh://nodename"),
#'   timeout = 5
#' )
#'
#' # SSH tunnelling example
#' ssh_config(
#'   remotes = c("ssh://10.75.32.90:222", "ssh://nodename"),
#'   tunnel = TRUE,
#'   host = "tls+tcp://127.0.0.1:5555"
#' )
#'
#' \dontrun{
#'
#' # launch 2 daemons on the remote machines 10.75.32.90 and 10.75.32.91 using
#' # SSH, connecting back directly to the host URL over a TLS connection:
#'
#' daemons(
#'   url = host_url(tls = TRUE),
#'   remote = ssh_config(
#'     remotes = c("ssh://10.75.32.90:222", "ssh://10.75.32.91:222"),
#'     timeout = 1
#'   )
#' )
#'
#' # launch 2 nodes on the remote machine 10.75.32.90 using SSH tunnelling over
#' # port 5555 ('url' hostname must be 'localhost' or '127.0.0.1'):
#'
#' cl <- make_cluster(
#'   url = "tcp://localhost:5555",
#'   remote = ssh_config(
#'     remotes = c("ssh://10.75.32.90", "ssh://10.75.32.90"),
#'     tunnel = TRUE,
#'     timeout = 1
#'   )
#' )
#' }
#'
#' @rdname remote_config
#' @export
#'
ssh_config <- function(remotes, tunnel = FALSE, timeout = 10, command = "ssh", rscript = "Rscript", host) {

  premotes <- lapply(remotes, parse_url)
  hostnames <- lapply(premotes, .subset2, "hostname")
  ports <- lapply(premotes, .subset2, "port")

  if (tunnel) {
    if (missing(host)) {
      host <- dynGet("url", ifnotfound = NULL)
      is.null(host) && stop(._[["correct_context"]])
    }
    purl <- lapply(host, parse_check_local_url)
    plen <- length(purl)
  }

  rlen <- length(remotes)
  args <- vector(mode = "list", length = if (tunnel) max(rlen, plen) else rlen)

  for (i in seq_along(args)) {
    args[[i]] <- c(
      if (tunnel) sprintf("-R %s:%s", purl[[min(i, plen)]][["port"]], purl[[min(i, plen)]][["host"]]),
      sprintf("-o ConnectTimeout=%s -fTp %s", as.character(timeout), ports[[min(i, rlen)]]),
      hostnames[[min(i, rlen)]],
      "."
    )
  }

  list(command = command, args = args, rscript = rscript, quote = TRUE)

}

#' URL Constructors
#'
#' \code{host_url} constructs a valid host URL (at which daemons may connect)
#' based on the computer's hostname. This may be supplied directly to the
#' \sQuote{url} argument of \code{\link{daemons}}.
#'
#' \code{host_url} relies on using the host name of the computer rather than an
#' IP address and typically works on local networks, although this is not always
#' guaranteed. If unsuccessful, substitute an IPv4 or IPv6 address in place of
#' the hostname.
#'
#' \code{local_url} generates a random URL for the platform's default
#' inter-process communications transport: abstract Unix domain sockets on
#' Linux, Unix domain sockets on MacOS, Solaris and other POSIX platforms, and
#' named pipes on Windows.
#'
#' @param ws [default FALSE] logical value whether to use a WebSockets 'ws://'
#'   or else TCP 'tcp://' scheme.
#' @param tls [default FALSE] logical value whether to use TLS in which case the
#'   scheme used will be either 'wss://' or 'tls+tcp://' accordingly.
#' @param port [default 0] numeric port to use. This should be open to
#'   connections from the network addresses the daemons are connecting from.
#'   \sQuote{0} is a wildcard value that automatically assigns a free ephemeral
#'   port.
#'
#' @return A character string comprising a valid URL.
#'
#' @examples
#' host_url()
#' host_url(ws = TRUE)
#' host_url(tls = TRUE)
#' host_url(ws = TRUE, tls = TRUE, port = 5555)
#'
#' @export
#'
host_url <- function(ws = FALSE, tls = FALSE, port = 0)
  sprintf(
    "%s://%s:%s",
    if (ws) { if (tls) "wss" else "ws" } else { if (tls) "tls+tcp" else "tcp" },
    Sys.info()[["nodename"]],
    as.character(port)
  )

#' URL Constructors
#'
#' \code{local_url} constructs a random URL suitable for local daemons.
#'
#' @examples
#' local_url()
#'
#' @rdname host_url
#' @export
#'
local_url <- function() sprintf("%s%s", .urlscheme, random(12L))

#' @export
#'
print.miraiLaunchCmd <- function(x, ...) {

  for (i in seq_along(x))
    cat(sprintf("[%d]\n%s\n\n", i, x[i]), file = stdout())
  invisible(x)

}

# internals --------------------------------------------------------------------

find_dot <- function(args) {
  sel <- args == "."
  any(sel) || stop(._[["dot_required"]], call. = FALSE)
  sel
}

process_url <- function(url, envir) {
  if (is.numeric(url)) {
    vec <- envir[["urls"]]
    all(url >= 1L, url <= length(vec)) || return(FALSE)
    url <- vec[url]
  } else {
    lapply(url, parse_url)
  }
  url
}

parse_check_local_url <- function(url) {
  purl <- parse_url(url)
  purl[["hostname"]] == "127.0.0.1" || purl[["hostname"]] == "localhost" || stop(._[["requires_local"]], call. = FALSE)
  purl
}
