# 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("convert array stream works", {
  stream0 <- basic_array_stream(list(), schema = na_struct(list(x = na_int32())))
  expect_identical(convert_array_stream(stream0), data.frame(x = integer()))

  stream1 <- basic_array_stream(list(data.frame(x = 1:5)))
  expect_identical(convert_array_stream(stream1), data.frame(x = 1:5))

  stream2 <- basic_array_stream(
    list(
      data.frame(x = 1:5),
      data.frame(x = 6:10)
    )
  )
  expect_identical(convert_array_stream(stream2), data.frame(x = 1:10))

  stream3 <- basic_array_stream(list(), schema = na_int32())
  expect_identical(convert_array_stream(stream3), integer())
})

test_that("convert array stream with explicit size works", {
  stream0 <- basic_array_stream(list(), schema = na_struct(list(x = na_int32())))
  expect_identical(
    convert_array_stream(stream0, size = 0),
    data.frame(x = integer())
  )

  stream1 <- basic_array_stream(list(data.frame(x = 1:5)))
  expect_identical(
    convert_array_stream(stream1, size = 5),
    data.frame(x = 1:5)
  )

  stream2 <- basic_array_stream(
    list(
      data.frame(x = 1:5),
      data.frame(x = 6:10)
    )
  )
  expect_identical(
    convert_array_stream(stream2, size = 10),
    data.frame(x = 1:10)
  )
})

test_that("convert array stream with functional ptype works", {
  skip_if_not_installed("tibble")

  tibble_or_bust <- function(array, ptype) {
    if (is.data.frame(ptype)) {
      ptype <- tibble::as_tibble(ptype)
      ptype[] <- Map(tibble_or_bust, list(NULL), ptype)
    }

    ptype
  }

  df_nested_df <- as.data.frame(
    tibble::tibble(a = 1L, b = "two", c = data.frame(a = 3))
  )
  stream_nested <- as_nanoarrow_array_stream(df_nested_df)
  expect_identical(
    convert_array_stream(stream_nested, tibble_or_bust),
    tibble::tibble(a = 1L, b = "two", c = tibble::tibble(a = 3))
  )
})

test_that("convert array stream works for nested data.frames", {
  skip_if_not_installed("tibble")

  tbl_nested_df <- tibble::tibble(a = 1L, b = "two", c = data.frame(a = 3))

  stream_nested <- as_nanoarrow_array_stream(tbl_nested_df)
  expect_identical(
    convert_array_stream(stream_nested, tbl_nested_df),
    tbl_nested_df
  )

  stream_nested <- as_nanoarrow_array_stream(tbl_nested_df)
  expect_identical(
    convert_array_stream(stream_nested, size = 1L),
    as.data.frame(tbl_nested_df)
  )

  stream_nested <- as_nanoarrow_array_stream(tbl_nested_df)
  expect_identical(
    convert_array_stream(stream_nested, tbl_nested_df, size = 1L),
    tbl_nested_df
  )
})

test_that("convert array stream works for struct-style vectors", {
  raw_posixlt <- as.data.frame(
    unclass(as.POSIXlt("2021-01-01", tz = "America/Halifax")),
    stringsAsFactors = FALSE
  )

  stream <- as_nanoarrow_array_stream(raw_posixlt)
  expect_identical(
    convert_array_stream(stream),
    raw_posixlt
  )

  stream <- as_nanoarrow_array_stream(raw_posixlt)
  expect_identical(
    convert_array_stream(stream, as.POSIXlt("2021-01-01", tz = "America/Halifax")),
    as.POSIXlt("2021-01-01", tz = "America/Halifax")
  )

  # Check with fixed size since this takes a different code path
  stream <- as_nanoarrow_array_stream(raw_posixlt)
  expect_identical(
    convert_array_stream(stream, size = 1L),
    raw_posixlt
  )

  stream <- as_nanoarrow_array_stream(raw_posixlt)
  expect_identical(
    convert_array_stream(
      stream,
      as.POSIXlt("2021-01-01", tz = "America/Halifax"),
      size = 1
    ),
    as.POSIXlt("2021-01-01", tz = "America/Halifax")
  )
})

