#' Read a Zarr array
#'
#' @param zarr_array_path Path to a Zarr array. A character vector of length 1.
#'   This can either be a location on a local file system or the URI to an array
#'   in S3 storage.
#' @param index A list of the same length as the number of dimensions in the
#'   Zarr array.  Each entry in the list provides the indices in that dimension
#'   that should be read from the array.  Setting a list entry to `NULL` will
#'   read everything in the associated dimension.  If this argument is missing
#'   the entirety of the the Zarr array will be read.
#' @param s3_client Object created by [paws.storage::s3()]. Only required for a
#'   file on S3. Leave as `NULL` for a file on local storage.
#'
#' @returns An array with the same number of dimensions as the input array. The
#'   extent of each dimension will correspond to the length of the values
#'   provided to the `index` argument.
#'
#' @examples
#'
#' ## Using a local file provided with the package
#' ## This array has 3 dimensions
#' z1 <- system.file("extdata", "zarr_examples", "row-first", "int32.zarr", package = "Rarr")
#'
#' ## read the entire array
#' read_zarr_array(zarr_array_path = z1)
#'
#' ## extract values for first 10 rows, all columns, first slice
#' read_zarr_array(zarr_array_path = z1, index = list(1:10, NULL, 1))
#'
#' \donttest{
#' ## using a Zarr file hosted on Amazon S3
#' ## This array has a single dimension with length 2729077
#' z2 <- "https://noaa-nwm-retro-v2-zarr-pds.s3.amazonaws.com/feature_id/"
#'
#' ## read the entire array
#' read_zarr_array(zarr_array_path = z2)
#'
#' ## read alternating elements
#' read_zarr_array(zarr_array_path = z2, index = list(seq(1, 576, 2)))
#' }
#'
#' @export
read_zarr_array <- function(zarr_array_path, index, s3_client = NULL) {
  zarr_array_path <- .normalize_array_path(zarr_array_path)
  ## determine if this is a local or S3 array
  s3_client <- s3_client %||% .create_s3_client(path = zarr_array_path)

  metadata <- .read_array_metadata(
    zarr_array_path,
    s3_client = s3_client
  )

  if (metadata$node_type == "group") {
    stop(
      "The provided path points to a Zarr group, but `read_zarr_array()` can ",
      "only read arrays. Please provide the path to an array within the group.",
      call. = FALSE
    )
  }

  metadata$configured_decoders <- .configure_codecs(
    metadata$codecs,
    operation = "decode"
  )

  ## if no index provided we will return everything
  if (missing(index)) {
    index <- vector(mode = "list", length = length(metadata$shape))
  }
  index <- check_index(index = index, metadata = metadata)

  res <- read_data(zarr_array_path, s3_client, index, metadata)

  if (!is.null(metadata$dimension_names)) {
    dimnames(res) <- setNames(
      vector("list", length = length(dim(res))),
      metadata$dimension_names
    )
  }

  return(res)
}


