#' Save plot as a static image
#' 
#' Static image exporting via [the kaleido python
#' package](https://github.com/plotly/Kaleido/). `kaleido()` imports
#' kaleido into a \pkg{reticulate}d Python session and returns a `$transform()`
#' method for converting R plots into static images. `save_image()` provides a convenience wrapper around `kaleido()$transform()`. 
#' 
#' @section Installation:
#' 
#' `kaleido()` requires [the kaleido python
#' package](https://github.com/plotly/Kaleido/) to be usable via the
#' \pkg{reticulate} package. If you're starting from scratch, you install
#' eveything you need with the following R code:
#' 
#' ```
#' install.packages("reticulate")
#' library(reticulate)
#' use_python(install_python())
#' py_install(c("kaleido", "plotly"))
#' ```
#' 
#' @param ... not currently used.
#' @param p a plot object.
#' @param file a file path with a suitable file extension (png, jpg, jpeg,
#'   webp, svg, or pdf).
#' @param width,height The width/height of the exported image in layout
#'   pixels. If `scale` is 1, this will also be the width/height of the exported
#'   image in physical pixels.
#' @param scale The scale factor to use when exporting
#'   the figure. A scale factor larger than 1.0 will increase the image
#'   resolution with respect to the figure's layout pixel dimensions. Whereas as
#'   scale factor of less than 1.0 will decrease the image resolution.
#' @export
#' @return For `save_image()`, the generated `file`. For `kaleido()`, an environment that contains:
#'   * `transform()`: a function to convert plots objects into static images. This function has the same signature (i.e., arguments) as `save_image()`
#'   * `shutdown()`: a function for shutting down any currently running subprocesses 
#'     that were launched via `transform()`
#'   * `scope`: a reference to the underlying `kaleido.scopes.plotly.PlotlyScope`
#'     python object. Modify this object to customize the underlying Chromium
#'     subprocess and/or configure other details such as URL to plotly.js, MathJax, etc. 
#' @examplesIf interactive() || !identical(.Platform$OS.type, "windows")
#' 
#' \dontrun{
#'   # Save a single image
#'   p <- plot_ly(x = 1:10)
#'   tmp <- tempfile(fileext = ".png")
#'   save_image(p, tmp)
#'   file.show(tmp)
#' 
#'   # Efficiently save multiple images
#'   scope <- kaleido()
#'   for (i in 1:5) {
#'     scope$transform(p, tmp)
#'   }
#'   # Remove and garbage collect to remove 
#'   # R/Python objects and shutdown subprocesses
#'   rm(scope); gc()
#' }
#' 
save_image <- function(p, file, ..., width = NULL, height = NULL, scale = NULL) {
  kaleido()$transform(
    p, file, ..., width = width, height = height, scale = scale
  )
}

#' @rdname save_image
#' @export
kaleido <- function(...) {
  rlang::check_installed("reticulate")
  
  call_env <- rlang::caller_env()
  
  if (!reticulate::py_available(TRUE)) {
    rlang::abort(c("`{reticulate}` wasn't able to find a Python environment.",
      i = "If you have an existing Python installation, use `reticulate::use_python()` to inform `{reticulate}` of it.",
      i = "To have `{reticulate}` install Python for you, `reticulate::install_python()`."
    ), call = call_env)
  }
  
  tryCatch(
    reticulate::import("plotly"),
    error = function(e) {
      rlang::abort(c(
        "The `plotly` Python package is required for static image exporting.",
        i = "Please install it via `reticulate::py_install('plotly')`."
      ), call = call_env)
    }
  )
  
  kaleido <- tryCatch(
    reticulate::import("kaleido"),
    error = function(e) {
      rlang::abort(c(
        "The `kaleido` Python package is required for static image exporting.",
        i = "Please install it via `reticulate::py_install('kaleido')`."
      ), call = call_env)
    }
  )
  
  # Check for kaleido v1 by looking for v1-specific function
  res <- if (reticulate::py_has_attr(kaleido, "write_fig_sync")) {
    newKaleidoScope(kaleido)
  } else {
    legacyKaleidoScope(kaleido)
  }
  
  class(res) <- "kaleidoScope"
  res
}

newKaleidoScope <- function(kaleido) {
  force(kaleido)
  list(
    scopes = NULL,
    transform = function(p, file, ..., width = NULL, height = NULL, scale = NULL) {
      # Perform JSON conversion exactly how the R package would do it
      fig_data <- plotly_build(p)$x[c("data", "layout", "config")]

      # Inject mapbox token into layout.mapbox.accesstoken if available
      # We use layout instead of config because Kaleido's parser preserves
      # layout but drops config. This handles the case where users set
      # MAPBOX_TOKEN env var but don't use plot_mapbox()
      mapbox <- Sys.getenv("MAPBOX_TOKEN", NA)
      if (!is.na(mapbox) && is.null(fig_data$layout$mapbox$accesstoken)) {
        fig_data$layout$mapbox$accesstoken <- mapbox
      }

      fig <- to_JSON(fig_data)

      # Write to JSON file
      tmp_json <- tempfile(fileext = ".json")
      on.exit(unlink(tmp_json))
      writeLines(fig, tmp_json)

      # Import it as a fig (dict)
      .py_run_string_with_context(
        "import json; fig = json.load(open(tmp_json_path))",
        context = list(tmp_json_path = tmp_json)
      )

      # Gather figure-level options
      opts <- list(
        format = tools::file_ext(file),
        width = reticulate::r_to_py(width),
        height = reticulate::r_to_py(height),
        scale = reticulate::r_to_py(scale)
      )

      # Pass the R plotly.js bundle path to Kaleido
      kopts <- list(plotlyjs = plotlyMainBundlePath())

      # Write the figure to a file using kaleido
      kaleido$write_fig_sync(reticulate::py$fig, file, opts = opts, kopts = kopts)
    },
    shutdown = function() {}
  )
}


