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

#include "storage/minio/MinioChunkManager.h"

#include <aws/core/auth/AWSCredentials.h>
#include <aws/core/auth/AWSCredentialsProviderChain.h>
#include <aws/s3/model/CreateBucketRequest.h>
#include <aws/s3/model/DeleteBucketRequest.h>
#include <aws/s3/model/DeleteObjectRequest.h>
#include <aws/s3/model/GetObjectRequest.h>
#include <aws/s3/model/HeadBucketRequest.h>
#include <aws/s3/model/HeadObjectRequest.h>
#include <aws/s3/model/ListObjectsRequest.h>
#include <aws/s3/model/PutObjectRequest.h>
#include <string.h>
#include <unistd.h>
#include <algorithm>
#include <chrono>
#include <exception>
#include <functional>
#include <sstream>

#include "aws/core/Aws.h"
#include "aws/core/auth/signer/AWSAuthV4Signer.h"
#include "aws/core/http/Scheme.h"
#include "aws/core/utils/Outcome.h"
#include "aws/core/utils/memory/AWSMemory.h"
#include "aws/core/utils/memory/stl/AWSStringStream.h"
#include "aws/s3/S3Client.h"
#include "aws/s3/model/Bucket.h"
#include "aws/s3/model/HeadObjectResult.h"
#include "aws/s3/model/ListBucketsResult.h"
#include "aws/s3/model/ListObjectsResult.h"
#include "aws/s3/model/Object.h"
#include "common/Consts.h"
#include "common/EasyAssert.h"
#include "fmt/core.h"
#include "google/cloud/internal/rest_client.h"
#include "google/cloud/storage/oauth2/compute_engine_credentials.h"
#include "google/cloud/version.h"
#include "log/Log.h"
#include "monitor/Monitor.h"
#include "prometheus/counter.h"
#include "prometheus/histogram.h"
#include "signal.h"
#include "storage/Types.h"
#include "storage/aliyun/AliyunCredentialsProvider.h"

