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

#include "io/fs/s3_file_system.h"

#include <fmt/format.h>

#include <cstddef>

#include "common/compiler_util.h" // IWYU pragma: keep

// IWYU pragma: no_include <bits/chrono.h>

#include <aws/core/utils/threading/Executor.h>

#include <aws/s3/S3Client.h>

#include <chrono> // IWYU pragma: keep

#include <filesystem>

#include <fstream> // IWYU pragma: keep

#include <future>

#include <memory>

#include <variant>

#include "common/check.h"

#include "common/config.h"

#include "common/logging.h"

#include "common/status.h"

#include "cpp/obj-client/obj_storage_client.h"

#include "cpp/obj-client/s3_common.h"

#include "cpp/sync_point.h"

#include "io/fs/err_utils.h"

#include "io/fs/file_system.h"

#include "io/fs/file_writer.h"

#include "io/fs/local_file_system.h"

#include "io/fs/remote_file_system.h"

#include "io/fs/s3_file_reader.h"

#include "io/fs/s3_file_writer.h"

#include "io/obj_storage/obj_client_factory.h"

#include "io/obj_storage/obj_storage_config_resolver.h"

#include "io/obj_storage/obj_storage_uri.h"

#include "runtime/exec_env.h"

#include "runtime/thread_context.h"

#include "util/s3_uri.h"

#include "util/s3_util.h"

namespace doris::io {

namespace {

constexpr std::string_view OSS_PRIVATE_ENDPOINT_SUFFIX = "-internal.aliyuncs.com";

constexpr int LEN_OF_OSS_PRIVATE_SUFFIX = 9; // length of "-internal"

#ifndef CHECK_S3_CLIENT

#define CHECK_S3_CLIENT(client)                                 \

    if (!client) {                                              \

        return Status::InvalidArgument("init s3 client error"); \

    }

#endif

} // namespace

ObjClientHolder::ObjClientHolder(S3ClientConf conf)

        : _state(std::make_shared<ObjClientState>(ObjClientState {

                  .config = ObjStorageConfigResolver::from_client_conf(conf),

                  .legacy_conf = std::move(conf),

                  .client = nullptr,

                  .generation = 0,

          })) {}

ObjClientHolder::ObjClientHolder(ObjClientConfig config)

