// 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 "runtime/small_file_mgr.h"

// IWYU pragma: no_include <bthread/errno.h>

#include <absl/strings/str_split.h>

#include <errno.h> // IWYU pragma: keep

#include <gen_cpp/HeartbeatService_types.h>

#include <gen_cpp/Types_types.h>

#include <glog/logging.h>

#include <stdint.h>

#include <stdio.h>

#include <unistd.h>

#include <cstring>

#include <memory>

#include <sstream>

#include <utility>

#include <vector>

#include "common/metrics/doris_metrics.h"

#include "common/metrics/metrics.h"

#include "common/status.h"

#include "io/fs/file_system.h"

#include "io/fs/local_file_system.h"

#include "runtime/exec_env.h"

#include "service/http/http_client.h"

#include "util/md5.h"

#include "util/string_util.h"

namespace doris {

DEFINE_GAUGE_METRIC_PROTOTYPE_2ARG(small_file_cache_count, MetricUnit::NOUNIT);

SmallFileMgr::SmallFileMgr(ExecEnv* env, const std::string& local_path)

        : _exec_env(env), _local_path(local_path) {

    REGISTER_HOOK_METRIC(small_file_cache_count, [this]() {

        // std::lock_guard<std::mutex> l(_lock);

        return _file_cache.size();

    });

}

SmallFileMgr::~SmallFileMgr() {

    DEREGISTER_HOOK_METRIC(small_file_cache_count);

}

Status SmallFileMgr::init() {

    RETURN_IF_ERROR(_load_local_files());

    return Status::OK();

}

Status SmallFileMgr::_load_local_files() {

    RETURN_IF_ERROR(io::global_local_filesystem()->create_directory(_local_path));

    auto scan_cb = [this](const io::FileInfo& file) {

        if (!file.is_file) {

            return true;

        }

        auto st = _load_single_file(_local_path, file.file_name);

        if (!st.ok()) {

            LOG(WARNING) << "load small file failed: " << st;

        }

        return true;

    };

    RETURN_IF_ERROR(io::global_local_filesystem()->iterate_directory(_local_path, scan_cb));

    return Status::OK();

}

Status SmallFileMgr::_load_single_file(const std::string& path, const std::string& file_name) {

    // file name format should be like:

    // file_id.md5

    std::vector<std::string> parts = absl::StrSplit(file_name, ".");

    if (parts.size() != 2) {

        return Status::InternalError("Not a valid file name: {}", file_name);

    }

    int64_t file_id = std::stol(parts[0]);

    std::string md5 = parts[1];

    if (_file_cache.find(file_id) != _file_cache.end()) {

        return Status::InternalError("File with same id is already been loaded: {}", file_id);

    }

    std::string file_md5;

    RETURN_IF_ERROR(io::global_local_filesystem()->md5sum(path + "/" + file_name, &file_md5));

    if (file_md5 != md5) {

        return Status::InternalError("Invalid md5 of file: {}", file_name);

    }

    CacheEntry entry;

    entry.path = path + "/" + file_name;

    entry.md5 = file_md5;

    _file_cache.emplace(file_id, entry);

    return Status::OK();

}

Status SmallFileMgr::get_file(int64_t file_id, const std::string& md5, std::string* file_path) {

    std::unique_lock<std::mutex> l(_lock);

    // find in cache

    auto it = _file_cache.find(file_id);

    if (it != _file_cache.end()) {

        // find the cached file, check it

        CacheEntry& entry = it->second;

        Status st = _check_file(entry, md5);

        if (!st.ok()) {

            // check file failed, we should remove this cache and download it from FE again

            if (remove(entry.path.c_str()) != 0) {

                return Status::InternalError("failed to remove file: {}, err: {}", file_id,

                                             std::strerror(errno));

            }

            _file_cache.erase(it);

        } else {

            // check ok, return the path

            *file_path = entry.path;

            return Status::OK();

        }

    }

    // file not found in cache. download it from FE

    RETURN_IF_ERROR(_download_file(file_id, md5, file_path));

    return Status::OK();

}

Status SmallFileMgr::_check_file(const CacheEntry& entry, const std::string& md5) {

    bool exists;

    RETURN_IF_ERROR(io::global_local_filesystem()->exists(entry.path, &exists));

    if (!exists) {

        return Status::InternalError("file not exist: {}", entry.path);

    }

    if (!iequal(md5, entry.md5)) {

        return Status::InternalError("invalid MD5 of file: {}", entry.path);

    }

    return Status::OK();

}

Status SmallFileMgr::_download_file(int64_t file_id, const std::string& md5,

                                    std::string* file_path) {

    std::stringstream ss;

    ss << _local_path << "/" << file_id << ".tmp";

    std::string tmp_file = ss.str();

    bool should_delete = true;

    auto fp_closer = [&tmp_file, &should_delete](FILE* fp) {

        fclose(fp);

        if (should_delete) remove(tmp_file.c_str());

    };

    std::unique_ptr<FILE, decltype(fp_closer)> fp(fopen(tmp_file.c_str(), "w"), fp_closer);

    if (fp == nullptr) {

        LOG(WARNING) << "fail to open file, file=" << tmp_file;

        return Status::InternalError("fail to open file");

    }

    ClusterInfo* cluster_info = _exec_env->cluster_info();

    // Small file download is the only BE→FE path that uses HTTP (not Thrift/RPC).

    // master_fe_http_port is set to https_port when enable_https=true (see HeartbeatMgr).

    // The ~1ms fallback overhead is acceptable; small file downloads are infrequent.

    const std::string host_port = cluster_info->master_fe_addr.hostname + ":" +

                                  std::to_string(cluster_info->master_fe_http_port);

    const std::string query = "/api/get_small_file?file_id=" + std::to_string(file_id) +

                              "&token=" + cluster_info->token;

    Status status;

    Md5Digest digest;

    auto download_cb = [&status, &tmp_file, &fp, &digest](const void* data, size_t length) {

        digest.update(data, length);

        auto res = fwrite(data, length, 1, fp.get());

        if (res != 1) {

            LOG(WARNING) << "fail to write data to file, file=" << tmp_file

                         << ", error=" << ferror(fp.get());

            status = Status::InternalError("fail to write data when download");

            return false;

        }

        return true;

    };

    std::string url = "http://" + host_port + query;

    LOG(INFO) << "download file from: " << url;

    HttpClient client;

    RETURN_IF_ERROR(client.init(url));

    Status execute_status = client.execute(download_cb);

    if (!execute_status.ok()) {

        rewind(fp.get());

        if (ftruncate(fileno(fp.get()), 0) != 0) {

            LOG(WARNING) << "fail to truncate temp file for https retry, errno=" << errno;

        }

        status = Status::OK();

        digest = Md5Digest();

        url = "https://" + host_port + query;

        LOG(INFO) << "HTTP failed, retrying with HTTPS: " << url;

        HttpClient https_client;

        RETURN_IF_ERROR(https_client.init(url));

        // Skip TLS cert verification: internal cluster traffic only; file integrity

        // is guaranteed independently by MD5 checksum verification below.

        https_client.use_untrusted_ssl();

        execute_status = https_client.execute(download_cb);

    }

    RETURN_IF_ERROR(execute_status);

    RETURN_IF_ERROR(status);

    digest.digest();

    if (!iequal(digest.hex(), md5)) {

        LOG(WARNING) << "file's checksum is not equal, download: " << digest.hex()

                     << ", expected: " << md5 << ", file: " << file_id;

        return Status::InternalError("download with invalid md5");

    }

    // close this file

    should_delete = false;

    fp.reset();

    // rename temporary file to library file

    std::stringstream real_ss;

    real_ss << _local_path << "/" << file_id << "." << md5;

    std::string real_file_path = real_ss.str();

    auto ret = rename(tmp_file.c_str(), real_file_path.c_str());

    if (ret != 0) {

        char buf[64];

        LOG(WARNING) << "fail to rename file from=" << tmp_file << ", to=" << real_file_path

                     << ", errno=" << errno << ", errmsg=" << strerror_r(errno, buf, 64);

        remove(tmp_file.c_str());

        remove(real_file_path.c_str());

        return Status::InternalError("fail to rename file");

    }

    // add to file cache

    CacheEntry entry;

    entry.path = real_file_path;

    entry.md5 = md5;

    _file_cache.emplace(file_id, entry);

    *file_path = real_file_path;

    LOG(INFO) << "finished to download file: " << file_path;

    return Status::OK();

}

} // end namespace doris