// Licensed to the LF AI & Data foundation 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.

#pragma once

#include <algorithm>
#include "common/FastMem.h"
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <exception>
#include <future>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <vector>

#include "arrow/array/builder_base.h"
#include "arrow/scalar.h"
#include "arrow/type.h"
#include "common/ArrowDataWrapper.h"
#include "common/FieldData.h"
#include "common/FieldDataInterface.h"
#include "common/FieldMeta.h"
#include "common/LoadInfo.h"
#include "common/Types.h"
#include "common/type_c.h"
#include "milvus-storage/common/metadata.h"
#include "milvus-storage/filesystem/fs.h"
#include "milvus-storage/properties.h"
#include "parquet/schema.h"
#include "storage/BinlogReader.h"
#include "storage/ChunkManager.h"
#include "storage/DataCodec.h"
#include "storage/Event.h"
#include "storage/MemFileManagerImpl.h"
#include "storage/PayloadReader.h"
#include "storage/PayloadStream.h"
#include "storage/ThreadPools.h"
#include "storage/Types.h"

namespace milvus::storage {

void
ReadMediumType(BinlogReaderPtr reader);

void
AddPayloadToArrowBuilder(std::shared_ptr<arrow::ArrayBuilder> builder,
                         const Payload& payload);

void
AddOneStringToArrowBuilder(std::shared_ptr<arrow::ArrayBuilder> builder,
                           const char* str,
                           int str_size);
void
AddOneBinaryToArrowBuilder(std::shared_ptr<arrow::ArrayBuilder> builder,
                           const uint8_t* data,
                           int length);

std::shared_ptr<arrow::ArrayBuilder>
CreateArrowBuilder(DataType data_type);

std::shared_ptr<arrow::ArrayBuilder>
CreateArrowBuilder(DataType data_type,
                   DataType element_type,
                   int dim,
                   bool nullable = false);

/// \brief Utility function to create arrow:Scalar from FieldMeta.default_value
///
/// Construct a arrow::Scalar based on input field meta
/// The data_type_ is checked to determine which `one_of` member of default value shall be used
/// Note that:
/// 1. default_value shall have value
/// 2. the type check shall be guaranteed(current by go side)
///
/// \param[in] field_meta the field meta object to construct arrow::Scalar from.
/// \return an std::shared_ptr of arrow::Scalar
std::shared_ptr<arrow::Scalar>
CreateArrowScalarFromDefaultValue(const FieldMeta& field_meta);

std::shared_ptr<arrow::Schema>
CreateArrowSchema(DataType data_type, bool nullable);

std::shared_ptr<arrow::Schema>
CreateArrowSchema(DataType data_type, int dim, bool nullable);

std::shared_ptr<arrow::Schema>
CreateArrowSchema(DataType data_type,
                  int dim,
                  DataType element_type,
                  bool nullable = false);

int
GetDimensionFromFileMetaData(const parquet::ColumnDescriptor* schema,
                             DataType data_type);

std::string
GenIndexPathIdentifier(int64_t build_id,
                       int64_t index_version,
                       int64_t segment_id,
                       int64_t field_id);

// is_temp: true for temporary path used during index building,
// false for path to store pre-built index contents downloaded from remote storage
std::string
GenIndexPathPrefixByType(ChunkManagerPtr cm,
                         int64_t build_id,
                         int64_t index_version,
                         int64_t segment_id,
                         int64_t field_id,
                         const std::string& index_type,
                         bool is_temp);

// is_temp: true for temporary path used during index building,
// false for path to store pre-built index contents downloaded from remote storage
std::string
GenIndexPathPrefix(ChunkManagerPtr cm,
                   int64_t build_id,
                   int64_t index_version,
                   int64_t segment_id,
                   int64_t field_id,
                   bool is_temp);

// is_temp: true for temporary path used during index building,
// false for path to store pre-built index contents downloaded from remote storage
std::string
GenTextIndexPathPrefix(ChunkManagerPtr cm,
                       int64_t build_id,
                       int64_t index_version,
                       int64_t segment_id,
                       int64_t field_id,
                       bool is_temp);

std::string
GenJsonStatsPathPrefix(ChunkManagerPtr cm,
                       int64_t build_id,
                       int64_t index_version,
                       int64_t segment_id,
                       int64_t field_id,
                       bool is_temp);

std::string
GenJsonStatsPathIdentifier(int64_t build_id,
                           int64_t index_version,
                           int64_t collection_id,
                           int64_t partition_id,
                           int64_t segment_id,
                           int64_t field_id);

std::string
GenRemoteJsonStatsPathPrefix(ChunkManagerPtr cm,
                             int64_t build_id,
                             int64_t index_version,
                             int64_t collection_id,
                             int64_t partition_id,
                             int64_t segment_id,
                             int64_t field_id);

std::string
GenNgramIndexPrefix(ChunkManagerPtr cm,
                    int64_t build_id,
                    int64_t index_version,
                    int64_t segment_id,
                    int64_t field_id,
                    bool is_temp);

std::string
GenFieldRawDataPathPrefix(ChunkManagerPtr cm,
                          int64_t segment_id,
                          int64_t field_id);

std::string
GetSegmentRawDataPathPrefix(ChunkManagerPtr cm, int64_t segment_id);

std::pair<std::string, size_t>
EncodeAndUploadIndexSlice(ChunkManager* chunk_manager,
                          uint8_t* buf,
                          int64_t batch_size,
                          IndexMeta index_meta,
                          FieldDataMeta field_meta,
                          std::string object_key,
                          std::shared_ptr<CPluginContext> plugin_context);

std::vector<std::future<std::unique_ptr<DataCodec>>>
GetObjectData(
    ChunkManager* remote_chunk_manager,
    const std::vector<std::string>& remote_files,
    milvus::ThreadPoolPriority priority = milvus::ThreadPoolPriority::HIGH,
    bool is_field_data = true);

// Helper function to wait for all futures and collect exceptions
// This ensures all background threads complete before rethrowing exception
inline void
WaitAllFutures(std::vector<std::future<void>>& futures) {
    std::exception_ptr first_exception = nullptr;

    for (auto& future : futures) {
        try {
            future.get();  // return type is void
        } catch (...) {
            if (!first_exception) {
                first_exception = std::current_exception();
            }
        }
    }

    if (first_exception) {
        std::rethrow_exception(first_exception);
    }
}

// Helper function to wait for all futures and collect exceptions
// This ensures all background threads complete before rethrowing exception
template <typename T>
std::vector<T>
WaitAllFutures(std::vector<std::future<T>> futures) {
    std::vector<T> results;
    results.reserve(futures.size());
    std::exception_ptr first_exception = nullptr;

    // Wait for all futures to complete and collect exceptions
    for (auto& future : futures) {
        try {
            results.emplace_back(future.get());
        } catch (...) {
            if (!first_exception) {
                first_exception = std::current_exception();
            }
            // Continue to wait for remaining futures to ensure all threads complete
        }
    }

    // Rethrow the first exception after all futures are processed
    if (first_exception) {
        std::rethrow_exception(first_exception);
    }

    return results;
}

// Process futures in order, invoking processor on each result as soon as it's
// ready.  This reduces peak memory compared to WaitAllFutures because each
// result can be consumed (and freed) before later futures resolve.
// On exception, we continue waiting for all remaining futures to avoid
// use-after-free on resources captured by background threads (see #46958).
template <typename T, typename Processor>
void
ProcessFuturesInOrder(std::vector<std::future<T>>& futures,
                      Processor&& processor) {
    std::exception_ptr first_exception = nullptr;
    for (auto& future : futures) {
        try {
            auto result = future.get();
            if (!first_exception) {
                processor(std::move(result));
            }
        } catch (...) {
            if (!first_exception) {
                first_exception = std::current_exception();
            }
        }
    }
    if (first_exception) {
        std::rethrow_exception(first_exception);
    }
}

// Best-effort drain for cleanup paths (catch blocks / scope guards): wait for
// every still-valid future and swallow its exception, so that background
// tasks referencing the caller's stack state finish before the caller
// unwinds (see #46958). Unlike WaitAllFutures, this never throws — use it
// when another exception is already propagating.
template <typename FutureContainer>
void
DrainFutures(FutureContainer& futures) noexcept {
    for (auto& future : futures) {
        if (!future.valid()) {
            continue;
        }
        try {
            future.get();
        } catch (...) {
        }
    }
}

template <typename T>
void
DrainFuture(std::future<T>& future) noexcept {
    if (!future.valid()) {
        return;
    }
    try {
        future.get();
    } catch (...) {
    }
}

std::vector<FieldDataPtr>
GetFieldDatasFromStorageV2(std::vector<std::vector<std::string>>& remote_files,
                           int64_t field_id,
                           DataType data_type,
                           DataType element_type,
                           int64_t dim,
                           milvus_storage::ArrowFileSystemPtr fs);

std::vector<FieldDataPtr>
GetFieldDatasFromManifest(
    const std::string& manifest_path,
    const std::shared_ptr<milvus_storage::api::Properties>& loon_ffi_properties,
    const FieldDataMeta& field_meta,
    std::optional<DataType> data_type,
    int64_t dim,
    std::optional<DataType> element_type,
    std::optional<StorageColumnMapping> storage_column_mapping = std::nullopt);

std::vector<FieldDataPtr>
GetTextFieldDatasFromManifest(
    const std::string& manifest_path,
    const std::shared_ptr<milvus_storage::api::Properties>& loon_ffi_properties,
    const FieldDataMeta& field_meta);

std::map<std::string, int64_t>
PutIndexData(ChunkManager* remote_chunk_manager,
             const std::vector<const uint8_t*>& data_slices,
             const std::vector<int64_t>& slice_sizes,
             const std::vector<std::string>& slice_names,
             FieldDataMeta& field_meta,
             IndexMeta& index_meta,
             std::shared_ptr<CPluginContext> plugin_context);

int64_t
GetTotalNumRowsForFieldDatas(const std::vector<FieldDataPtr>& field_datas);

size_t
GetNumRowsForLoadInfo(const LoadFieldDataInfo& load_info);

void
ReleaseArrowUnused();

ChunkManagerPtr
CreateChunkManager(const StorageConfig& storage_config);

milvus_storage::ArrowFileSystemPtr
InitArrowFileSystem(milvus::storage::StorageConfig storage_config);

FieldDataPtr
CreateFieldData(const DataType& type,
                const DataType& element_type,
                bool nullable = false,
                int64_t dim = 1,
                int64_t total_num_rows = 0);

int64_t
GetByteSizeOfFieldDatas(const std::vector<FieldDataPtr>& field_datas);

std::vector<FieldDataPtr>
CollectFieldDataChannel(FieldDataChannelPtr& channel);

FieldDataPtr
MergeFieldData(std::vector<FieldDataPtr>& data_array);

int64_t
ExtractGroupIdFromPath(const std::string& path);

template <typename T, typename = void>
struct has_native_type : std::false_type {};
template <typename T>
struct has_native_type<T, std::void_t<typename T::NativeType>>
    : std::true_type {};
template <DataType T>
using DataTypeNativeOrVoid =
    typename std::conditional<has_native_type<TypeTraits<T>>::value,
                              typename TypeTraits<T>::NativeType,
                              void>::type;
template <DataType T>
using DataTypeToOffsetMap =
    std::unordered_map<DataTypeNativeOrVoid<T>, int64_t>;

std::vector<FieldDataPtr>
FetchFieldData(ChunkManager* cm, const std::vector<std::string>& batch_files);

inline void
SortByPath(std::vector<std::string>& paths) {
    std::sort(paths.begin(),
              paths.end(),
              [](const std::string& a, const std::string& b) {
                  return std::stol(a.substr(a.find_last_of("/") + 1)) <
                         std::stol(b.substr(b.find_last_of("/") + 1));
              });
}

// same as SortByPath, but with pair of path and entries_nums
inline void
SortByPath(std::vector<std::pair<std::string, int64_t>>& paths) {
    std::sort(
        paths.begin(),
        paths.end(),
        [](const std::pair<std::string, int64_t>& a,
           const std::pair<std::string, int64_t>& b) {
            return std::stol(a.first.substr(a.first.find_last_of("/") + 1)) <
                   std::stol(b.first.substr(b.first.find_last_of("/") + 1));
        });
}

template <typename T>
inline void
SortByPath(std::vector<T>& paths) {
    std::sort(paths.begin(), paths.end(), [](const T& a, const T& b) {
        return std::stol(
                   a.file_path.substr(a.file_path.find_last_of("/") + 1)) <
               std::stol(b.file_path.substr(b.file_path.find_last_of("/") + 1));
    });
}

std::vector<FieldDataPtr>
CacheRawDataAndFillMissing(const MemFileManagerImplPtr& file_manager,
                           const Config& config);

// used only for test
inline std::shared_ptr<ArrowDataWrapper>
ConvertFieldDataToArrowDataWrapper(const FieldDataPtr& field_data) {
    BaseEventData event_data;
    event_data.payload_reader = std::make_shared<PayloadReader>(field_data);
    auto event_data_bytes = event_data.Serialize();

    std::shared_ptr<uint8_t[]> file_data(new uint8_t[event_data_bytes.size()]);
    milvus::fastmem::FastMemcpy(
        file_data.get(), event_data_bytes.data(), event_data_bytes.size());

    storage::BinlogReaderPtr reader = std::make_shared<storage::BinlogReader>(
        file_data, event_data_bytes.size());
    event_data = storage::BaseEventData(reader,
                                        event_data_bytes.size(),
                                        field_data->get_data_type(),
                                        field_data->IsNullable(),
                                        false);
    return std::make_shared<ArrowDataWrapper>(
        event_data.payload_reader->get_reader(),
        event_data.payload_reader->get_file_reader(),
        file_data);
}

milvus_storage::FieldIDList
GetFieldIDList(FieldId column_group_id,
               const std::string& filepath,
               const std::shared_ptr<arrow::Schema>& arrow_schema,
               milvus_storage::ArrowFileSystemPtr fs);

// Convert a single row of an Arrow ListArray to a protobuf ScalarField.
// The element type is inferred from the ListArray's value type.
// Supported value types: BOOL, INT8, INT16, INT32, INT64, FLOAT, DOUBLE, STRING.
proto::schema::ScalarField
ArrowListToScalarFieldProto(const std::shared_ptr<arrow::ListArray>& list_array,
                            int64_t row_index);

// Convert a timestamp value to microseconds based on its Arrow TimeUnit.
inline int64_t
ConvertToMicroseconds(int64_t value, arrow::TimeUnit::type unit) {
    switch (unit) {
        case arrow::TimeUnit::SECOND:
            return value * 1000000;
        case arrow::TimeUnit::MILLI:
            return value * 1000;
        case arrow::TimeUnit::MICRO:
            return value;
        case arrow::TimeUnit::NANO:
            return value / 1000;
        default:
            return value;
    }
}

// ============================================================
// Arrow type normalization helpers for external table loading.
// Schemaless Parquet reader returns native Arrow types; these
// helpers convert them to the types each consumer expects.
// ============================================================

// --- Tier 1: Atomic conversion helpers ---

// Convert FixedSizeBinaryArray -> BinaryArray (for nullable vectors).
// Null rows are skipped in the data buffer; offsets encode the gaps.
arrow::ArrayVector
ConvertFixedSizeBinaryToBinary(const arrow::ArrayVector& arrays);

// Zero-copy: reinterpret StringArray as BinaryArray (identical layout).
arrow::ArrayVector
ConvertStringArrayToBinary(const arrow::ArrayVector& arrays);

// Convert StringArray (WKT) -> BinaryArray (WKB) via GEOS.
arrow::ArrayVector
ConvertWKTStringArrayToWKBBinary(const arrow::ArrayVector& arrays);

// Convert TimestampArray -> Int64Array (microseconds).
arrow::ArrayVector
ConvertTimestampToInt64(const arrow::ArrayVector& arrays);

// --- Tier 2: Per-consumer entry points ---

// Unified single-array normalize: converts external Parquet arrow types to
// Milvus internal arrow types. All consumer-specific entry points delegate here.
//
// Conversions:
//   VARCHAR/STRING/TEXT/JSON: String -> Binary (zero-copy)
//   Geometry: String(WKT) -> Binary(WKB) via GEOS; Binary stays as-is
//   Timestamptz: Timestamp -> Int64 (microseconds)
//   Array: List(element) -> Binary (protobuf)
//   Vectors: various -> FixedSizeBinary
//   VectorArray: List<List<scalar>> -> List<FixedSizeBinary>
//
std::shared_ptr<arrow::Array>
NormalizeExternalArrow(const std::shared_ptr<arrow::Array>& array,
                       const FieldMeta& field_meta);

// Load path: batch wrapper around NormalizeExternalArrow.
arrow::ArrayVector
NormalizeArrowForChunkWriter(const arrow::ArrayVector& arrays,
                             const FieldMeta& field_meta);

// Coerce any binary-like array (LARGE_BINARY / BINARY_VIEW /
// LARGE_STRING / STRING_VIEW / STRING) to canonical BinaryArray.
// Required because vortex schemaless mode emits view variants for the
// whole variable-length family; downstream paths assume canonical layout.
arrow::ArrayVector
CoerceToBinary(const arrow::ArrayVector& arrays);

// Coerce LARGE_LIST / LIST_VIEW to canonical (32-bit offset) ListArray.
// Vortex schemaless mode may emit list variants for the same logical
// List<T>; downstream code expects arrow::ListArray.
arrow::ArrayVector
CoerceToList(const arrow::ArrayVector& arrays);

// Single source of truth for view/large-variant elimination.
// STRING_VIEW/LARGE_STRING -> STRING, BINARY_VIEW/LARGE_BINARY -> BINARY,
// LARGE_LIST/LIST_VIEW -> LIST (recursive into List/FixedSizeList inner).
// Pure type/layout normalization, no semantic conversion. Callers ingesting
// schemaless readers (vortex) should route arrays through this helper before
// any type-id dispatch (refresh / load / sample / etc).
std::shared_ptr<arrow::Array>
CanonicalizeArrowVariants(const std::shared_ptr<arrow::Array>& array);

}  // namespace milvus::storage