namespace milvus::storage {

std::atomic<size_t> MinioChunkManager::init_count_(0);
std::mutex MinioChunkManager::client_mutex_;

void
ConfigureGoogleCloudIAMHttpClientFactory(Aws::SDKOptions& sdk_options) {
    sdk_options.httpOptions.httpClientFactory_create_fn = []() {
        auto credentials =
            std::make_shared<google::cloud::oauth2_internal::
                                 GOOGLE_CLOUD_CPP_NS::ComputeEngineCredentials>(
                google::cloud::Options{},
                [](google::cloud::Options const& opts) {
                    return google::cloud::rest_internal::MakeDefaultRestClient(
                        "", opts);
                });
        return Aws::MakeShared<GoogleHttpClientFactory>(
            GOOGLE_CLIENT_FACTORY_ALLOCATION_TAG, credentials);
    };
}

static void
SwallowHandler(int signal) {
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-result"
    switch (signal) {
        case SIGPIPE:
            // cannot use log or stdio
            write(1, "SIGPIPE Swallowed\n", 18);
            break;
        default:
            // cannot use log or stdio
            write(2, "Unexpected signal\n", 18);
    }
#pragma GCC diagnostic pop
}

/**
 * @brief convert std::string to Aws::String
 * because Aws has String type internally
 * but has a copy of string content unfortunately
 * TODO: remove this convert
 * @param str
 * @return Aws::String
 */
inline Aws::String
ConvertToAwsString(const std::string& str) {
    return Aws::String(str.c_str(), str.size());
}

/**
 * @brief convert Aws::string to std::string
 * @param aws_str
 * @return std::string
 */
inline std::string
ConvertFromAwsString(const Aws::String& aws_str) {
    return std::string(aws_str.c_str(), aws_str.size());
}

void
AwsLogger::ProcessFormattedStatement(Aws::String&& statement) {
    LOG_INFO("[AWS LOG] {}", statement);
}

void
MinioChunkManager::InitSDKAPI(RemoteStorageType type,
                              bool useIAM,
                              const std::string& log_level_str,
                              const std::string& tls_min_version) {
    std::scoped_lock lock{client_mutex_};
    const size_t initCount = init_count_++;
    if (initCount == 0) {
        // sdk_options_.httpOptions.installSigPipeHandler = true;
        struct sigaction psa;
        memset(&psa, 0, sizeof psa);
        psa.sa_handler = SwallowHandler;
        psa.sa_flags = psa.sa_flags | SA_ONSTACK;
        sigaction(SIGPIPE, &psa, 0);
        // block multiple SIGPIPE concurrently processing
        sigemptyset(&psa.sa_mask);
        sigaddset(&psa.sa_mask, SIGPIPE);
        sigaction(SIGPIPE, &psa, 0);
        bool need_tls =
            !tls_min_version.empty() && tls_min_version != "default";
        if (type == RemoteStorageType::GOOGLE_CLOUD && useIAM) {
            ConfigureGoogleCloudIAMHttpClientFactory(sdk_options_);
        } else if (need_tls) {
            sdk_options_.httpOptions.httpClientFactory_create_fn =
                [tls_min_version]() {
                    return Aws::MakeShared<TlsHttpClientFactory>(
                        TLS_FACTORY_ALLOCATION_TAG, tls_min_version);
                };
        }
        LOG_INFO("init aws with log level:{}", log_level_str);
        auto get_aws_log_level = [](const std::string& level_str) {
            Aws::Utils::Logging::LogLevel level =
                Aws::Utils::Logging::LogLevel::Off;
            if (level_str == "fatal") {
                level = Aws::Utils::Logging::LogLevel::Fatal;
            } else if (level_str == "error") {
                level = Aws::Utils::Logging::LogLevel::Error;
            } else if (level_str == "warn") {
                level = Aws::Utils::Logging::LogLevel::Warn;
            } else if (level_str == "info") {
                level = Aws::Utils::Logging::LogLevel::Info;
            } else if (level_str == "debug") {
                level = Aws::Utils::Logging::LogLevel::Debug;
            } else if (level_str == "trace") {
                level = Aws::Utils::Logging::LogLevel::Trace;
            }
            return level;
        };
        auto log_level = get_aws_log_level(log_level_str);
        sdk_options_.loggingOptions.logLevel = log_level;
        sdk_options_.loggingOptions.logger_create_fn = [log_level]() {
            return std::make_shared<AwsLogger>(log_level);
        };
        Aws::InitAPI(sdk_options_);
    }
}

void
MinioChunkManager::InitSDKAPIDefault(const std::string& log_level_str,
                                     const std::string& tls_min_version) {
    std::scoped_lock lock{client_mutex_};
    const size_t initCount = init_count_++;
    if (initCount == 0) {
        // sdk_options_.httpOptions.installSigPipeHandler = true;
        struct sigaction psa;
        memset(&psa, 0, sizeof psa);
        psa.sa_handler = SwallowHandler;
        psa.sa_flags = psa.sa_flags | SA_ONSTACK;
        sigaction(SIGPIPE, &psa, 0);
        // block multiple SIGPIPE concurrently processing
        sigemptyset(&psa.sa_mask);
        sigaddset(&psa.sa_mask, SIGPIPE);
        sigaction(SIGPIPE, &psa, 0);
        bool need_tls =
            !tls_min_version.empty() && tls_min_version != "default";
        if (need_tls && !sdk_options_.httpOptions.httpClientFactory_create_fn) {
            sdk_options_.httpOptions.httpClientFactory_create_fn =
                [tls_min_version]() {
                    return Aws::MakeShared<TlsHttpClientFactory>(
                        TLS_FACTORY_ALLOCATION_TAG, tls_min_version);
                };
        }
        LOG_INFO("init aws with log level:{}", log_level_str);
        auto get_aws_log_level = [](const std::string& level_str) {
            Aws::Utils::Logging::LogLevel level =
                Aws::Utils::Logging::LogLevel::Off;
            if (level_str == "fatal") {
                level = Aws::Utils::Logging::LogLevel::Fatal;
            } else if (level_str == "error") {
                level = Aws::Utils::Logging::LogLevel::Error;
            } else if (level_str == "warn") {
                level = Aws::Utils::Logging::LogLevel::Warn;
            } else if (level_str == "info") {
                level = Aws::Utils::Logging::LogLevel::Info;
            } else if (level_str == "debug") {
                level = Aws::Utils::Logging::LogLevel::Debug;
            } else if (level_str == "trace") {
                level = Aws::Utils::Logging::LogLevel::Trace;
            }
            return level;
        };
        auto log_level = get_aws_log_level(log_level_str);
        sdk_options_.loggingOptions.logLevel = log_level;
        sdk_options_.loggingOptions.logger_create_fn = [log_level]() {
            return std::make_shared<AwsLogger>(log_level);
        };
        Aws::InitAPI(sdk_options_);
    }
}

void
MinioChunkManager::ShutdownSDKAPI() {
    std::scoped_lock lock{client_mutex_};
    const size_t initCount = --init_count_;
    if (initCount == 0) {
        Aws::ShutdownAPI(sdk_options_);
    }
}

void
MinioChunkManager::BuildS3Client(
    const StorageConfig& storage_config,
    const Aws::Client::ClientConfiguration& config) {
    if (storage_config.useIAM) {
        auto provider =
            std::make_shared<Aws::Auth::DefaultAWSCredentialsProviderChain>();
        auto aws_credentials = provider->GetAWSCredentials();
        AssertInfo(!aws_credentials.GetAWSAccessKeyId().empty(),
                   "if use iam, access key id should not be empty");
        AssertInfo(!aws_credentials.GetAWSSecretKey().empty(),
                   "if use iam, secret key should not be empty");
        AssertInfo(!aws_credentials.GetSessionToken().empty(),
                   "if use iam, token should not be empty");

        client_ = std::make_shared<Aws::S3::S3Client>(
            provider,
            config,
            Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Never,
            storage_config.useVirtualHost);
    } else {
        BuildAccessKeyClient(storage_config, config);
    }
}

void
MinioChunkManager::PreCheck(const StorageConfig& config) {
    LOG_INFO("start to precheck chunk manager with configuration: {}",
             config.ToString());
    try {
        // Just test connection not check real list, avoid cost resource.
        std::string check_prefix = "justforconnectioncheck";
        if (!remote_root_path_.empty()) {
            check_prefix = remote_root_path_ + "/" + check_prefix;
        }
        ListWithPrefix(check_prefix);
    } catch (const SegcoreError& e) {
        auto err_message = fmt::format(
            "precheck chunk manager client failed, "
            "error:{}, "
            "configuration:{}",
            e.what(),
            config.ToString());
        LOG_ERROR("{}", err_message);
        throw SegcoreError(e.get_error_code(), err_message);
    } catch (std::exception& e) {
        throw e;
    }
};

void
MinioChunkManager::BuildAccessKeyClient(
    const StorageConfig& storage_config,
    const Aws::Client::ClientConfiguration& config) {
    AssertInfo(!storage_config.access_key_id.empty(),
               "if not use iam, access key should not be empty");
    AssertInfo(!storage_config.access_key_value.empty(),
               "if not use iam, access value should not be empty");

    client_ = std::make_shared<Aws::S3::S3Client>(
        Aws::Auth::AWSCredentials(
            ConvertToAwsString(storage_config.access_key_id),
            ConvertToAwsString(storage_config.access_key_value)),
        config,
        Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Never,
        storage_config.useVirtualHost);
}

void
MinioChunkManager::BuildAliyunCloudClient(
    const StorageConfig& storage_config,
    const Aws::Client::ClientConfiguration& config) {
    // For aliyun oss, support use virtual host mode
    StorageConfig mutable_config = storage_config;
    mutable_config.useVirtualHost = true;
    if (storage_config.useIAM) {
        auto aliyun_provider = Aws::MakeShared<
            Aws::Auth::AliyunSTSAssumeRoleWebIdentityCredentialsProvider>(
            "AliyunSTSAssumeRoleWebIdentityCredentialsProvider");
        auto aliyun_credentials = aliyun_provider->GetAWSCredentials();
        AssertInfo(!aliyun_credentials.GetAWSAccessKeyId().empty(),
                   "if use iam, access key id should not be empty");
        AssertInfo(!aliyun_credentials.GetAWSSecretKey().empty(),
                   "if use iam, secret key should not be empty");
        AssertInfo(!aliyun_credentials.GetSessionToken().empty(),
                   "if use iam, token should not be empty");
        client_ = std::make_shared<Aws::S3::S3Client>(
            aliyun_provider,
            config,
            Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Never,
            mutable_config.useVirtualHost);
    } else {
        BuildAccessKeyClient(mutable_config, config);
    }
}

void
MinioChunkManager::BuildGoogleCloudClient(
    const StorageConfig& storage_config,
    const Aws::Client::ClientConfiguration& config) {
    if (storage_config.useIAM) {
        // Using S3 client instead of google client because of compatible protocol
        client_ = std::make_shared<Aws::S3::S3Client>(
            config,
            Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Never,
            storage_config.useVirtualHost);
    } else {
        BuildAccessKeyClient(storage_config, config);
    }
}

void
MinioChunkManager::ApplyChecksumConfigOverrides(
    Aws::Client::ClientConfiguration& config) {
    // Non-AWS S3-compatible APIs (GCP, Aliyun OSS, Tencent COS, Huawei OBS) do
    // not accept the extra x-amz-checksum-* headers / aws-chunked streaming
    // that AWS SDK >= 1.11.x sends by default (WHEN_SUPPORTED). The V4 signer
    // emits Transfer-Encoding: aws-chunked + x-amz-content-sha256:
    // STREAMING-UNSIGNED-PAYLOAD-TRAILER, which Aliyun OSS rejects with
    // x-oss-ec=0017-00000804, Huawei OBS with XAmzContentSHA256Mismatch, etc.
    // Restrict to WHEN_REQUIRED so the SDK only adds checksums when the
    // operation model mandates them (PutObject does not), restoring the
    // simple PUT behavior we had with AWS SDK 1.11.352. Same approach as
    // milvus-storage PR #500.
    config.checksumConfig.requestChecksumCalculation =
        Aws::Client::RequestChecksumCalculation::WHEN_REQUIRED;
    config.checksumConfig.responseChecksumValidation =
        Aws::Client::ResponseChecksumValidation::WHEN_REQUIRED;
}

bool
MinioChunkManager::NeedChecksumOverride(const std::string& cloud_provider) {
    return cloud_provider == "gcp" || cloud_provider == "aliyun" ||
           cloud_provider == "tencent" || cloud_provider == "huawei";
}

MinioChunkManager::MinioChunkManager(const StorageConfig& storage_config)
    : default_bucket_name_(storage_config.bucket_name),
      use_crc32c_checksum_(storage_config.use_crc32c_checksum) {
    remote_root_path_ = storage_config.root_path;
    RemoteStorageType storageType;
    if (storage_config.address.find("google") != std::string::npos) {
        storageType = RemoteStorageType::GOOGLE_CLOUD;
    } else if (storage_config.address.find("aliyun") != std::string::npos) {
        storageType = RemoteStorageType::ALIYUN_CLOUD;
    } else {
        storageType = RemoteStorageType::S3;
    }

    InitSDKAPI(storageType,
               storage_config.useIAM,
               storage_config.log_level,
               storage_config.tls_min_version);

    // The ClientConfiguration default constructor will take a long time.
    // For more details, please refer to https://github.com/aws/aws-sdk-cpp/issues/1440
    static Aws::Client::ClientConfiguration g_config;
    Aws::Client::ClientConfiguration config = g_config;
    config.endpointOverride = ConvertToAwsString(storage_config.address);

    // Three cases:
    // 1. no ssl, verifySSL=false
    // 2. self-signed certificate, verifySSL=false
    // 3. CA-signed certificate, verifySSL=true
    if (storage_config.useSSL) {
        config.scheme = Aws::Http::Scheme::HTTPS;
        config.verifySSL = true;
        if (!storage_config.sslCACert.empty()) {
            config.caPath = ConvertToAwsString(storage_config.sslCACert);
            config.verifySSL = false;
        }
    } else {
        config.scheme = Aws::Http::Scheme::HTTP;
        config.verifySSL = false;
    }

    config.requestTimeoutMs = storage_config.requestTimeoutMs == 0
                                  ? DEFAULT_CHUNK_MANAGER_REQUEST_TIMEOUT_MS
                                  : storage_config.requestTimeoutMs;

    if (storage_config.max_connections > 0) {
        config.maxConnections = storage_config.max_connections;
    }

    if (!storage_config.region.empty()) {
        config.region = ConvertToAwsString(storage_config.region);
    }

    if (NeedChecksumOverride(storage_config.cloud_provider)) {
        ApplyChecksumConfigOverrides(config);
    }

    if (storageType == RemoteStorageType::S3) {
        BuildS3Client(storage_config, config);
    } else if (storageType == RemoteStorageType::ALIYUN_CLOUD) {
        BuildAliyunCloudClient(storage_config, config);
    } else if (storageType == RemoteStorageType::GOOGLE_CLOUD) {
        BuildGoogleCloudClient(storage_config, config);
    }

    PreCheck(storage_config);

    LOG_INFO(
        "init MinioChunkManager with "
        "parameter[endpoint={}][bucket_name={}][root_path={}][use_secure={}]"
        "[tls_min_version={}]",
        storage_config.address,
        storage_config.bucket_name,
        storage_config.root_path,
        storage_config.useSSL,
        storage_config.tls_min_version.empty()
            ? "default"
            : storage_config.tls_min_version);
}

MinioChunkManager::~MinioChunkManager() {
    client_.reset();
    // ShutdownSDKAPI();
}

uint64_t
MinioChunkManager::Size(const std::string& filepath) {
    return GetObjectSize(default_bucket_name_, filepath);
}

bool
MinioChunkManager::Exist(const std::string& filepath) {
    return ObjectExists(default_bucket_name_, filepath);
}

void
MinioChunkManager::Remove(const std::string& filepath) {
    DeleteObject(default_bucket_name_, filepath);
}

std::vector<std::string>
MinioChunkManager::ListWithPrefix(const std::string& filepath) {
    return ListObjects(default_bucket_name_, filepath);
}

uint64_t
MinioChunkManager::Read(const std::string& filepath, void* buf, uint64_t size) {
    return GetObjectBuffer(default_bucket_name_, filepath, buf, size);
}

void
MinioChunkManager::Write(const std::string& filepath,
                         void* buf,
                         uint64_t size) {
    PutObjectBuffer(default_bucket_name_, filepath, buf, size);
}

bool
MinioChunkManager::BucketExists(const std::string& bucket_name) {
    Aws::S3::Model::HeadBucketRequest request;
    request.SetBucket(bucket_name.c_str());

    auto outcome = client_->HeadBucket(request);

    if (!outcome.IsSuccess()) {
        const auto& err = outcome.GetError();
        auto error_type = err.GetErrorType();
        // only throw if the error is not nosuchbucket
        // if bucket not exist, HeadBucket return errorType RESOURCE_NOT_FOUND
        if (error_type != Aws::S3::S3Errors::NO_SUCH_BUCKET &&
            error_type != Aws::S3::S3Errors::RESOURCE_NOT_FOUND) {
            ThrowS3Error("BucketExists", err, "params, bucket={}", bucket_name);
        }
        return false;
    }
    return true;
}

std::vector<std::string>
MinioChunkManager::ListBuckets() {
    std::vector<std::string> buckets;
    auto outcome = client_->ListBuckets();

    if (!outcome.IsSuccess()) {
        const auto& err = outcome.GetError();
        ThrowS3Error("ListBuckets", err, "params");
    }
    for (auto&& b : outcome.GetResult().GetBuckets()) {
        buckets.emplace_back(b.GetName());
    }
    return buckets;
}

bool
MinioChunkManager::CreateBucket(const std::string& bucket_name) {
    Aws::S3::Model::CreateBucketRequest request;
    request.SetBucket(bucket_name.c_str());

    auto outcome = client_->CreateBucket(request);

    if (!outcome.IsSuccess()) {
        const auto& err = outcome.GetError();
        if (err.GetErrorType() !=
            Aws::S3::S3Errors::BUCKET_ALREADY_OWNED_BY_YOU) {
            ThrowS3Error("CreateBucket", err, "params, bucket={}", bucket_name);
        }
        return false;
    }
    return true;
}

bool
MinioChunkManager::DeleteBucket(const std::string& bucket_name) {
    Aws::S3::Model::DeleteBucketRequest request;
    request.SetBucket(bucket_name.c_str());

    auto outcome = client_->DeleteBucket(request);

    if (!outcome.IsSuccess()) {
        const auto& err = outcome.GetError();
        auto error_type = err.GetErrorType();
        if (error_type != Aws::S3::S3Errors::NO_SUCH_BUCKET &&
            error_type != Aws::S3::S3Errors::RESOURCE_NOT_FOUND) {
            ThrowS3Error("DeleteBucket", err, "params, bucket={}", bucket_name);
        }
        return false;
    }
    return true;
}

bool
MinioChunkManager::ObjectExists(const std::string& bucket_name,
                                const std::string& object_name) {
    Aws::S3::Model::HeadObjectRequest request;
    request.SetBucket(bucket_name.c_str());
    request.SetKey(object_name.c_str());

    auto start = std::chrono::system_clock::now();
    auto outcome = client_->HeadObject(request);
    milvus::monitor::internal_storage_request_latency_stat.Observe(
        std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::system_clock::now() - start)
            .count());

