# 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.

test_that("read_nanoarrow() works for raw vectors", {
  stream <- read_nanoarrow(example_ipc_stream())
  expect_s3_class(stream, "nanoarrow_array_stream")
  expect_identical(
    as.data.frame(stream),
    data.frame(some_col = c(0L, 1L, 2L))
  )
})

test_that("read_nanoarrow() works for open connections", {
  con <- rawConnection(example_ipc_stream())
  on.exit(close(con))

  stream <- read_nanoarrow(con)
  expect_s3_class(stream, "nanoarrow_array_stream")
  expect_identical(
    as.data.frame(stream),
    data.frame(some_col = c(0L, 1L, 2L))
  )
})

test_that("write_nanoarrow() works for open connections", {
  tf <- tempfile()
  on.exit(unlink(tf))

  con <- rawConnection(raw(), "wb")
  on.exit(close(con))

  write_nanoarrow(data.frame(), con)
  expect_identical(
    as.data.frame(read_nanoarrow(rawConnectionValue(con))),
    data.frame()
  )
})

test_that("read_nanoarrow() works for unopened connections", {
  tf <- tempfile()
  on.exit(unlink(tf))

  con <- file(tf, "wb")
  writeBin(example_ipc_stream(), con)
  close(con)

  con <- file(tf)
  # Don't close on exit, because we're supposed to do that

  stream <- read_nanoarrow(con)
  expect_true(isOpen(con))
  stream$release()
  expect_error(
    close(con),
    "invalid connection"
  )
})

test_that("write_nanoarrow() works for unopened connections", {
  tf <- tempfile()
  on.exit(unlink(tf))

  con <- file(tf)
  # Don't close on exit, because we're supposed to do that

  write_nanoarrow(data.frame(), con)
  expect_error(
    close(con),
    "invalid connection"
  )
})

test_that("read_nanoarrow() works for file paths", {
  tf <- tempfile()
  on.exit(unlink(tf))

  con <- file(tf, "wb")
  writeBin(example_ipc_stream(), con)
  close(con)

  stream <- read_nanoarrow(tf)
  expect_identical(
    as.data.frame(stream),
    data.frame(some_col = c(0L, 1L, 2L))
  )
})

test_that("write_nanoarrow() works for file paths", {
  tf <- tempfile()
  on.exit(unlink(tf))

  df <- data.frame(letters = letters, stringsAsFactors = FALSE)
  expect_identical(write_nanoarrow(df, tf), df)
  expect_identical(as.data.frame(read_nanoarrow(tf)), df)
})

test_that("read_nanoarrow() works for URLs", {
  tf <- tempfile()
  on.exit(unlink(tf))

  con <- file(tf, "wb")
  writeBin(example_ipc_stream(), con)
  close(con)

  stream <- read_nanoarrow(paste0("file://", tf))
  expect_identical(
    as.data.frame(stream),
    data.frame(some_col = c(0L, 1L, 2L))
  )
})

test_that("read_nanoarrow() works with zstd buffer compression", {
  skip_if_not(nanoarrow_with_zstd())

  stream <- read_nanoarrow(example_ipc_stream(compression = "zstd"))
  expect_s3_class(stream, "nanoarrow_array_stream")
  expect_identical(
    as.data.frame(stream),
    data.frame(some_col = c(0L, 1L, 2L))
  )
})

test_that("read_nanoarrow() works with lz4 buffer compression", {
  skip_if_not(nanoarrow_with_lz4())

  stream <- read_nanoarrow(example_ipc_stream(compression = "lz4"))
  expect_s3_class(stream, "nanoarrow_array_stream")
  expect_identical(
    as.data.frame(stream),
    data.frame(some_col = c(0L, 1L, 2L))
  )
})

test_that("write_nanoarrow() works for URLs", {
  tf <- tempfile()
  on.exit(unlink(tf))

  df <- data.frame(letters = letters, stringsAsFactors = FALSE)
  expect_identical(write_nanoarrow(df, paste0("file://", tf)), df)
  expect_identical(as.data.frame(read_nanoarrow(tf)), df)
})

test_that("read_nanoarrow() works for compressed .gz file paths", {
  tf <- tempfile(fileext = ".gz")
  on.exit(unlink(tf))

  con <- gzfile(tf, "wb")
  writeBin(example_ipc_stream(), con)
  close(con)

  stream <- read_nanoarrow(tf)
  expect_identical(
    as.data.frame(stream),
    data.frame(some_col = c(0L, 1L, 2L))
  )
})

test_that("write_nanoarrow() works for compressed .gz file paths", {
  tf <- tempfile(fileext = ".gz")
  on.exit(unlink(tf))

  df <- data.frame(letters = letters, stringsAsFactors = FALSE)
  expect_identical(write_nanoarrow(df, tf), df)
  expect_identical(as.data.frame(read_nanoarrow(tf)), df)
})

test_that("read_nanoarrow() works for compressed .bz2 file paths", {
  tf <- tempfile(fileext = ".bz2")
  on.exit(unlink(tf))

  con <- bzfile(tf, "wb")
  writeBin(example_ipc_stream(), con)
  close(con)

  stream <- read_nanoarrow(tf)
  expect_identical(
    as.data.frame(stream),
    data.frame(some_col = c(0L, 1L, 2L))
  )
})