test_that("convert array stream works for fixed_size_list_of() -> matrix()", {
  mat <- matrix(1:6, ncol = 2, byrow = TRUE)
  array <- as_nanoarrow_array(mat)
  stream <- basic_array_stream(list(array, array))
  expect_identical(
    convert_array_stream(stream, matrix(integer(), ncol = 2)),
    rbind(mat, mat)
  )

  # Check with non-default ptype
  stream <- basic_array_stream(list(array, array))
  expect_identical(
    convert_array_stream(stream, matrix(double(), ncol = 2)),
    rbind(
      matrix(as.double(1:6), ncol = 2, byrow = TRUE),
      matrix(as.double(1:6), ncol = 2, byrow = TRUE)
    )
  )
})

test_that("convert array stream works for empty fixed_size_list_of() -> matrix()", {
  stream <- basic_array_stream(list(), schema = na_fixed_size_list(na_int32(), 2))
  expect_identical(
    convert_array_stream(stream, matrix(integer(), ncol = 2)),
    matrix(integer(), ncol = 2)
  )
})

test_that("convert array stream works for nested fixed_size_list_of() -> matrix()", {
  df <- data.frame(x = 1:3)
  df$mat <- matrix(1:6, ncol = 2, byrow = TRUE)

  expected <- df[c(1:3, 1:3),]
  row.names(expected) <- 1:6

  array <- as_nanoarrow_array(df)
  stream <- basic_array_stream(list(array, array))
  expect_identical(
    convert_array_stream(stream, df[integer(0),]),
    expected
  )
})

test_that("convert array stream works for fixed_size_list_of() with non-zero offsets -> matrix() ", {
  mat <- matrix(1:6, ncol = 2, byrow = TRUE)

  # Non-zero parent offset
  array <- as_nanoarrow_array(mat)
  array <- nanoarrow_array_modify(array, list(offset = 1, length = 2))
  stream <- basic_array_stream(list(array, array))
  expect_identical(
    convert_array_stream(stream, matrix(integer(), ncol = 2)),
    rbind(mat[2:3, ], mat[2:3, ])
  )


  # Non-zero child offset
  array <- as_nanoarrow_array(mat)
  array <- nanoarrow_array_modify(
    array,
    list(
      length = 2,
      children = list(
        nanoarrow_array_modify(array$children[[1]], list(offset = 2, length = 4))
      )
    )
  )
  stream <- basic_array_stream(list(array, array))
  expect_identical(
    convert_array_stream(stream, matrix(integer(), ncol = 2)),
    rbind(mat[2:3, ], mat[2:3, ])
  )

  # Non-zero child offset and non-zero parent offset
  array <- nanoarrow_array_modify(
    array,
    list(
      offset = 1,
      length = 1,
      children = list(
        nanoarrow_array_modify(array$children[[1]], list(offset = 2, length = 4))
      )
    )
  )
  stream <- basic_array_stream(list(array, array))
  expect_identical(
    convert_array_stream(stream, matrix(integer(), ncol = 2)),
    rbind(mat[3, ], mat[3, ])
  )
})

test_that("convert array stream works for fixed_size_list_of() with parent nulls -> matrix()", {
  mat <- matrix(1:6, ncol = 2, byrow = TRUE)
  array <- as_nanoarrow_array(mat)
  array <- nanoarrow_array_modify(
    array,
    list(
      null_count = 1,
      buffers = list(
        as_nanoarrow_array(c(TRUE, TRUE, FALSE))$buffers[[2]]
      )
    )
  )

  stream <- basic_array_stream(list(array, array))
  expect_identical(
    convert_array_stream(stream, matrix(integer(), ncol = 2)),
    rbind(mat[c(1, 2, NA), ], mat[c(1, 2, NA), ])
  )
})