    if (!outcome.IsSuccess()) {
        const auto& err = outcome.GetError();
        if (!IsNotFound(err.GetErrorType())) {
            milvus::monitor::internal_storage_op_count_stat_fail.Increment();
            ThrowS3Error("ObjectExists",
                         err,
                         "params, bucket={}, object={}",
                         bucket_name,
                         object_name);
        }
        milvus::monitor::internal_storage_op_count_stat_suc.Increment();
        return false;
    }
    milvus::monitor::internal_storage_op_count_stat_suc.Increment();
    return true;
}

uint64_t
MinioChunkManager::GetObjectSize(const std::string& bucket_name,
                                 const std::string& object_name) {
    Aws::S3::Model::HeadObjectRequest request;
    request.SetBucket(bucket_name.c_str());
    request.SetKey(object_name.c_str());

    auto start = std::chrono::system_clock::now();
    auto outcome = client_->HeadObject(request);
    milvus::monitor::internal_storage_request_latency_stat.Observe(
        std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::system_clock::now() - start)
            .count());
    if (!outcome.IsSuccess()) {
        milvus::monitor::internal_storage_op_count_stat_fail.Increment();
        const auto& err = outcome.GetError();
        ThrowS3Error("GetObjectSize",
                     err,
                     "params, bucket={}, object={}",
                     bucket_name,
                     object_name);
    }
    milvus::monitor::internal_storage_op_count_stat_suc.Increment();
    return outcome.GetResult().GetContentLength();
}