test_that("write_nanoarrow() works for compressed .bz2 file paths", {
  tf <- tempfile(fileext = ".bz2")
  on.exit(unlink(tf))

  df <- data.frame(letters = letters, stringsAsFactors = FALSE)
  expect_identical(write_nanoarrow(df, tf), df)
  expect_identical(as.data.frame(read_nanoarrow(tf)), df)
})

test_that("read_nanoarrow() works for compressed .zip file paths", {
  tf <- tempfile(fileext = ".zip")
  tdir <- tempfile()
  on.exit(unlink(c(tf, tdir), recursive = TRUE))

  dir.create(tdir)
  uncompressed <- file.path(tdir, "file.arrows")
  con <- file(uncompressed, "wb")
  writeBin(example_ipc_stream(), con)
  close(con)

  local({
    wd <- getwd()
    on.exit(setwd(wd))
    setwd(tdir)
    tryCatch(
      zip(tf, "file.arrows", extras = "-q"),
      error = function(...) skip("zip() not supported")
    )
  })

  stream <- read_nanoarrow(tf)
  expect_identical(
    as.data.frame(stream),
    data.frame(some_col = c(0L, 1L, 2L))
  )
})

test_that("write_nanoarrow() errors for compressed .zip file paths", {
  tf <- tempfile(fileext = ".zip")
  on.exit(unlink(tf))

  df <- data.frame(letters = letters, stringsAsFactors = FALSE)
  expect_error(
    write_nanoarrow(df, tf),
    "zip compression not supported"
  )
})

test_that("read_nanoarrow() errors for compressed URL paths", {
  expect_error(
    read_nanoarrow("https://something.zip"),
    "Reading compressed streams from URLs"
  )
})

test_that("read|write_nanoarrow() errors for input with length != 1", {
  expect_error(
    read_nanoarrow(character(0)),
    "Can't interpret character"
  )
  expect_error(
    write_nanoarrow(data.frame(), character(0)),
    "Can't interpret character"
  )
})

test_that("read_nanoarrow() errors zip archives that contain files != 1", {
  tf <- tempfile(fileext = ".zip")
  tdir <- tempfile()
  on.exit(unlink(c(tf, tdir), recursive = TRUE))

  dir.create(tdir)
  file.create(file.path(tdir, c("file1", "file2")))
  local({
    wd <- getwd()
    on.exit(setwd(wd))
    setwd(tdir)
    tryCatch(
      zip(tf, c("file1", "file2"), extras = "-q"),
      error = function(...) skip("zip() not supported")
    )
  })

  expect_error(
    read_nanoarrow(tf),
    "Unzip only supported of archives with exactly one file"
  )
})

test_that("read_nanoarrow() reports errors from readBin", {
  tf <- tempfile()
  on.exit(unlink(tf))
  writeLines("this is not a binary file", tf)

  con <- file(tf, open = "r")
  on.exit(close(con), add = TRUE)

  expect_error(
    read_nanoarrow(con),
    "R execution error"
  )
})

test_that("write_nanoarrow() reports errors from writeBin", {
  tf <- tempfile()
  on.exit(unlink(tf))
  con <- file(tf, open = "w")
  on.exit(close(con), add = TRUE)

  expect_error(
    write_nanoarrow(data.frame(), con),
    "R execution error"
  )
})

test_that("read_nanoarrow() respects lazy argument", {
  expect_error(
    read_nanoarrow(raw(0), lazy = FALSE),
    "No data available on stream"
  )

  reader <- read_nanoarrow(raw(0), lazy = TRUE)
  expect_error(
    reader$get_next(),
    "No data available on stream"
  )

  tf <- tempfile()
  con <- rawConnection(raw(0))
  on.exit({
    close(con)
    unlink(tf)
  })

  expect_error(
    read_nanoarrow(con, lazy = FALSE),
    "No data available on stream"
  )

  reader <- read_nanoarrow(con, lazy = TRUE)
  expect_error(
    reader$get_next(),
    "No data available on stream"
  )

  file.create(tf)
  expect_error(
    read_nanoarrow(tf, lazy = FALSE),
    "No data available on stream"
  )

  reader <- read_nanoarrow(tf, lazy = TRUE)
  expect_error(
    reader$get_next(),
    "No data available on stream"
  )
})

test_that("read_nanoarrow() from connection errors when called from another thread", {
  skip_if_not_installed("arrow")
  skip_if_not(arrow::arrow_info()$capabilities["dataset"])
  skip_if_not_installed("dplyr")

  tf <- tempfile()
  tf_out <- tempfile()
  on.exit(unlink(c(tf, tf_out), recursive = TRUE))

  con <- file(tf, "wb")
  writeBin(example_ipc_stream(), con)
  close(con)

  stream <- read_nanoarrow(tf)
  reader <- arrow::as_record_batch_reader(stream)

  # There is an internal MakeSafeRecordBatchReader that ensures all read
  # calls happen on the R thread (used in DuckDB integration), but for now
  # this should at least error and not crash.
  expect_error(
    arrow::write_dataset(reader, tf_out),
    "Can't read from R connection on a non-R thread"
  )
})

