# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

#' Read/write serialized streams of Arrow data
#'
#' Reads/writes connections, file paths, URLs, or raw vectors from/to serialized
#' Arrow data. Arrow documentation typically refers to this format as "Arrow
#' IPC", since its origin was as a means to transmit tables between processes
#' (e.g., multiple R sessions). This format can also be written to and read
#' from files or URLs and is essentially a high performance equivalent of
#' a CSV file that does a better job maintaining types.
#'
#' The nanoarrow package implements an IPC writer; however, you can also
#' use [arrow::write_ipc_stream()] to write data from R, or use
#' the equivalent writer from another Arrow implementation in Python, C++,
#' Rust, JavaScript, Julia, C#, and beyond.
#'
#' The media type of an Arrow stream is `application/vnd.apache.arrow.stream`
#' and the recommended file extension is `.arrows`.
#'
#' @param x A `raw()` vector, connection, or file path from which to read
#'   binary data. Common extensions indicating compression (.gz, .bz2, .zip)
#'   are automatically uncompressed.
#' @param data An object to write as an Arrow IPC stream, converted using
#'   [as_nanoarrow_array_stream()]. Notably, this includes a [data.frame()].
#' @param lazy By default, `read_nanoarrow()` will read and discard a copy of
#'   the reader's schema to ensure that invalid streams are discovered as
#'   soon as possible. Use `lazy = TRUE` to defer this check until the reader
#'   is actually consumed.
#' @param ... Currently unused.
#'
#' @return A [nanoarrow_array_stream][as_nanoarrow_array_stream]
#' @export
#'
#' @examples
#' as.data.frame(read_nanoarrow(example_ipc_stream()))
#'
read_nanoarrow <- function(x, ..., lazy = FALSE) {
  UseMethod("read_nanoarrow")
}

#' @export
read_nanoarrow.raw <- function(x, ..., lazy = FALSE) {
  buffer <- as_nanoarrow_buffer(x)
  reader <- .Call(nanoarrow_c_ipc_array_reader_buffer, buffer)
  check_stream_if_requested(reader, lazy)
}

#' @export
read_nanoarrow.character <- function(x, ..., lazy = FALSE) {
  if (length(x) != 1) {
    stop(sprintf("Can't interpret character(%d) as file path", length(x)))
  }

  con_type <- guess_connection_type(x)
  if (con_type == "unz") {
    con <- do.call(con_type, list(x, filename = guess_zip_filename(x)))
  } else {
    con <- do.call(con_type, list(x))
  }

  # Helps with error reporting when reading invalid files
  reader <- read_nanoarrow(con, lazy = TRUE)
  check_stream_if_requested(reader, lazy)
}

#' @export
read_nanoarrow.connection <- function(x, ..., lazy = FALSE) {
  if (!isOpen(x)) {
    # Unopened connections should be opened in binary mode
    open(x, "rb")

    stream <- tryCatch(
      .Call(nanoarrow_c_ipc_array_reader_connection, x),
      error = function(e) {
        close(x)
        stop(e)
      }
    )

    # Close the connection when the array stream is released
    stream_finalizer <- function() {
      close(x)
    }

    finalizer_env <- new.env(parent = baseenv())
    finalizer_env$x <- x
    environment(stream_finalizer) <- finalizer_env

    reader <- array_stream_set_finalizer(stream, stream_finalizer)
  } else {
    reader <- .Call(nanoarrow_c_ipc_array_reader_connection, x)
  }

  check_stream_if_requested(reader, lazy)
}

#' @rdname read_nanoarrow
#' @export
write_nanoarrow <- function(data, x, ...) {
  UseMethod("write_nanoarrow", x)
}

#' @export
write_nanoarrow.connection <- function(data, x, ...) {
  if (!isOpen(x)) {
    open(x, "wb")
    on.exit(close(x))
  }

  writer <- .Call(nanoarrow_c_ipc_writer_connection, x)
  stream <- as_nanoarrow_array_stream(data)
  on.exit(nanoarrow_pointer_release(stream), add = TRUE)

  .Call(nanoarrow_c_ipc_writer_write_stream, writer, stream)
  invisible(data)
}