bool
MinioChunkManager::DeleteObject(const std::string& bucket_name,
                                const std::string& object_name) {
    Aws::S3::Model::DeleteObjectRequest request;
    request.SetBucket(bucket_name.c_str());
    request.SetKey(object_name.c_str());

    auto start = std::chrono::system_clock::now();
    auto outcome = client_->DeleteObject(request);
    milvus::monitor::internal_storage_request_latency_remove.Observe(
        std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::system_clock::now() - start)
            .count());

    if (!outcome.IsSuccess()) {
        const auto& err = outcome.GetError();
        if (!IsNotFound(err.GetErrorType())) {
            milvus::monitor::internal_storage_op_count_remove_fail.Increment();
            ThrowS3Error("DeleteObject",
                         err,
                         "params, bucket={}, object={}",
                         bucket_name,
                         object_name);
        }
        milvus::monitor::internal_storage_op_count_remove_suc.Increment();
        return false;
    }
    milvus::monitor::internal_storage_op_count_remove_suc.Increment();
    return true;
}

bool
MinioChunkManager::PutObjectBuffer(const std::string& bucket_name,
                                   const std::string& object_name,
                                   void* buf,
                                   uint64_t size) {
    Aws::S3::Model::PutObjectRequest request;
    request.SetBucket(bucket_name.c_str());
    request.SetKey(object_name.c_str());
    if (use_crc32c_checksum_) {
        request.SetChecksumAlgorithm(Aws::S3::Model::ChecksumAlgorithm::CRC32C);
    }

    const std::shared_ptr<Aws::IOStream> input_data =
        Aws::MakeShared<Aws::StringStream>("");

    input_data->write(reinterpret_cast<char*>(buf), size);
    request.SetBody(input_data);

    auto start = std::chrono::system_clock::now();
    auto outcome = client_->PutObject(request);
    milvus::monitor::internal_storage_request_latency_put.Observe(
        std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::system_clock::now() - start)
            .count());
    milvus::monitor::internal_storage_kv_size_put.Observe(size);

    if (!outcome.IsSuccess()) {
        milvus::monitor::internal_storage_op_count_put_fail.Increment();
        const auto& err = outcome.GetError();
        ThrowS3Error("PutObjectBuffer",
                     err,
                     "params, bucket={}, object={}",
                     bucket_name,
                     object_name);
    }
    milvus::monitor::internal_storage_op_count_put_suc.Increment();
    return true;
}