test_that("read_nanoarrow() works for dictionary arrays produced by arrow", {
  skip_if_not_installed("arrow")

  # Generate data with many small batches to stress the shared dictionary
  levels <- LETTERS
  df <- data.frame(
    idx = 1:100,
    col = factor(rep_len(rev(LETTERS), 100), levels = levels)
  )
  batch <- arrow::as_record_batch(df)

  tf <- tempfile()
  on.exit(unlink(tf))

  # Write in small batches using RecordBatchStreamWriter
  tbl <- arrow::Table$create(df)
  sink <- arrow::FileOutputStream$create(tf)

  writer <- arrow::RecordBatchStreamWriter$create(
    sink,
    batch$schema
  )

  for (i in 1:10) {
    writer$write_batch(batch)
  }

  writer$close()
  sink$close()

  # Ensure we have a dictionary type
  reader <- read_nanoarrow(tf)
  schema <- reader$get_schema()
  expect_true(batch$schema$Equals(arrow::as_schema(schema)))

  # Ensure we have the expected number of batches
  batches <- collect_array_stream(reader)
  expect_identical(length(batches), 10L)

  # Ensure all dictionaries are pointing to the same buffers (shared)
  dictionaries <- lapply(batches, function(b) b$children$col$dictionary)
  offsets_addr <- nanoarrow_pointer_addr_chr(dictionaries[[1]]$buffers[[2]]$data)
  data_addr <- nanoarrow_pointer_addr_chr(dictionaries[[1]]$buffers[[3]]$data)
  for (dictionary in dictionaries[2:10]) {
    offsets_addr <- nanoarrow_pointer_addr_chr(dictionary$buffers[[2]]$data)
    data_addr <- nanoarrow_pointer_addr_chr(dictionary$buffers[[3]]$data)
  }

  # Ensure the result is correctly read into a factor
  result <- convert_array_stream(read_nanoarrow(tf), to = df)
  expect_identical(result, do.call(rbind, rep(list(df), 10)))

  # Ensure the default read works too (by default dictionaries are expanded)
  result <- convert_array_stream(read_nanoarrow(tf))
  df$col <- as.character(df$col)
  expect_identical(result, do.call(rbind, rep(list(df), 10)))
})

test_that("read_nanoarrow() works for dictionaries of structs of dictionaries", {
  skip_if_not_installed("arrow")
  skip_if_not_installed("tibble")

  # Create a record batch with a dictionary inside a struct. Use tibble
  # for slightly easier nested data frame construction.
  df <- tibble::tibble(
    idx = 1:26,
    letter = letters,
    nested = data.frame(idx2 = 1:26, LETTER = LETTERS)
  )

  schema <- arrow::schema(
    idx = arrow::int32(),
    letter = arrow::utf8(),
    nested = arrow::struct(
      idx2 = arrow::int32(),
      LETTER = arrow::dictionary(
        index_type = arrow::int32(),
        value_type = arrow::large_utf8()
      )
    )
  )

  batch <- arrow::record_batch(df)$cast(schema)

  # Modify the batch so that the inner struct is the values member of a dict
  na_batch <- as_nanoarrow_array(batch)
  na_batch$children$nested <- nanoarrow_array_modify(
    nanoarrow_array_init(
      as_nanoarrow_schema(
        arrow::dictionary(
          index_type = arrow::int32(),
          value_type = arrow::struct(
            idx2 = arrow::int32(),
            LETTER = arrow::dictionary(
              index_type = arrow::int8(),
              value_type = arrow::large_utf8()
            )
          )
        )
      )
    ),
    list(
      length = 26,
      null_count = 0,
      buffers = list(
        NULL,
        0:25
      ),
      dictionary = na_batch$children$nested
    ),
  )

  # Create a reversed batch as well (same dictionary, different indices)
  batch_with_nested_dictionary <- arrow::as_record_batch(na_batch)
  reversed_batch_with_nested_dictionary <- batch_with_nested_dictionary$Take(
    (batch$num_rows - 1):0
  )

  # Test IPC roundtrip with multiple batches
  tf <- tempfile()
  on.exit(unlink(tf))

  # Write multiple batches using RecordBatchStreamWriter
  sink <- arrow::FileOutputStream$create(tf)
  writer <- arrow::RecordBatchStreamWriter$create(
    sink,
    batch_with_nested_dictionary$schema
  )

  # Write the forward batch and the reversed batch
  writer$write_batch(batch_with_nested_dictionary)
  writer$write_batch(reversed_batch_with_nested_dictionary)
  writer$close()
  sink$close()

  # Read back via nanoarrow and verify result
  result <- convert_array_stream(read_nanoarrow(tf))

  expected <- rbind(
    as.data.frame(df),
    as.data.frame(df[rev(seq_len(nrow(df))), ])
  )
  row.names(expected) <- seq_len(nrow(expected))

  expect_identical(result, expected)
})