legacyKaleidoScope <- function(kaleido) {
  py <- reticulate::py
  scope_name <- paste0("scope_", new_id())
  py[[scope_name]] <- kaleido$scopes$plotly$PlotlyScope(
    plotlyjs = plotlyMainBundlePath()
  )
  
  scope <- py[[scope_name]]
  
  mapbox <- Sys.getenv("MAPBOX_TOKEN", NA)
  if (!is.na(mapbox)) {
    scope$mapbox_access_token <- mapbox
  }
  
  res <- list2env(list(
    scope = scope,
    # https://github.com/plotly/Kaleido/blob/6a46ecae/repos/kaleido/py/kaleido/scopes/plotly.py#L78-L106
    transform = function(p, file = "figure.png", ..., width = NULL, height = NULL, scale = NULL) {
      # Perform JSON conversion exactly how the R package would do it
      # (this is essentially plotly_json(), without the additional unneeded info)
      # and attach as an attribute on the python scope object
      scope[["_last_plot"]] <- to_JSON(
        plotly_build(p)$x[c("data", "layout", "config")]
      )
      # On the python side, _last_plot is a string, so use json.loads() to 
      # convert to dict(). This should be fine since json is a dependency of the
      # BaseScope() https://github.com/plotly/Kaleido/blob/586be5/repos/kaleido/py/kaleido/scopes/base.py#L2
      transform_cmd <- sprintf(
        "%s.transform(sys.modules['json'].loads(%s._last_plot), format='%s', width=%s, height=%s, scale=%s)",
        scope_name, scope_name, tools::file_ext(file), 
        reticulate::r_to_py(width), reticulate::r_to_py(height), 
        reticulate::r_to_py(scale)
      )
      # Write the base64 encoded string that transform() returns to disk
      # https://github.com/plotly/Kaleido/blame/master/README.md#L52
      .py_run_string_with_context(
        sprintf("import sys; open(output_file, 'wb').write(%s)", transform_cmd),
        context = list(output_file = file)
      )
      
      invisible(file)
    },
    # Shutdown the kaleido subprocesses
    # https://github.com/plotly/Kaleido/blob/586be5c/repos/kaleido/py/kaleido/scopes/base.py#L71-L72
    shutdown = function() {
      reticulate::py_run_string(paste0(scope_name, ".__del__()"))
    }
  ))
  
  # Shutdown subprocesses and delete python scope when 
  # this object is garbage collected by R
  reg.finalizer(res, onexit = TRUE, function(x) {
    x$shutdown()
    reticulate::py_run_string(paste("del", scope_name))
  })
  
  res
}

.py_run_string_with_context <- function(code, context = list(), convert = TRUE) {
  context_names <- names(context)
  old_values <- vector("list", length(context))
  had_value <- logical(length(context))
  was_set <- logical(length(context))
  
  if (length(context) > 0) {
    if (is.null(context_names) || any(context_names == "")) {
      rlang::abort("`context` must be a named list.")
    }
    if (any(!grepl("^[A-Za-z_][A-Za-z0-9_]*$", context_names))) {
      rlang::abort("`context` names must be valid Python identifiers.")
    }

    py <- reticulate::py
    on.exit({
      for (i in rev(which(was_set))) {
        name <- context_names[[i]]
        if (had_value[[i]]) {
          reticulate::py_set_attr(py, name, old_values[[i]])
        } else {
          reticulate::py_del_attr(py, name)
        }
      }
    }, add = TRUE)
    
    for (i in seq_along(context)) {
      name <- context_names[[i]]
      had_value[[i]] <- reticulate::py_has_attr(py, name)
      if (had_value[[i]]) {
        old_values[[i]] <- reticulate::py_get_attr(py, name)
      }
      reticulate::py_set_attr(py, name, context[[i]])
      was_set[[i]] <- TRUE
    }
  }
  
  reticulate::py_run_string(code, convert = convert)
}


#' Print method for kaleido
#' 
#' S3 method for [kaleido()].
#' 
#' @param x a [kaleido()] object.
#' @param ... currently unused.
#' @export
#' @importFrom utils capture.output
#' @keywords internal
print.kaleidoScope <- function(x, ...) {
  args <- formals(x$transform)
  cat("$transform: function(", paste(names(args), collapse = ", "), ")\n", sep = "")
  cat("$shutdown: function()\n")
  cat("$scope: ", utils::capture.output(x$scope))
}