class AwsStreambuf : public std::streambuf {
 public:
    AwsStreambuf(char* buffer, std::streamsize buffer_size) {
        setp(buffer, buffer + buffer_size - 1);
    }

 protected:
    int_type
    overflow(int_type ch) override {
        if (ch != traits_type::eof()) {
            *pptr() = ch;
            pbump(1);
        }
        return ch;
    }
};

class AwsResponseStream : public Aws::IOStream {
 public:
    /**
     * Creates a stream for get response from server
     * @param buffer the buffer address from user space
     * @param size length of the underlying buffer.
     */
    AwsResponseStream(char* buffer, int64_t size)
        : Aws::IOStream(&aws_streambuf), aws_streambuf(buffer, size) {
    }

 private:
    AwsResponseStream(const AwsResponseStream&) = delete;
    AwsResponseStream(AwsResponseStream&&) = delete;
    AwsResponseStream&
    operator=(const AwsResponseStream&) = delete;
    AwsResponseStream&
    operator=(AwsResponseStream&&) = delete;

    AwsStreambuf aws_streambuf;
};

uint64_t
MinioChunkManager::GetObjectBuffer(const std::string& bucket_name,
                                   const std::string& object_name,
                                   void* buf,
                                   uint64_t size) {
    Aws::S3::Model::GetObjectRequest request;
    request.SetBucket(bucket_name.c_str());
    request.SetKey(object_name.c_str());

    request.SetResponseStreamFactory([buf, size]() {
    // For macOs, pubsetbuf interface not implemented
#ifdef __linux__
        std::unique_ptr<Aws::StringStream> stream(
            Aws::New<Aws::StringStream>(""));
        stream->rdbuf()->pubsetbuf(static_cast<char*>(buf), size);
#else
        std::unique_ptr<Aws::IOStream> stream(Aws::New<AwsResponseStream>(
            "AwsResponseStream", static_cast<char*>(buf), size));
#endif
        return stream.release();
    });
    auto start = std::chrono::system_clock::now();
    auto outcome = client_->GetObject(request);
    milvus::monitor::internal_storage_request_latency_get.Observe(
        std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::system_clock::now() - start)
            .count());
    milvus::monitor::internal_storage_kv_size_get.Observe(size);

    if (!outcome.IsSuccess()) {
        milvus::monitor::internal_storage_op_count_get_fail.Increment();
        const auto& err = outcome.GetError();
        ThrowS3Error("GetObjectBuffer",
                     err,
                     "params, bucket={}, object={}",
                     bucket_name,
                     object_name);
    }
    milvus::monitor::internal_storage_op_count_get_suc.Increment();
    return size;
}

std::vector<std::string>
MinioChunkManager::ListObjects(const std::string& bucket_name,
                               const std::string& prefix) {
    std::vector<std::string> objects_vec;
    Aws::S3::Model::ListObjectsRequest request;
    request.WithBucket(bucket_name);
    if (prefix != "") {
        request.SetPrefix(prefix);
    }

    auto start = std::chrono::system_clock::now();
    auto outcome = client_->ListObjects(request);
    milvus::monitor::internal_storage_request_latency_list.Observe(
        std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::system_clock::now() - start)
            .count());

    if (!outcome.IsSuccess()) {
        milvus::monitor::internal_storage_op_count_list_fail.Increment();
        const auto& err = outcome.GetError();
        ThrowS3Error("ListObjects",
                     err,
                     "params, bucket={}, prefix={}",
                     bucket_name,
                     prefix);
    }
    milvus::monitor::internal_storage_op_count_list_suc.Increment();
    auto objects = outcome.GetResult().GetContents();
    for (auto& obj : objects) {
        objects_vec.emplace_back(obj.GetKey());
    }
    return objects_vec;
}

}  // namespace milvus::storage
