// 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 "common/config.h"

#include "common/logging.h"

#include "common/status.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_common.h"

#include "io/fs/s3_file_reader.h"

#include "io/fs/s3_file_writer.h"

#include "io/fs/s3_obj_storage_client.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

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

    // 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();

}

} // namespace

ObjClientHolder::ObjClientHolder(S3ClientConf conf) : _conf(std::move(conf)) {}

ObjClientHolder::~ObjClientHolder() = default;

Status ObjClientHolder::init() {

    _client = S3ClientFactory::instance().create(_conf);

    if (!_client) {

        return Status::InvalidArgument("failed to init s3 client with conf {}", _conf.to_string());

    }

    return Status::OK();

}

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

    S3ClientConf reset_conf;

    {

        std::shared_lock lock(_mtx);

        if (conf.get_hash() == _conf.get_hash()) {

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

        }

        reset_conf = _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.role_arn = conf.role_arn;

        reset_conf.external_id = conf.external_id;

        reset_conf.cred_provider_type = conf.cred_provider_type;

        // Should check endpoint here?

    }

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

    if (!client) {

        return Status::InvalidArgument("failed to init s3 client with conf {}", conf.to_string());

    }

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

    {

        std::lock_guard lock(_mtx);

        _client = std::move(client);

        _conf = std::move(reset_conf);

    }

    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("{}/{}/{}", _conf.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;

}

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);

            }

        }

    }

}

Status S3FileSystem::init() {

    return _client->init();

}

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) {

    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('/');

    }

    auto resp = client->delete_objects_recursively({

            .path = full_s3_path(prefix),

            .bucket = _bucket,

            .prefix = prefix,

    });

    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('/');

    }

    // clang-format off

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

    // clang-format on

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

        for (auto&& file : *files) {

            file.file_name.erase(0, prefix.size());

        }

    }

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

}

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;

    if (is_public_endpoint &&

        _client->s3_client_conf().endpoint.ends_with(OSS_PRIVATE_ENDPOINT_SUFFIX)) {

        auto new_s3_conf = _client->s3_client_conf();

        new_s3_conf.endpoint.erase(

                _client->s3_client_conf().endpoint.size() - OSS_PRIVATE_ENDPOINT_SUFFIX.size(),

                LEN_OF_OSS_PRIVATE_SUFFIX);

        client = S3ClientFactory::instance().create(new_s3_conf);

    } else {

        client = _client->get();

    }

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

                                          _client->s3_client_conf());

}

} // namespace doris::io