#' @export
write_nanoarrow.character <- function(data, x, ...) {
  if (length(x) != 1) {
    stop(sprintf("Can't interpret character(%d) as file path", length(x)))
  }

  con_type <- guess_connection_type(x)
  if (con_type == "unz") {
    stop("zip compression not supported for write_nanoarrow()")
  }

  con <- do.call(con_type, list(x))
  write_nanoarrow(data, con)
}

#' Example Arrow IPC Data
#'
#' An example stream that can be used for testing or examples.
#'
#' @param compression One of "none", "zstd", or "lz4"
#'
#' @return A raw vector that can be passed to [read_nanoarrow()]
#' @export
#'
#' @examples
#' as.data.frame(read_nanoarrow(example_ipc_stream()))
example_ipc_stream <- function(compression = c("none", "zstd", "lz4")) {
  compression <- match.arg(compression)

  # data.frame(some_col = c(0L, 1L, 2L)) as a serialized schema/batch
  schema <- as.raw(c(
    0xff, 0xff, 0xff, 0xff, 0x10, 0x01, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x0a, 0x00, 0x0e, 0x00, 0x06, 0x00, 0x05, 0x00, 0x08, 0x00, 0x0a, 0x00, 0x00, 0x00,
    0x00, 0x01, 0x04, 0x00, 0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x0c, 0x00,
    0x00, 0x00, 0x04, 0x00, 0x08, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x3c, 0x00, 0x00, 0x00,
    0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x84, 0xff,
    0xff, 0xff, 0x18, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00,
    0x73, 0x6f, 0x6d, 0x65, 0x5f, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x00, 0x00, 0x08, 0x00,
    0x00, 0x00, 0x73, 0x6f, 0x6d, 0x65, 0x5f, 0x6b, 0x65, 0x79, 0x00, 0x00, 0x00, 0x00,
    0x01, 0x00, 0x00, 0x00, 0x18, 0x00, 0x00, 0x00, 0x00, 0x00, 0x12, 0x00, 0x18, 0x00,
    0x08, 0x00, 0x06, 0x00, 0x07, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x10, 0x00, 0x14, 0x00,
    0x12, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x02, 0x14, 0x00, 0x00, 0x00, 0x70, 0x00,
    0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x18, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x08, 0x00, 0x00, 0x00, 0x73, 0x6f, 0x6d, 0x65, 0x5f, 0x63, 0x6f, 0x6c, 0x00, 0x00,
    0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x08, 0x00, 0x0c, 0x00,
    0x04, 0x00, 0x08, 0x00, 0x08, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x04, 0x00,
    0x00, 0x00, 0x10, 0x00, 0x00, 0x00, 0x73, 0x6f, 0x6d, 0x65, 0x5f, 0x76, 0x61, 0x6c,
    0x75, 0x65, 0x5f, 0x66, 0x69, 0x65, 0x6c, 0x64, 0x00, 0x00, 0x00, 0x00, 0x0e, 0x00,
    0x00, 0x00, 0x73, 0x6f, 0x6d, 0x65, 0x5f, 0x6b, 0x65, 0x79, 0x5f, 0x66, 0x69, 0x65,
    0x6c, 0x64, 0x00, 0x00, 0x08, 0x00, 0x0c, 0x00, 0x08, 0x00, 0x07, 0x00, 0x08, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
  ))

  if (identical(compression, "zstd")) {
    batch <- as.raw(c(
      0xff, 0xff, 0xff, 0xff, 0xa0, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x0c, 0x00, 0x18, 0x00, 0x06, 0x00, 0x05, 0x00, 0x08, 0x00, 0x0c, 0x00,
      0x0c, 0x00, 0x00, 0x00, 0x00, 0x03, 0x04, 0x00, 0x1c, 0x00, 0x00, 0x00, 0x20, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x1e, 0x00,
      0x10, 0x00, 0x04, 0x00, 0x08, 0x00, 0x0c, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x50, 0x00,
      0x00, 0x00, 0x24, 0x00, 0x00, 0x00, 0x18, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x06, 0x00, 0x08, 0x00,
      0x07, 0x00, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x02, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x1d, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x03, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x0c, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x28, 0xb5, 0x2f, 0xfd, 0x20, 0x0c,
      0x61, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00
    ))
  } else if (identical(compression, "lz4")) {
    batch <- as.raw(c(
      0xff, 0xff, 0xff, 0xff, 0x98, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x0c, 0x00, 0x18, 0x00, 0x06, 0x00, 0x05, 0x00, 0x08, 0x00, 0x0c, 0x00,
      0x0c, 0x00, 0x00, 0x00, 0x00, 0x03, 0x04, 0x00, 0x1c, 0x00, 0x00, 0x00, 0x28, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x1c, 0x00,
      0x10, 0x00, 0x04, 0x00, 0x08, 0x00, 0x0c, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x48, 0x00,
      0x00, 0x00, 0x1c, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0x00, 0x04, 0x00, 0x04, 0x00,
      0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x23, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x01, 0x00, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0c, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x04, 0x22, 0x4d, 0x18, 0x60, 0x40, 0x82, 0x0c, 0x00, 0x00, 0x80, 0x00, 0x00, 0x00,
      0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00
    ))
  } else {
    batch <- as.raw(c(
      0xff, 0xff, 0xff, 0xff, 0x88, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x0c, 0x00, 0x16, 0x00, 0x06, 0x00, 0x05, 0x00, 0x08, 0x00, 0x0c, 0x00,
      0x0c, 0x00, 0x00, 0x00, 0x00, 0x03, 0x04, 0x00, 0x18, 0x00, 0x00, 0x00, 0x10, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x18, 0x00, 0x0c, 0x00,
      0x04, 0x00, 0x08, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x3c, 0x00, 0x00, 0x00, 0x10, 0x00,
      0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x0c, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00,
      0x00, 0x00, 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02, 0x00,
      0x00, 0x00, 0x00, 0x00, 0x00, 0x00
    ))
  }

  c(schema, batch)
}