        : _state(std::make_shared<ObjClientState>(ObjClientState {

                  .config = std::move(config),

                  .legacy_conf = std::nullopt,

                  .client = nullptr,

                  .generation = 0,

          })) {}

ObjClientHolder::~ObjClientHolder() = default;

Status ObjClientHolder::init() {

    auto current = state();

    std::shared_ptr<ObjStorageClient> client;

    if (current->legacy_conf.has_value()) {

        // A legacy conf still goes through S3ClientFactory, which applies the legacy

        // validation rules (e.g. a non-empty region) and the test client hook.

        client = DORIS_TRY(S3ClientFactory::instance().create(*current->legacy_conf));

    } else {

        client = DORIS_TRY(ObjClientFactory::instance().get_or_create(current->config));

    }

    return _publish(current, current->config, current->legacy_conf, std::move(client));

}

S3ClientConf ObjClientHolder::s3_client_conf() const {

    auto current = state();

    DORIS_CHECK(current->legacy_conf.has_value());

    return *current->legacy_conf;

}

Status ObjClientHolder::reset(const S3ClientConf& conf) {

    auto current = state();

    DORIS_CHECK(current->legacy_conf.has_value());

    S3ClientConf reset_conf = *current->legacy_conf;

    reset_conf.ak = conf.ak;

    reset_conf.sk = conf.sk;

    reset_conf.token = conf.token;

    reset_conf.bucket = conf.bucket;

    reset_conf.connect_timeout_ms = conf.connect_timeout_ms;

    reset_conf.max_connections = conf.max_connections;

    reset_conf.request_timeout_ms = conf.request_timeout_ms;

    reset_conf.use_virtual_addressing = conf.use_virtual_addressing;

    reset_conf.is_internal_bucket = conf.is_internal_bucket;

    reset_conf.role_arn = conf.role_arn;

    reset_conf.external_id = conf.external_id;

    reset_conf.cred_provider_type = conf.cred_provider_type;

    // Compare full-field equality of the merged conf, not get_hash(): the hash is

    // an XOR of crc32s and distinct configurations can collide, which would skip a

    // required client rebuild (e.g. a credential update).

    if (reset_conf == *current->legacy_conf) {

        return Status::OK(); // Same conf

    }

    auto client = DORIS_TRY(S3ClientFactory::instance().create(reset_conf));

    LOG(WARNING) << "reset s3 client with new conf: " << conf.to_string();

    auto config = ObjStorageConfigResolver::from_client_conf(reset_conf);

    return _publish(current, std::move(config), std::move(reset_conf), std::move(client));

}

Status ObjClientHolder::_publish(const std::shared_ptr<const ObjClientState>& expected,

                                 ObjClientConfig config, std::optional<S3ClientConf> legacy_conf,

                                 std::shared_ptr<ObjStorageClient> client) {

    // Declared before the lock: after the swap below it holds the replaced state, and

    // possibly the last reference to its client, which is then released after unlocking.

    std::shared_ptr<const ObjClientState> next = std::make_shared<ObjClientState>(ObjClientState {

            .config = std::move(config),

            .legacy_conf = std::move(legacy_conf),

            .client = std::move(client),

            .generation = expected->generation + 1,

    });

    {

        std::lock_guard lock(_mtx);

        // Two resets that both start from generation N merge their conf onto N. Publishing

        // the second one on top of N + 1 would silently drop the first update, so it is

        // rejected and the caller retries from the new state.

        if (_state->generation != expected->generation) {

            return Status::Aborted(

                    "s3 client was updated concurrently, expected generation {}, current {}",

                    expected->generation, _state->generation);

        }

        _state.swap(next);

    }

    return Status::OK();

}

Result<int64_t> ObjClientHolder::object_file_size(const std::string& bucket,

                                                  const std::string& key) const {

    auto client = get();

    if (!client) {

        return ResultError(Status::InvalidArgument("init s3 client error"));

    }

    auto resp = client->head_object({

            .bucket = bucket,

            .key = key,

    });

    if (resp.resp.status.code != ErrorCode::OK) {

        return ResultError(std::move(Status(resp.resp.status.code, std::move(resp.resp.status.msg))

                                             .append(fmt::format("failed to head s3 file {}",

                                                                 full_s3_path(bucket, key)))));

    }

    return resp.file_size;

}

std::string ObjClientHolder::full_s3_path(std::string_view bucket, std::string_view key) const {

    return fmt::format("{}/{}/{}", state()->config.endpoint(), bucket, key);

}

std::string S3FileSystem::full_s3_path(std::string_view key) const {

    return _client->full_s3_path(_bucket, key);

}

Result<std::shared_ptr<S3FileSystem>> S3FileSystem::create(S3Conf s3_conf, std::string id) {

    std::shared_ptr<S3FileSystem> fs(new S3FileSystem(std::move(s3_conf), std::move(id)));

    RETURN_IF_ERROR_RESULT(fs->init());

    return fs;

}

Result<std::shared_ptr<S3FileSystem>> S3FileSystem::create(ObjClientConfig config,

                                                           std::string bucket, std::string id) {

    std::shared_ptr<S3FileSystem> fs(

            new S3FileSystem(std::move(config), std::move(bucket), std::move(id)));

    RETURN_IF_ERROR_RESULT(fs->init());

    return fs;

}

S3FileSystem::S3FileSystem(S3Conf s3_conf, std::string id)

        : RemoteFileSystem(s3_conf.prefix, std::move(id), FileSystemType::S3),

          _bucket(std::move(s3_conf.bucket)),

          _prefix(std::move(s3_conf.prefix)),

          _client(std::make_shared<ObjClientHolder>(std::move(s3_conf.client_conf))) {

    // FIXME(plat1ko): Normalize prefix

    // remove the first and last '/'

    if (!_prefix.empty()) {

        size_t start = _prefix.find_first_not_of('/');

        if (start == std::string::npos) {

            _prefix = "";

        } else {

            size_t end = _prefix.find_last_not_of('/');

            if (start > 0 || end < _prefix.size() - 1) {

                _prefix = _prefix.substr(start, end - start + 1);

            }

        }

    }

}

S3FileSystem::S3FileSystem(ObjClientConfig config, std::string bucket, std::string id)

        : RemoteFileSystem("", std::move(id), FileSystemType::S3),

          _bucket(std::move(bucket)),

          _native_azure_location(std::holds_alternative<AzureClientConfig>(config.backend)),

