# 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 x parallel -------------------------------------------------------------

#' Make Mirai Cluster
#'
#' \code{make_cluster} creates a cluster of type \sQuote{miraiCluster}, which
#' may be used as a cluster object for any function in the \pkg{parallel} base
#' package such as \code{\link[parallel]{clusterApply}} or
#' \code{\link[parallel]{parLapply}}.
#'
#' @param n integer number of nodes (automatically launched on the local machine
#'   unless \sQuote{url} is supplied).
#' @param url [default NULL] (specify for remote nodes) the character URL on the
#'   host for remote nodes to dial into, including a port accepting incoming
#'   connections, e.g. 'tcp://10.75.37.40:5555'. Specify a URL with the scheme
#'   \sQuote{tls+tcp://} to use secure TLS connections.
#' @param remote [default NULL] (specify to launch remote nodes) a remote launch
#'   configuration generated by \code{\link{remote_config}} or
#'   \code{\link{ssh_config}}. If not supplied, nodes may be deployed manually
#'   on remote resources.
#' @param ... additional arguments passed onto \code{\link{daemons}}.
#'
#' @return For \strong{make_cluster}: An object of class \sQuote{miraiCluster}
#'   and \sQuote{cluster}. Each \sQuote{miraiCluster} has an automatically
#'   assigned ID and \sQuote{n} nodes of class \sQuote{miraiNode}. If
#'   \sQuote{url} is supplied but not \sQuote{remote}, the shell commands for
#'   deployment of nodes on remote resources are printed to the console.
#'
#'   For \strong{stop_cluster}: invisible NULL.
#'
#' @section Remote Nodes:
#'
#' Specify \sQuote{url} and \sQuote{n} to set up a host connection for remote
#' nodes to dial into. \sQuote{n} defaults to one if not specified.
#'
#' Also specify \sQuote{remote} to launch the nodes using a configuration
#' generated by \code{\link{remote_config}} or \code{\link{ssh_config}}. In this
#' case, the number of nodes is inferred from the configuration provided and
#' \sQuote{n} is disregarded.
#'
#' If \sQuote{remote} is not supplied, the shell commands for deploying nodes
#' manually on remote resources are automatically printed to the console.
#'
#' \code{\link{launch_remote}} may be called at any time on a
#' \sQuote{miraiCluster} to return the shell commands for deployment of all
#' nodes, or on a \sQuote{miraiNode} to return the command for a single node.
#'
#' @section Status:
#'
#' Call \code{\link{status}} on a \sQuote{miraiCluster} to check the number of
#' currently active connections as well as the host URL.
#'
#' @section Errors:
#'
#' Errors are thrown by the \sQuote{parallel} mechanism if one or more nodes
#' failed (quit unexpectedly). The resulting \sQuote{errorValue} returned is 19
#' (Connection reset). Other types of error, e.g. in evaluation, should result
#' in the usual \sQuote{miraiError} being returned.
#'
#' @note The default behaviour of clusters created by this function is designed
#'   to map as closely as possible to clusters created by the \pkg{parallel}
#'   package. However, \sQuote{...} arguments are passed onto
#'   \code{\link{daemons}} for additional customisation if desired, although
#'   resultant behaviour may not always be supported.
#'
#' @examples
#' if (interactive()) {
#' # Only run examples in interactive R sessions
#'
#' cl <- make_cluster(2)
#' cl
#' cl[[1L]]
#'
#' Sys.sleep(0.5)
#' status(cl)
#'
#' stop_cluster(cl)
#'
#' }
#'
#' @export
#'
make_cluster <- function(n, url = NULL, remote = NULL, ...) {

  id <- sprintf("`%d`", length(..))
  cvs <- cv()

  if (is.character(url)) {

    length(url) == 1L || stop(._[["single_url"]])
    daemons(url = url, remote = remote, dispatcher = "none", cleanup = FALSE, ..., .compute = id)

    if (is.null(remote)) {
      if (missing(n)) n <- 1L
      is.numeric(n) || stop(._[["numeric_n"]])
      cat("Shell commands for deployment on nodes:\n\n", file = stdout())
      print(launch_remote(rep(..[[id]][["urls"]], n), .compute = id))
    } else {
      args <- remote[["args"]]
      n <- if (is.list(args)) length(args) else 1L
    }

  } else {
    is.numeric(n) || stop(._[["numeric_n"]])
    n >= 1L || stop(._[["n_one"]])
    daemons(n = n, dispatcher = "none", cleanup = FALSE, ..., .compute = id)
  }

  `[[<-`(..[[id]], "cvs", cvs)

  cl <- lapply(seq_len(n), create_node, id = id)
  `attributes<-`(cl, list(class = c("miraiCluster", "cluster"), id = id))

}

