// SPDX-License-Identifier: Apache-2.0

// SPDX-FileCopyrightText: Copyright the Vortex contributors

#include <nanoarrow/nanoarrow.h>

#include <catch2/catch_test_macros.hpp>

#include <nanoarrow/nanoarrow.hpp>

typedef struct ArrowSchema FFI_ArrowSchema;

typedef struct ArrowArray FFI_ArrowArray;

typedef struct ArrowArrayStream FFI_ArrowArrayStream;

#define USE_OWN_ARROW 1

#include <string>

#include <string_view>

#include <vector>

#include <vortex/data_source.hpp>

#include "common.hpp"

using namespace vortex;

using namespace std::string_view_literals;

using vortex_test::TempPath;

namespace {

Array strings_from_arrow(std::span<const std::string_view> values, bool with_null) {

nanoarrow::UniqueSchema schema;

REQUIRE(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRING) == NANOARROW_OK);

nanoarrow::UniqueArray arr;

REQUIRE(ArrowArrayInitFromSchema(arr.get(), schema.get(), nullptr) == NANOARROW_OK);

REQUIRE(ArrowArrayStartAppending(arr.get()) == NANOARROW_OK);

for (const auto &value : values) {

const ArrowStringView view {value.data(), static_cast<int64_t>(value.size())};

REQUIRE(ArrowArrayAppendString(arr.get(), view) == NANOARROW_OK);

}

if (with_null) {

REQUIRE(ArrowArrayAppendNull(arr.get(), 1) == NANOARROW_OK);

}

REQUIRE(ArrowArrayFinishBuildingDefault(arr.get(), nullptr) == NANOARROW_OK);

ArrowArray raw_arr = {};

ArrowSchema raw_schema = {};

ArrowArrayMove(arr.get(), &raw_arr);

ArrowSchemaMove(schema.get(), &raw_schema);

return Array::from_arrow(&raw_arr, &raw_schema, true);

}

Array bytes_from_arrow(std::span<const std::string_view> values) {

nanoarrow::UniqueSchema schema;

REQUIRE(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_BINARY) == NANOARROW_OK);

nanoarrow::UniqueArray arr;

REQUIRE(ArrowArrayInitFromSchema(arr.get(), schema.get(), nullptr) == NANOARROW_OK);

REQUIRE(ArrowArrayStartAppending(arr.get()) == NANOARROW_OK);

for (const auto &value : values) {

const ArrowBufferView view {{reinterpret_cast<const uint8_t *>(value.data())},

static_cast<int64_t>(value.size())};

REQUIRE(ArrowArrayAppendBytes(arr.get(), view) == NANOARROW_OK);

}

REQUIRE(ArrowArrayFinishBuildingDefault(arr.get(), nullptr) == NANOARROW_OK);

ArrowArray raw_arr = {};

ArrowSchema raw_schema = {};

ArrowArrayMove(arr.get(), &raw_arr);

ArrowSchemaMove(schema.get(), &raw_schema);

return Array::from_arrow(&raw_arr, &raw_schema, true);

}

TEST_CASE("String view over utf8 array", "[strings]") {

Session session;

const std::string long1(40, 'x');

const std::vector<std::string_view> values = {"short"sv, "Широкая строка"sv, long1, ""sv};

Array array = strings_from_arrow(values, true);

StringView view = array.strings(session);

REQUIRE(view.size() == values.size() + 1);

for (size_t i = 0; i < values.size(); ++i) {

REQUIRE_FALSE(view.is_null(i));

REQUIRE(view[i] == values[i]);

}

REQUIRE(view.is_null(values.size()));

REQUIRE_THROWS_AS(view[view.size()], VortexException);

}

TEST_CASE("Bytes view over binary array", "[strings]") {

Session session;

// Includes a NUL byte

const std::vector<std::string_view> values = {std::string_view {"abc\0def", 7}, "ffff"sv};

Array array = bytes_from_arrow(values);

BytesView view = array.bytes(session);

REQUIRE(view.size() == 2);

for (size_t i = 0; i < values.size(); ++i) {

const auto bytes = view[i];

REQUIRE(std::string_view(reinterpret_cast<const char *>(bytes.data()), bytes.size()) == values[i]);

}

}

TEST_CASE("Strings roundtrip", "[strings]") {

Session session;

TempPath path = TempPath::unique();

const std::string long1(64, 'y');

const std::vector<std::string_view> values = {"inlined"sv, long1};

Array strings = strings_from_arrow(values, false);

Writer writer =

Writer::open(session, path.string(), dtype::struct_({{"s", dtype::utf8(dtype::Nullable)}}));

writer.push(make_struct({{"s", strings}}));

writer.finish();

DataSource ds = DataSource::open(session, {path.string()});

Scan scan = ds.scan();

auto partition = scan.next_partition();

REQUIRE(partition.has_value());

auto batch = partition->next();

REQUIRE(batch.has_value());

StringView view = batch->field(0).strings(session);

REQUIRE(view.size() == 2);

REQUIRE(view[0] == "inlined"sv);

REQUIRE(view[1] == long1);

}

TEST_CASE("strings() on a non-utf8 array", "[strings]") {

Session session;

std::vector<int32_t> data = {1};

Array a = Array::primitive<int32_t>(data);

REQUIRE_THROWS_AS(a.strings(session), VortexException);

REQUIRE_THROWS_AS(a.bytes(session), VortexException);

}

} // namespace