          _client(std::make_shared<ObjClientHolder>(std::move(config))) {}

Status S3FileSystem::init() {

    return _client->init();

}

Result<std::string> S3FileSystem::_get_key(const Path& full_path) const {

    if (_native_azure_location) {

        // Typed Azure access keeps the location FE got from the table, e.g.

        // abfss://c1@acct.dfs.core.windows.net/db/t/f.parquet, which S3URI cannot parse.

        auto uri = ObjStorageUri::parse_azure(full_path.native());

        if (!uri.has_value()) {

            return ResultError(std::move(uri).error());

        }

        return std::move(uri.value().key);

    }

    // FIXME(plat1ko): Check bucket in full path and support relative path

    S3URI uri(full_path.native());

    RETURN_IF_ERROR_RESULT(uri.parse());

    return uri.get_key();

}

S3FileSystem::~S3FileSystem() = default;

Status S3FileSystem::create_file_impl(const Path& file, FileWriterPtr* writer,

                                      const FileWriterOptions* opts) {

    auto client = _client->get();

    CHECK_S3_CLIENT(client);

    auto key = DORIS_TRY(_get_key(file));

    *writer = std::make_unique<S3FileWriter>(_client, _bucket, std::move(key), opts);

    return Status::OK();

}

Status S3FileSystem::open_file_internal(const Path& file, FileReaderSPtr* reader,

                                        const FileReaderOptions& opts) {

    TEST_SYNC_POINT_CALLBACK("S3FileSystem::open_file_internal", &file, &opts);

    auto key = DORIS_TRY(_get_key(file));

    *reader = DORIS_TRY(S3FileReader::create(_client, _bucket, key, opts.file_size, nullptr));

    return Status::OK();

}

Status S3FileSystem::create_directory_impl(const Path& dir, bool failed_if_exists) {

    return Status::OK();

}

Status S3FileSystem::delete_file_impl(const Path& file) {

    auto client = _client->get();

    CHECK_S3_CLIENT(client);

    auto key = DORIS_TRY(_get_key(file));

    auto resp = client->delete_object({.bucket = _bucket, .key = key});

    if (resp.status.code == ErrorCode::OK || resp.status.code == ErrorCode::NOT_FOUND) {

        return Status::OK();

    }

    return std::move(Status(resp.status.code, std::move(resp.status.msg))

                             .append(fmt::format("failed to delete file {}", full_s3_path(key))));

}

Status S3FileSystem::delete_directory_impl(const Path& dir) {

    auto client = _client->get();

    CHECK_S3_CLIENT(client);

    auto prefix = DORIS_TRY(_get_key(dir));

    if (!prefix.empty() && prefix.back() != '/') {

        prefix.push_back('/');

    }

    ObjStoragePath delete_path {

            .path = full_s3_path(prefix),

            .bucket = _bucket,

            .prefix = prefix,

    };

    auto resp = delete_objects_recursively(std::move(client), delete_path);

    return {resp.status.code, std::move(resp.status.msg)};

}

Status S3FileSystem::batch_delete_impl(const std::vector<Path>& remote_files) {

    auto client = _client->get();

    CHECK_S3_CLIENT(client);

    // `DeleteObjectsRequest` can only contain 1000 keys at most.

    constexpr size_t max_delete_batch = 1000;

    auto path_iter = remote_files.begin();

    do {

        std::vector<std::string> objects;

        auto path_begin = path_iter;

        for (; path_iter != remote_files.end() && (path_iter - path_begin < max_delete_batch);

             ++path_iter) {

            auto key = DORIS_TRY(_get_key(*path_iter));

            objects.emplace_back(std::move(key));

        }

        if (objects.empty()) {

            return Status::OK();

        }

        // clang-format off

        if (auto resp = client->delete_objects( {.bucket = _bucket,}, std::move(objects)); resp.status.code != ErrorCode::OK) {

            return {resp.status.code, std::move(resp.status.msg)};

        }

        // clang-format on

    } while (path_iter != remote_files.end());

    return Status::OK();

}

Status S3FileSystem::exists_impl(const Path& path, bool* res) const {

    auto client = _client->get();

    CHECK_S3_CLIENT(client);

    auto key = DORIS_TRY(_get_key(path));

    VLOG_DEBUG << "key:" << key << " path:" << path;

    auto resp = client->head_object({.bucket = _bucket, .key = key});

    if (resp.resp.status.code == ErrorCode::OK) {

        *res = true;

    } else if (resp.resp.status.code == ErrorCode::NOT_FOUND) {

        *res = false;

    } else {

        return std::move(

                Status(resp.resp.status.code, std::move(resp.resp.status.msg))

                        .append(fmt::format(" failed to check exists {}", full_s3_path(key))));

    }

    return Status::OK();

}

Status S3FileSystem::file_size_impl(const Path& file, int64_t* file_size) const {

    auto key = DORIS_TRY(_get_key(file));

    *file_size = DORIS_TRY(_client->object_file_size(_bucket, key));

    return Status::OK();

}

Status S3FileSystem::list_impl(const Path& dir, bool only_file, std::vector<FileInfo>* files,

                               bool* exists) {

    // For object storage, this path is always not exist.

    // So we ignore this property and set exists to true.

    *exists = true;

    auto client = _client->get();

    CHECK_S3_CLIENT(client);

    auto prefix = DORIS_TRY(_get_key(dir));

    if (!prefix.empty() && prefix.back() != '/') {

        prefix.push_back('/');

    }

    std::vector<ObjectMeta> objects;

    auto resp = client->list_objects({.bucket = _bucket, .prefix = prefix}, &objects);

    if (!resp.ok()) {

        files->clear();

        return {resp.status.code, std::move(resp.status.msg)};

    }

    files->reserve(files->size() + objects.size());

    for (auto& obj : objects) {

        obj.key.erase(0, prefix.size());

        bool is_dir = obj.key.empty() || obj.key.back() == '/';

        files->emplace_back(FileInfo {

                .file_name = std::move(obj.key), .file_size = obj.size, .is_file = !is_dir});

    }

    return Status::OK();

}

Status S3FileSystem::rename_impl(const Path& orig_name, const Path& new_name) {

    return Status::NotSupported("S3FileSystem::rename_impl");

}

Status S3FileSystem::upload_impl(const Path& local_file, const Path& remote_file) {

    auto client = _client->get();

    CHECK_S3_CLIENT(client);

    auto key = DORIS_TRY(_get_key(remote_file));

    auto start = std::chrono::steady_clock::now();

    FileWriterPtr obj_writer;

    RETURN_IF_ERROR(create_file_impl(key, &obj_writer, nullptr));

    FileReaderSPtr local_reader;

    RETURN_IF_ERROR(io::global_local_filesystem()->open_file(local_file, &local_reader));

    size_t local_buffer_size = config::s3_file_system_local_upload_buffer_size;

    std::unique_ptr<char[]> write_buffer =

            std::make_unique_for_overwrite<char[]>(local_buffer_size);

    size_t cur_read = 0;

    while (cur_read < local_reader->size()) {

        size_t bytes_read = 0;

        RETURN_IF_ERROR(local_reader->read_at(

                cur_read, Slice {write_buffer.get(), local_buffer_size}, &bytes_read));

        RETURN_IF_ERROR(obj_writer->append({write_buffer.get(), bytes_read}));

        cur_read += bytes_read;

    }

    RETURN_IF_ERROR(obj_writer->close());

    auto duration = std::chrono::duration<float>(std::chrono::steady_clock::now() - start);

    auto size = local_reader->size();

    LOG(INFO) << "Upload " << local_file.native() << " to " << full_s3_path(key)

              << ", duration=" << duration.count() << ", bytes=" << size;

    return Status::OK();

}

Status S3FileSystem::batch_upload_impl(const std::vector<Path>& local_files,

                                       const std::vector<Path>& remote_files) {

    auto client = _client->get();

    CHECK_S3_CLIENT(client);

    if (local_files.size() != remote_files.size()) {

        return Status::InvalidArgument("local_files.size({}) != remote_files.size({})",

                                       local_files.size(), remote_files.size());

    }

    std::vector<FileWriterPtr> obj_writers(local_files.size());

    auto upload_task = [&, this](size_t idx) {

        SCOPED_ATTACH_TASK(ExecEnv::GetInstance()->s3_file_buffer_tracker());

        const auto& local_file = local_files[idx];

        const auto& remote_file = remote_files[idx];

        auto& obj_writer = obj_writers[idx];

        auto key = DORIS_TRY(_get_key(remote_file));

        LOG(INFO) << "Start to upload " << local_file.native() << " to " << full_s3_path(key);

        RETURN_IF_ERROR(create_file_impl(key, &obj_writer, nullptr));

        FileReaderSPtr local_reader;

        RETURN_IF_ERROR(io::global_local_filesystem()->open_file(local_file, &local_reader));

        size_t local_buffer_size = config::s3_file_system_local_upload_buffer_size;

        std::unique_ptr<char[]> write_buffer =

                std::make_unique_for_overwrite<char[]>(local_buffer_size);

        size_t cur_read = 0;

        while (cur_read < local_reader->size()) {

            size_t bytes_read = 0;

            RETURN_IF_ERROR(local_reader->read_at(

                    cur_read, Slice {write_buffer.get(), local_buffer_size}, &bytes_read));

            RETURN_IF_ERROR((*obj_writer).append({write_buffer.get(), bytes_read}));

            cur_read += bytes_read;

        }

        RETURN_IF_ERROR((*obj_writer).close());

        return Status::OK();

    };

    Status s = Status::OK();

    std::vector<std::future<Status>> futures;

    for (int i = 0; i < local_files.size(); ++i) {

        auto task = std::make_shared<std::packaged_task<Status(size_t idx)>>(upload_task);

        futures.emplace_back(task->get_future());

        auto st = ExecEnv::GetInstance()->s3_file_system_thread_pool()->submit_func(

                [t = std::move(task), idx = i]() mutable { (*t)(idx); });

        // We shouldn't return immediately since the previous submitted tasks might still be running in the thread pool

        if (!st.ok()) {

            s = st;

            break;

        }

    }

    for (auto&& f : futures) {

        auto cur_s = f.get();

        if (!cur_s.ok()) {

            s = std::move(cur_s);

        }

    }

    return s;

}

Status S3FileSystem::download_impl(const Path& remote_file, const Path& local_file) {

    auto client = _client->get();

    CHECK_S3_CLIENT(client);

    auto key = DORIS_TRY(_get_key(remote_file));

    int64_t size;

    RETURN_IF_ERROR(file_size(remote_file, &size));

    std::unique_ptr<char[]> buf = std::make_unique_for_overwrite<char[]>(size);

    size_t bytes_read = 0;

    // clang-format off

    auto resp = client->get_object( {.bucket = _bucket, .key = key,},

            buf.get(), 0, size, &bytes_read);

    // clang-format on

    if (resp.status.code != ErrorCode::OK) {

        return {resp.status.code, std::move(resp.status.msg)};

    }

    Aws::OFStream local_file_s;

    local_file_s.open(local_file, std::ios::out | std::ios::binary);

    if (local_file_s.good()) {

        local_file_s << StringViewStream(buf.get(), size).rdbuf();

    } else {

        return localfs_error(errno, fmt::format("failed to write file {}", local_file.native()));

    }

    return Status::OK();

}

// oss has public endpoint and private endpoint, is_public_endpoint determines

// whether to return a public endpoint.

std::string S3FileSystem::generate_presigned_url(const Path& path, int64_t expiration_secs,

                                                 bool is_public_endpoint) const {

    std::string key = fmt::format("{}/{}", _prefix, path.native());

    std::shared_ptr<ObjStorageClient> client;

    // Read the conf and the client from one snapshot so a concurrent reset cannot pair the

    // endpoint of one generation with the client of another.

    auto state = _client->state();

    if (is_public_endpoint && state->config.endpoint().ends_with(OSS_PRIVATE_ENDPOINT_SUFFIX)) {

        // Only S3-protocol backends have an OSS private endpoint. Their file systems are

        // created from an S3Conf (the typed constructor is used for Azure only), so the holder

        // keeps the legacy conf that the public-endpoint client is derived from.

        DORIS_CHECK(state->legacy_conf.has_value());

        auto new_s3_conf = *state->legacy_conf;

        new_s3_conf.endpoint.erase(new_s3_conf.endpoint.size() - OSS_PRIVATE_ENDPOINT_SUFFIX.size(),

                                   LEN_OF_OSS_PRIVATE_SUFFIX);

        auto client_result = S3ClientFactory::instance().create(new_s3_conf);

        if (!client_result) {

            LOG(WARNING) << "failed to create S3 client for presigned URL: "

                         << client_result.error();

            return {};

        }

        client = std::move(client_result).value();

    } else {

        client = state->client;

    }

    return client->generate_presigned_url({.bucket = _bucket, .key = key, .prefix = ""},

                                          expiration_secs);

}

} // namespace doris::io