#' Stop Mirai Cluster
#'
#' \code{stop_cluster} stops a cluster created by \code{make_cluster}.
#'
#' @param cl a \sQuote{miraiCluster}.
#'
#' @rdname make_cluster
#' @export
#'
stop_cluster <- function(cl)
  daemons(n = 0L, .compute = attr(cl, "id")) || return(invisible())

#' @exportS3Method parallel::stopCluster
#'
stopCluster.miraiCluster <- stop_cluster

#' @exportS3Method parallel::sendData
#'
sendData.miraiNode <- function(node, data) {

  id <- attr(node, "id")
  envir <- ..[[id]]
  is.null(envir) && stop(._[["cluster_inactive"]])

  value <- data[["data"]]
  tagged <- !is.null(value[["tag"]])
  if (tagged) set_cv(envir) else unset_cv(envir)

  m <- mirai(do.call(node, data, quote = TRUE), node = value[["fun"]], data = value[["args"]], .compute = id)
  if (tagged) assign("tag", value[["tag"]], m)
  `[[<-`(node, "mirai", m)

}

#' @exportS3Method parallel::recvData
#'
recvData.miraiNode <- function(node) call_aio(.subset2(node, "mirai"))

#' @exportS3Method parallel::recvOneData
#'
recvOneData.miraiCluster <- function(cl) {

  wait(..[[attr(cl, "id")]][["cv"]])
  node <- which.min(lapply(cl, node_unresolved))
  m <- .subset2(.subset2(cl, node), "mirai")
  list(node = node, value = `class<-`(m, NULL))

}

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

  id <- attr(x, "id")
  cat(sprintf("< miraiCluster | ID: %s nodes: %d active: %s >\n", id, length(x), as.logical(length(..[[id]]))), file = stdout())
  invisible(x)

}

#' @export
#'
`[.miraiCluster` <- function(x, ...) .subset(x, ...)

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

  cat(sprintf("< miraiNode | node: %d cluster ID: %s >\n", attr(x, "node"), attr(x, "id")), file = stdout())
  invisible(x)

}

#' Register Mirai Cluster
#'
#' Registers 'miraiCluster' with the \pkg{parallel} package as cluster type
#' 'MIRAI' and (optionally) makes it the default cluster type.
#'
#' @param default [default TRUE] logical value whether to also register
#'   'miraiCluster' as the default cluster type.
#'
#' @return Invisible NULL
#'
#' @note The underlying implementation in the \pkg{parallel} package this
#'   function relies on iis currently only available in R-devel (4.5).
#'
#' @examples
#' tryCatch(mirai::register_cluster(), error = identity)
#'
#' @keywords internal
#' @export
#'
register_cluster <- function(default = TRUE) {

  register <- .getNamespace("parallel")[["registerClusterType"]]
  is.null(register) && stop(._[["register_cluster"]])
  register("MIRAI", make_cluster, make.default = isTRUE(default)) # nocov

}

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

create_node <- function(node, id)
  `attributes<-`(
    new.env(hash = FALSE, parent = emptyenv()),
    list(class = "miraiNode", node = node, id = id)
  )

set_cv <- function(envir) `[[<-`(envir, "cv", envir[["cvs"]])

unset_cv <- function(envir) `[[<-`(envir, "cv", NULL)

node_unresolved <- function(node) {
  m <- .subset2(node, "mirai")
  unresolved(m) || !is.object(m)
}