read_data <- function(
  zarr_array_path,
  s3_client,
  index,
  metadata
) {
  if (!is.null(s3_client)) {
    parsed_s3_url <- parse_s3_path(zarr_array_path)
    bucket <- parsed_s3_url$bucket
    root <- parsed_s3_url$object
  } else {
    bucket <- NULL
    root <- zarr_array_path
  }

  ## precompute, for each chunk, the positions in `index` that belong to it
  chunk_dim <- unlist(metadata$chunk_grid$configuration$chunk_shape)
  chunk_positions <- .chunk_positions_by_chunk(
    index,
    metadata,
    chunk_dim
  )

  chunk_names <- names(chunk_positions)
  # In the DelayedArray framework, we can have integer(0) indices
  # https://github.com/Huber-group-EMBL/Rarr/issues/112
  chunk_paths <- paste0(root, chunk_names, recycle0 = TRUE)

  ## Vectorized check for chunk existence
  chunk_exists <- .store_check_exist(zarr_array_path, chunk_names, s3_client)
  existing_idx <- which(chunk_exists)

  warnings <- list()
  ## hopefully we can eventually do this in parallel
  chunk_selections <- withCallingHandlers(
    lapply(
      # We skip missing chunks here since they will just be filled with the fill value
      # when initializing the consolidated array.
      existing_idx,
      function(i) {
        .extract_elements(
          chunk_name = chunk_names[i],
          current_chunk_path = chunk_paths[i],
          metadata = metadata,
          chunk_dim = chunk_dim,
          s3_client = s3_client,
          s3_bucket = bucket,
          chunk_positions = chunk_positions
        )
      }
    ),
    warning = function(w) {
      warnings <<- c(warnings, list(w)) # nolint: undesirable_operator_linter.
      invokeRestart("muffleWarning")
    }
  )
  for (w in unique(warnings)) {
    warning(w)
  }

  ## predefine our array to be populated from the read chunks
  output <- array(metadata$fill_value, dim = lengths(index))

  ## proceed in serial and update the output with each chunk selection in turn
  for (i in seq_along(chunk_selections)) {
    index_in_result <- chunk_selections[[i]][[2L]]
    rlang::inject(output[!!!index_in_result] <- chunk_selections[[i]][[1L]]) # nolint: implicit_assignment_linter.
    if (
      is.list(metadata$data_type) &&
        metadata$data_type$name %in% c("struct", "structured")
    ) {
      # Assigning a list drops the dim attribute so we have to continuously add it again
      dim(output) <- lengths(index)
    }
  }
  return(output)
}

.extract_elements <- function(
  chunk_name,
  current_chunk_path,
  metadata,
  chunk_dim,
  s3_client,
  s3_bucket,
  chunk_positions
) {
  ## find elements to select from the chunk and what in the output we replace
  chunk_info <- chunk_positions[[chunk_name]]
  index_in_result <- chunk_info$positions
  index_in_chunk <- chunk_info$index_in_chunk

  # When we get here, we know the chunk exists, so we can read it without worrying about
  # handling missing.
  if (nzchar(Sys.getenv("RARR_DEBUG"))) {
    message(current_chunk_path)
  }

  if (is.null(s3_client)) {
    size <- file.size(current_chunk_path)
    raw_chunk <- readBin(con = current_chunk_path, what = "raw", n = size)
  } else {
    raw_chunk <- s3_client$get_object(
      Bucket = s3_bucket,
      Key = current_chunk_path
    )$Body
  }

  ## read this chunk
  chunk <- read_chunk(
    chunk_bytes = raw_chunk,
    chunk_dim = chunk_dim,
    decoders = metadata$configured_decoders,
    datatype = metadata$datatype,
    fill_value = metadata$fill_value
  )

  ## extract the required elements from the chunk
  # FIXME: optimization: skip this step if we are taking everything in the chunk
  chunk <- .extract_chunk(chunk, index_in_chunk)
  return(list(chunk, index_in_result))
}

#' Read a single Zarr chunk
#'
#' @param chunk_bytes A raw vector containing the bytes of the chunk to be read.
#' @param chunk_dim The dimensions of the chunk to be read. A numeric vector.
#' @param decoders A list of configured decoders for the array.
#' @param datatype A list describing the datatype of the array.
#'
#' @returns An array containing the decompressed chunk values.
#'
#' @keywords internal
#' @noRd
read_chunk <- function(
  chunk_bytes,
  chunk_dim,
  decoders,
  datatype,
  fill_value
) {
  # Bytes -> Bytes codecs
  for (codec in decoders[["bytes_bytes"]]) {
    chunk_bytes <- codec(
      bytes = chunk_bytes
    )
  }

  # Bytes -> Array codecs
  for (codec in decoders[["array_bytes"]]) {
    converted_chunk <- codec(
      chunk_bytes,
      chunk_dim,
      datatype,
      fill_value
    )
  }
  # Array -> Array codecs
  for (codec in decoders[["array_array"]]) {
    converted_chunk <- do.call(codec, list(converted_chunk))
  }

  return(converted_chunk)
}
