#include "arrow/array.h"

#include <cstring>
#include <string>
#include <utility>
#include <vector>

namespace arrow_c {
namespace {

// Producer-owned storage for one schema node. The struct's `format`, `name` and
// `children` pointers point in here, which is why it outlives the call that
// built it and is freed by the release callback rather than by the builder.
struct SchemaData {
    std::string format;
    std::string name;
    std::vector<ArrowSchema*> children;
};

// Producer-owned storage for one array node. `owned` holds the bytes this
// node's own buffers point into - offset arrays, and scalars that have no
// storage elsewhere. The payload itself is borrowed from the caller.
struct ArrayData {
    std::vector<uint8_t> owned;
    std::vector<const void*> buffers;
    std::vector<ArrowArray*> children;
};

// A node's release callback releases its children before freeing itself, so a
// consumer that only releases the root - which is all the interface requires -
// still frees the whole tree.
void release_schema(ArrowSchema* schema) {
    auto* data = static_cast<SchemaData*>(schema->private_data);

    for (ArrowSchema* child : data->children) {
        if (child->release != nullptr) {
            child->release(child);
        }
        delete child;
    }

    delete data;
    schema->release = nullptr;
}

void release_array(ArrowArray* array) {
    auto* data = static_cast<ArrayData*>(array->private_data);

    for (ArrowArray* child : data->children) {
        if (child->release != nullptr) {
            child->release(child);
        }
        delete child;
    }

    delete data;
    array->release = nullptr;
}

void set_schema(ArrowSchema& schema, const char* format, const char* name,
                std::vector<ArrowSchema*> children) {
    auto* data = new SchemaData{format, name, std::move(children)};

    schema.format = data->format.c_str();
    schema.name = data->name.c_str();
    schema.metadata = nullptr;
    // Nullable everywhere, which is what pyarrow's defaults produce: the same
    // struct built by a Python node has to compare equal to this one.
    schema.flags = ARROW_FLAG_NULLABLE;
    schema.n_children = static_cast<int64_t>(data->children.size());
    schema.children = data->children.empty() ? nullptr : data->children.data();
    schema.dictionary = nullptr;
    schema.release = release_schema;
    schema.private_data = data;
}

// `buffers` may point into `owned`: moving a vector keeps its heap block, so
// pointers taken before the move stay valid after it.
void set_array(ArrowArray& array, int64_t length, std::vector<uint8_t> owned,
               std::vector<const void*> buffers, std::vector<ArrowArray*> children) {
    auto* data = new ArrayData{std::move(owned), std::move(buffers), std::move(children)};

    array.length = length;
    // Nothing this node sends is null, so the validity bitmap - the first
    // buffer of every node below - is left out entirely.
    array.null_count = 0;
    array.offset = 0;
    array.n_buffers = static_cast<int64_t>(data->buffers.size());
    array.n_children = static_cast<int64_t>(data->children.size());
    array.buffers = data->buffers.data();
    array.children = data->children.empty() ? nullptr : data->children.data();
    array.dictionary = nullptr;
    array.release = release_array;
    array.private_data = data;
}

// One field of a struct, owned by its parent.
struct Child {
    ArrowSchema* schema;
    ArrowArray* array;
};

Child new_child() {
    return {new ArrowSchema{}, new ArrowArray{}};
}

void fill_uint32(ArrowSchema& schema, ArrowArray& array, const char* name, uint32_t value) {
    std::vector<uint8_t> owned(sizeof(value));
    std::memcpy(owned.data(), &value, sizeof(value));
    const void* values = owned.data();

    set_schema(schema, "I", name, {});
    set_array(array, 1, std::move(owned), {nullptr, values}, {});
}

// One string, offsets and characters in a single allocation.
void fill_utf8(ArrowSchema& schema, ArrowArray& array, const char* name, const char* text,
               size_t size) {
    const int32_t offsets[2] = {0, static_cast<int32_t>(size)};

    std::vector<uint8_t> owned(sizeof(offsets) + size);
    std::memcpy(owned.data(), offsets, sizeof(offsets));
    std::memcpy(owned.data() + sizeof(offsets), text, size);
    const void* offsets_buffer = owned.data();
    const void* values_buffer = owned.data() + sizeof(offsets);

    set_schema(schema, "u", name, {});
    set_array(array, 1, std::move(owned), {nullptr, offsets_buffer, values_buffer}, {});
}

// One byte string. Only the offsets are owned; the bytes are the caller's.
void fill_large_binary(ArrowSchema& schema, ArrowArray& array, const char* name,
                       const uint8_t* data, size_t size) {
    const int64_t offsets[2] = {0, static_cast<int64_t>(size)};

    std::vector<uint8_t> owned(sizeof(offsets));
    std::memcpy(owned.data(), offsets, sizeof(offsets));
    const void* offsets_buffer = owned.data();

    set_schema(schema, "Z", name, {});
    set_array(array, 1, std::move(owned), {nullptr, offsets_buffer, data}, {});
}

} // namespace

Exported image(uint32_t width, uint32_t height, const char* encoding, const uint8_t* data,
               size_t size) {
    Child w = new_child();
    fill_uint32(*w.schema, *w.array, "width", width);

    Child h = new_child();
    fill_uint32(*h.schema, *h.array, "height", height);

    Child e = new_child();
    fill_utf8(*e.schema, *e.array, "encoding", encoding, std::strlen(encoding));

    Child d = new_child();
    fill_large_binary(*d.schema, *d.array, "data", data, size);

    Exported out;
    set_schema(out.schema, "+s", "", {w.schema, h.schema, e.schema, d.schema});
    set_array(out.array, 1, {}, {nullptr}, {w.array, h.array, e.array, d.array});
    return out;
}

Exported float32(const float* values, size_t count) {
    Exported out;
    set_schema(out.schema, "f", "", {});
    set_array(out.array, static_cast<int64_t>(count), {}, {nullptr, values}, {});
    return out;
}

Exported string(const char* text, size_t size) {
    Exported out;
    fill_utf8(out.schema, out.array, "", text, size);
    return out;
}

} // namespace arrow_c