test_that("complex nested struct/map combination can be read", {
  skip_if_not_installed("arrow")
  skip_if_not_installed("vctrs")

  complex_map_file <- system.file("test-data/complex-map.arrows", package = "nanoarrow")
  df_from_arrow <- as.data.frame(arrow::read_ipc_stream(complex_map_file))
  df_from_nanoarrow <- as.data.frame(read_nanoarrow(complex_map_file))

  expect_identical(
    df_from_nanoarrow$idx,
    as.double(df_from_arrow$idx)
  )

  expect_identical(
    df_from_nanoarrow$names$primary,
    df_from_arrow$names$primary
  )

  expect_identical(
    df_from_nanoarrow$names$common[[3000]],
    as.data.frame(df_from_arrow$names$common[[3000]])
  )

  expect_identical(
    lengths(df_from_nanoarrow$names$rules),
    lengths(df_from_arrow$names$rules)
  )

  expect_identical(
    df_from_nanoarrow$names$rules[[4238]][c("variant", "language", "value")],
    as.data.frame(df_from_arrow$names$rules[[4238]][c("variant", "language", "value")])
  )
})

test_that("convert array stream respects the value of n", {
  batches <- list(
    data.frame(x = 1:5),
    data.frame(x = 6:10),
    data.frame(x = 11:15)
  )

  stream3 <- basic_array_stream(batches)
  expect_identical(
    convert_array_stream(stream3, n = 0),
    data.frame(x = integer())
  )

  stream3 <- basic_array_stream(batches)
  expect_identical(
    convert_array_stream(stream3, n = 1),
    data.frame(x = 1:5)
  )

  stream3 <- basic_array_stream(batches)
  expect_identical(
    convert_array_stream(stream3, n = 2),
    data.frame(x = 1:10)
  )
})

test_that("fixed-size convert array stream respects the value of n", {
  batches <- list(
    data.frame(x = 1:5),
    data.frame(x = 6:10),
    data.frame(x = 11:15)
  )

  stream3 <- basic_array_stream(batches)
  expect_identical(
    convert_array_stream(stream3, n = 0, size = 0),
    data.frame(x = integer())
  )

  stream3 <- basic_array_stream(batches)
  expect_identical(
    convert_array_stream(stream3, n = 1, size = 5),
    data.frame(x = 1:5)
  )

  stream3 <- basic_array_stream(batches)
  expect_identical(
    convert_array_stream(stream3, n = 2, size = 10),
    data.frame(x = 1:10)
  )
})

test_that("fixed-size stream conversion errors when the output has insufficient size", {
  stream <- as_nanoarrow_array_stream(data.frame(x = 1:100))
  expect_error(
    convert_array_stream(stream, size = 2),
    "Expected to materialize 100 values in batch 1 but materialized 2"
  )
})

test_that("convert array stream 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
    ),
  )

  # Split into multiple batches (first 13 rows and last 13 rows)
  batch_with_nested_dictionary <- arrow::as_record_batch(na_batch)
  batch1 <- batch_with_nested_dictionary$Take(0:12)
  batch2 <- batch_with_nested_dictionary$Take(13:25)

  # Create a stream with multiple batches
  stream <- basic_array_stream(
    list(
      as_nanoarrow_array(batch1),
      as_nanoarrow_array(batch2)
    )
  )

  # Ensure the stream converts back to the original
  expect_identical(
    convert_array_stream(stream),
    as.data.frame(df)
  )

  # Also test with reversed batches
  reversed_batch1 <- batch_with_nested_dictionary$Take(12:0)
  reversed_batch2 <- batch_with_nested_dictionary$Take(25:13)

  reversed_stream <- basic_array_stream(
    list(
      as_nanoarrow_array(reversed_batch1),
      as_nanoarrow_array(reversed_batch2)
    )
  )

  expected_reversed <- rbind(
    as.data.frame(df[13:1, ]),
    as.data.frame(df[26:14, ])
  )
  row.names(expected_reversed) <- 1:26

  expect_identical(
    convert_array_stream(reversed_stream),
    expected_reversed
  )
})