check_stream_if_requested <- function(reader, lazy) {
  if (!lazy) {
    # Report error as coming from read_nanoarrow() always
    cnd_call <- sys.call(-1)
    tryCatch(
      reader$get_schema(),
      error = function(e) {
        reader$release()
        e$call <- cnd_call
        stop(e)
      }
    )
  }

  reader
}

guess_connection_type <- function(x) {
  is_url <- grepl("://", x)

  compressed_con <- switch(
    tools::file_ext(x),
    "gz" = "gzfile",
    "bz2" = "bzfile",
    "zip" = "unz"
  )

  if (is_url && !is.null(compressed_con)) {
    stop("Reading compressed streams from URLs is not supported")
  }

  if (is_url) {
    "url"
  } else if (is.null(compressed_con)) {
    "file"
  } else {
    compressed_con
  }
}

guess_zip_filename <- function(x) {
  files <- utils::unzip(x, list = TRUE)[[1]]
  if (length(files) != 1) {
    stop(
      sprintf(
        "Unzip only supported of archives with exactly one file (found %d)",
        length(files)
      )
    )
  }

  files
}

# The C-level R_tryCatch() does not provide for handling interrupts (or
# I couldn't figure out how to make it work), so instead we provide wrappers
# around readBin() and writeBin() that convert interrupt conditions to errors
# (which the C code does know how to handle).
read_bin_wrapper <- function(con, what, n) {
  withCallingHandlers(
    readBin(con, what, n),
    interrupt = function(e) {
      stop("user interrupt")
    }
  )
}

write_bin_wrapper <- function(object, con) {
  withCallingHandlers(
    writeBin(object, con),
    interrupt = function(e) {
      stop("user interrupt")
    }
  )
}
