#include "client.h"

#include <glog/logging.h>

#include <algorithm>
#include <cassert>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <iomanip>
#include <optional>
#include <ranges>
#include <stdexcept>
#include <string>
#include <thread>

#include "transfer_engine.h"
#include "transfer_task.h"
#include "transport/transport.h"
#include "config.h"
#include "types.h"
#include "local_hot_cache.h"

namespace mooncake {

[[nodiscard]] size_t CalculateSliceSize(const std::vector<Slice>& slices) {
    size_t slice_size = 0;
    for (const auto& slice : slices) {
        slice_size += slice.size;
    }
    return slice_size;
}

[[nodiscard]] size_t CalculateSliceSize(std::span<const Slice> slices) {
    size_t slice_size = 0;
    for (const auto& slice : slices) {
        slice_size += slice.size;
    }
    return slice_size;
}

Client::Client(const std::string& local_hostname,
               const std::string& metadata_connstring)
    : metrics_(ClientMetric::Create()),
      master_client_(metrics_ ? &metrics_->master_client_metric : nullptr),
      local_hostname_(local_hostname),
      metadata_connstring_(metadata_connstring),
      write_thread_pool_(2) {
    client_id_ = generate_uuid();
    LOG(INFO) << "client_id=" << client_id_;

    if (metrics_) {
        if (metrics_->GetReportingInterval() > 0) {
            LOG(INFO) << "Client metrics enabled with reporting thread started "
                         "(interval: "
                      << metrics_->GetReportingInterval() << "s)";
        } else {
            LOG(INFO)
                << "Client metrics enabled but reporting disabled (interval=0)";
        }
    } else {
        LOG(INFO) << "Client metrics disabled (set MC_STORE_CLIENT_METRIC=1 to "
                     "enable)";
    }
}

Client::~Client() {
    // Make a copy of mounted_segments_ to avoid modifying while iterating
    std::vector<Segment> segments_to_unmount;
    {
        std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
        segments_to_unmount.reserve(mounted_segments_.size());
        for (auto& entry : mounted_segments_) {
            segments_to_unmount.emplace_back(entry.second);
        }
    }

    for (auto& segment : segments_to_unmount) {
        auto result =
            UnmountSegment(reinterpret_cast<void*>(segment.base), segment.size);
        if (!result) {
            LOG(ERROR) << "Failed to unmount segment: "
                       << toString(result.error());
        }
    }

    // Clear any remaining segments
    {
        std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
        mounted_segments_.clear();
    }

    // Stop hot cache handler and hot cache
    hot_cache_handler_.reset();
    hot_cache_.reset();

    // Stop ping thread only after no need to contact master anymore
    if (ping_running_) {
        ping_running_ = false;
        if (ping_thread_.joinable()) {
            ping_thread_.join();
        }
    }
}

static std::optional<bool> get_auto_discover() {
    const char* ev_ad = std::getenv("MC_MS_AUTO_DISC");
    if (ev_ad) {
        int iv = std::stoi(ev_ad);
        if (iv == 1) {
            LOG(INFO) << "auto discovery set by env MC_MS_AUTO_DISC";
            return true;
        } else if (iv == 0) {
            LOG(INFO) << "auto discovery not set by env MC_MS_AUTO_DISC";
            return false;
        } else {
            LOG(WARNING)
                << "invalid MC_MS_AUTO_DISC value: " << ev_ad
                << ", should be 0 or 1, using default: auto discovery not set";
        }
    }
    return std::nullopt;
}

static inline void ltrim(std::string& s) {
    s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](unsigned char ch) {
                return !std::isspace(ch);
            }));
}

static inline void rtrim(std::string& s) {
    s.erase(std::find_if(s.rbegin(), s.rend(),
                         [](unsigned char ch) { return !std::isspace(ch); })
                .base(),
            s.end());
}

static std::vector<std::string> get_auto_discover_filters(bool auto_discover) {
    std::vector<std::string> whitelst_filters;
    char* ev_ad = std::getenv("MC_MS_FILTERS");
    if (ev_ad) {
        if (!auto_discover) {
            LOG(WARNING)
                << "auto discovery not set, but find whitelist filters: "
                << ev_ad;
            return whitelst_filters;
        }
        LOG(INFO) << "whitelist filters: " << ev_ad;
        char delimiter = ',';
        char* end = ev_ad + std::strlen(ev_ad);
        char *start = ev_ad, *pos = ev_ad;
        while ((pos = std::find(start, end, delimiter)) != end) {
            std::string str(start, pos);
            ltrim(str);
            rtrim(str);
            whitelst_filters.emplace_back(std::move(str));
            start = pos + 1;
        }
        if (start != (end + 1)) {
            std::string str(start, end);
            ltrim(str);
            rtrim(str);
            whitelst_filters.emplace_back(std::move(str));
        }
    }
    return whitelst_filters;
}

tl::expected<void, ErrorCode> CheckRegisterMemoryParams(const void* addr,
                                                        size_t length) {
    if (addr == nullptr) {
        LOG(ERROR) << "addr is nullptr";
        return tl::unexpected(ErrorCode::INVALID_PARAMS);
    }
    if (length == 0) {
        LOG(ERROR) << "length is 0";
        return tl::unexpected(ErrorCode::INVALID_PARAMS);
    }
    // Tcp is not limited by max_mr_size, but we ignore it for now.
    auto max_mr_size = globalConfig().max_mr_size;  // Max segment size
    if (length > max_mr_size) {
        LOG(ERROR) << "length " << length
                   << " is larger than max_mr_size: " << max_mr_size;
        return tl::unexpected(ErrorCode::INVALID_PARAMS);
    }
    return {};
}

ErrorCode Client::ConnectToMaster(const std::string& master_server_entry) {
    if (master_server_entry.find("etcd://") == 0) {
        std::string etcd_entry = master_server_entry.substr(strlen("etcd://"));

        // Get master address from etcd
        auto err = master_view_helper_.ConnectToEtcd(etcd_entry);
        if (err != ErrorCode::OK) {
            LOG(ERROR) << "Failed to connect to etcd";
            return err;
        }
        std::string master_address;
        ViewVersionId master_version = 0;
        err = master_view_helper_.GetMasterView(master_address, master_version);
        if (err != ErrorCode::OK) {
            LOG(ERROR) << "Failed to get master address";
            return err;
        }

        err = master_client_.Connect(master_address);
        if (err != ErrorCode::OK) {
            LOG(ERROR) << "Failed to connect to master";
            return err;
        }

        // Start ping thread to monitor master health and trigger remount if
        // needed.
        ping_running_ = true;
        bool is_ha_mode = true;
        std::string current_master_address = master_address;
        ping_thread_ = std::thread([this, is_ha_mode,
                                    current_master_address]() mutable {
            this->PingThreadMain(is_ha_mode, std::move(current_master_address));
        });

        return ErrorCode::OK;
    } else {
        auto err = master_client_.Connect(master_server_entry);
        if (err != ErrorCode::OK) {
            return err;
        }
        // Non-HA mode also enables heartbeat/ping
        ping_running_ = true;
        bool is_ha_mode = false;
        std::string current_master_address = master_server_entry;
        ping_thread_ = std::thread([this, is_ha_mode,
                                    current_master_address]() mutable {
            this->PingThreadMain(is_ha_mode, std::move(current_master_address));
        });
        return ErrorCode::OK;
    }
}

ErrorCode Client::InitTransferEngine(
    const std::string& local_hostname, const std::string& metadata_connstring,
    const std::string& protocol,
    const std::optional<std::string>& device_names) {
    if (!te_initialized_) {
        // get auto_discover and filters from env
        std::optional<bool> env_auto_discover = get_auto_discover();
        bool auto_discover = false;
        if (env_auto_discover.has_value()) {
            // Use user-specified auto-discover setting
            auto_discover = env_auto_discover.value();
        } else {
            // Enable auto-discover for RDMA if no devices are specified
            if (protocol == "rdma" && !device_names.has_value()) {
                LOG(INFO)
                    << "Set auto discovery ON by default for RDMA protocol, "
                       "since no "
                       "device names provided";
                auto_discover = true;
            }
        }
        transfer_engine_->setAutoDiscover(auto_discover);

        auto [hostname, port] = parseHostNameWithPort(local_hostname);
        int rc = transfer_engine_->init(metadata_connstring, local_hostname,
                                        hostname, port);
        if (rc != 0) {
            LOG(ERROR) << "Failed to initialize transfer engine, rc=" << rc;
            return ErrorCode::INTERNAL_ERROR;
        }

        if (auto_discover) {
            LOG(INFO)
                << "Transfer engine auto discovery is enabled for protocol: "
                << protocol;
            auto filters = get_auto_discover_filters(auto_discover);
            transfer_engine_->setWhitelistFilters(std::move(filters));
        } else {
            LOG(INFO)
                << "Transfer engine auto discovery is disabled for protocol: "
                << protocol;

            Transport* transport = nullptr;

            if (protocol == "rdma") {
                if (!device_names.has_value() || device_names->empty()) {
                    LOG(ERROR)
                        << "RDMA protocol requires device names when auto "
                           "discovery is disabled";
                    return ErrorCode::INVALID_PARAMS;
                }

                LOG(INFO) << "Using specified RDMA devices: "
                          << device_names.value();

                std::vector<std::string> devices =
                    splitString(device_names.value(), ',', /*skip_empty=*/true);

                // Manually discover topology with specified devices only
                auto topology = transfer_engine_->getLocalTopology();
                if (topology) {
                    topology->discover(devices);
                    LOG(INFO) << "Topology discovery complete with specified "
                                 "devices. Found "
                              << topology->getHcaList().size() << " HCAs";
                }

                transport = transfer_engine_->installTransport("rdma", nullptr);
                if (!transport) {
                    LOG(ERROR)
                        << "Failed to install RDMA transport with specified "
                           "devices";
                    return ErrorCode::INTERNAL_ERROR;
                }
            } else if (protocol == "tcp") {
                if (device_names.has_value()) {
                    LOG(WARNING)
                        << "TCP protocol does not use device names, ignoring";
                }

                try {
                    transport =
                        transfer_engine_->installTransport("tcp", nullptr);
                } catch (std::exception& e) {
                    LOG(ERROR)
                        << "tcp_transport_install_failed error_message=\""
                        << e.what() << "\"";
                    return ErrorCode::INTERNAL_ERROR;
                }

                if (!transport) {
                    LOG(ERROR) << "Failed to install TCP transport";
                    return ErrorCode::INTERNAL_ERROR;
                }
            } else if (protocol == "ascend") {
                if (device_names.has_value()) {
                    LOG(WARNING) << "Ascend protocol does not use device "
                                    "names, ignoring";
                }
                try {
                    transport =
                        transfer_engine_->installTransport("ascend", nullptr);
                } catch (std::exception& e) {
                    LOG(ERROR)
                        << "ascend_transport_install_failed error_message=\""
                        << e.what() << "\"";
                    return ErrorCode::INTERNAL_ERROR;
                }

                if (!transport) {
                    LOG(ERROR) << "Failed to install Ascend transport";
                    return ErrorCode::INTERNAL_ERROR;
                }
            } else if (protocol == "ub") {
                try {
                    transport = transfer_engine_->installTransport("ub", nullptr);
                } catch (std::exception& e) {
                    LOG(ERROR) << "ub_transport_install_failed error_message=\"" << e.what() << "\"";
                    return ErrorCode::INTERNAL_ERROR;
                }

                if (!transport) {
                    LOG(ERROR) << "Failed to install ub transport";
                    return ErrorCode::INTERNAL_ERROR;
                }
            } else {
                LOG(ERROR) << "unsupported_protocol protocol=" << protocol;
                return ErrorCode::INVALID_PARAMS;
            }
        }
    }

    // Initialize TransferSubmitter after transfer engine is ready
    // Keep using logical local_hostname for name-based behaviors; endpoint is
    // used separately where needed.
    transfer_submitter_ = std::make_unique<TransferSubmitter>(
        *transfer_engine_, storage_backend_,
        metrics_ ? &metrics_->transfer_metric : nullptr);

    return ErrorCode::OK;
}

std::optional<std::shared_ptr<Client>> Client::Create(
    const std::string& local_hostname, const std::string& metadata_connstring,
    const std::string& protocol, const std::optional<std::string>& device_names,
    const std::string& master_server_entry,
    const std::shared_ptr<TransferEngine>& transfer_engine) {
    auto client = std::shared_ptr<Client>(
        new Client(local_hostname, metadata_connstring));

    ErrorCode err = client->ConnectToMaster(master_server_entry);
    if (err != ErrorCode::OK) {
        return std::nullopt;
    }

    // Initialize storage backend if storage_root_dir is valid
    auto response = client->master_client_.GetFsdir();
    if (!response) {
        LOG(ERROR) << "Failed to get fsdir from master";
    } else if (response.value().empty()) {
        LOG(INFO) << "Storage root directory is not set. persisting data is "
                     "disabled.";
    } else {
        auto dir_string = response.value();
        size_t pos = dir_string.find_last_of('/');
        if (pos != std::string::npos) {
            std::string storage_root_dir = dir_string.substr(0, pos);
            std::string fs_subdir = dir_string.substr(pos + 1);
            LOG(INFO) << "Storage root directory is: " << storage_root_dir;
            LOG(INFO) << "Fs subdir is: " << fs_subdir;
            // Initialize storage backend
            client->PrepareStorageBackend(storage_root_dir, fs_subdir);
        } else {
            LOG(ERROR) << "Invalid fsdir format: " << dir_string;
        }
    }

    // Initialize transfer engine
    if (transfer_engine == nullptr) {
        client->transfer_engine_ = std::make_shared<TransferEngine>();
    } else {
        client->transfer_engine_ = transfer_engine;
        client->te_initialized_ = true;
        LOG(INFO) << "Use exist transfer engine instance";
    }
    err = client->InitTransferEngine(local_hostname, metadata_connstring,
                                     protocol, device_names);
    if (err != ErrorCode::OK) {
        LOG(ERROR) << "Failed to initialize transfer engine";
        return std::nullopt;
    }

    // Initialize local hot cache
    err = client->InitLocalHotCache();
    if (err != ErrorCode::OK) {
        LOG(ERROR) << "Failed to initialize local hot cache";
    }

    return client;
}

tl::expected<void, ErrorCode> Client::Get(const std::string& object_key,
                                          std::vector<Slice>& slices) {
    auto query_result = Query(object_key);
    if (!query_result) {
        return tl::unexpected(query_result.error());
    }
    return Get(object_key, query_result.value(), slices);
}

std::vector<tl::expected<void, ErrorCode>> Client::BatchGet(
    const std::vector<std::string>& object_keys,
    std::unordered_map<std::string, std::vector<Slice>>& slices) {
    auto batched_query_results = BatchQuery(object_keys);

    // If any queries failed, return error results immediately for failed
    // queries
    std::vector<tl::expected<void, ErrorCode>> results;
    results.reserve(object_keys.size());

    std::vector<QueryResult> valid_query_results;
    std::vector<size_t> valid_indices;
    std::vector<std::string> valid_keys;

    for (size_t i = 0; i < batched_query_results.size(); ++i) {
        if (batched_query_results[i]) {
            valid_query_results.emplace_back(batched_query_results[i].value());
            valid_indices.emplace_back(i);
            valid_keys.emplace_back(object_keys[i]);
            results.emplace_back();  // placeholder for successful results
        } else {
            results.emplace_back(
                tl::unexpected(batched_query_results[i].error()));
        }
    }

    // If we have any valid queries, process them
    if (!valid_keys.empty()) {
        std::unordered_map<std::string, std::vector<Slice>> valid_slices;
        for (const auto& key : valid_keys) {
            auto it = slices.find(key);
            if (it != slices.end()) {
                valid_slices[key] = it->second;
            }
        }

        auto valid_results =
            BatchGet(valid_keys, valid_query_results, valid_slices);

        // Merge results back
        for (size_t i = 0; i < valid_indices.size(); ++i) {
            results[valid_indices[i]] = valid_results[i];
        }
    }

    return results;
}

tl::expected<std::unordered_map<std::string, std::vector<Replica::Descriptor>>,
             ErrorCode>
Client::QueryByRegex(const std::string& str) {
    auto result = master_client_.GetReplicaListByRegex(str);
    return result;
}

tl::expected<QueryResult, ErrorCode> Client::Query(
    const std::string& object_key) {
    std::chrono::steady_clock::time_point start_time =
        std::chrono::steady_clock::now();
    auto result = master_client_.GetReplicaList(object_key);
    if (!result) {
        return tl::unexpected(result.error());
    }
    return QueryResult(
        std::move(result.value().replicas),
        start_time + std::chrono::milliseconds(result.value().lease_ttl_ms));
}

std::vector<tl::expected<QueryResult, ErrorCode>> Client::BatchQuery(
    const std::vector<std::string>& object_keys) {
    std::chrono::steady_clock::time_point start_time =
        std::chrono::steady_clock::now();
    auto response = master_client_.BatchGetReplicaList(object_keys);

    // Check if we got the expected number of responses
    if (response.size() != object_keys.size()) {
        LOG(ERROR) << "BatchQuery response size mismatch. Expected: "
                   << object_keys.size() << ", Got: " << response.size();
        // Return vector of RPC_FAIL errors
        std::vector<tl::expected<QueryResult, ErrorCode>> results;
        results.reserve(object_keys.size());
        for (size_t i = 0; i < object_keys.size(); ++i) {
            results.emplace_back(tl::unexpected(ErrorCode::RPC_FAIL));
        }
        return results;
    }
    std::vector<tl::expected<QueryResult, ErrorCode>> results;
    results.reserve(response.size());
    for (size_t i = 0; i < response.size(); ++i) {
        if (response[i]) {
            results.emplace_back(QueryResult(
                std::move(response[i].value().replicas),
                start_time + std::chrono::milliseconds(
                                 response[i].value().lease_ttl_ms)));
        } else {
            results.emplace_back(tl::unexpected(response[i].error()));
        }
    }
    return results;
}

tl::expected<void, ErrorCode> Client::Get(const std::string& object_key,
                                          const QueryResult& query_result,
                                          std::vector<Slice>& slices) {
    // Find the first complete replica
    Replica::Descriptor replica;
    ErrorCode err = FindFirstCompleteReplica(query_result.replicas, replica);
    if (err != ErrorCode::OK) {
        if (err == ErrorCode::INVALID_REPLICA) {
            LOG(ERROR) << "no_complete_replicas_found key=" << object_key;
        }
        return tl::unexpected(err);
    }

    // Check local hot cache and update replica descriptor if cache hit
    size_t cache_hits = 0;
    size_t total_blocks = 0;
    if (hot_cache_ && replica.is_memory_replica()) {
        auto& mem_desc = replica.get_memory_descriptor();
        total_blocks = mem_desc.buffer_descriptors.size();
        cache_hits = updateReplicaDescriptorFromCache(object_key, replica);
    }

    auto t0_get = std::chrono::steady_clock::now();
    err = TransferRead(replica, slices);

    // Asynchronously update local hot cache with TE transfer slices
    if (hot_cache_ && hot_cache_handler_) {
        ProcessSlicesAsync(object_key, slices, replica);
    }
    
    auto us_get = std::chrono::duration_cast<std::chrono::microseconds>(
                      std::chrono::steady_clock::now() - t0_get)
                      .count();
    if (metrics_) {
        metrics_->transfer_metric.get_latency_us.observe(us_get);
    }

    if (err != ErrorCode::OK) {
        LOG(ERROR) << "transfer_read_failed key=" << object_key;
        return tl::unexpected(err);
    }

    if (query_result.IsLeaseExpired()) {
        LOG(WARNING) << "lease_expired_before_data_transfer_completed key="
                     << object_key;
        return tl::unexpected(ErrorCode::LEASE_EXPIRED);
    }
    // Log cache hit statistics
    if (hot_cache_ && replica.is_memory_replica()) {
        VLOG(1) << "Get completed: key=" << object_key
                << " cache_hits=" << cache_hits
                << " total_blocks=" << total_blocks;
    }

    return {};
}

struct BatchGetOperation {
    std::vector<Replica::Descriptor> replicas;
    std::vector<std::vector<Slice>> batched_slices;
    std::vector<size_t> key_indexes;
    std::vector<TransferFuture> futures;
};

std::vector<tl::expected<void, ErrorCode>> Client::BatchGetWhenPreferSameNode(
    const std::vector<std::string>& object_keys,
    const std::vector<QueryResult>& query_results,
    std::unordered_map<std::string, std::vector<Slice>>& slices) {
    std::vector<tl::expected<void, ErrorCode>> results;
    results.resize(object_keys.size());

    std::unordered_map<std::string, BatchGetOperation> seg_to_op_map{};
    for (size_t i = 0; i < object_keys.size(); ++i) {
        const auto& key = object_keys[i];
        const auto& replica_list = query_results[i].replicas;
        auto slices_it = slices.find(key);
        if (slices_it == slices.end()) {
            LOG(ERROR) << "Slices not found for key: " << key;
            results[i] = tl::unexpected(ErrorCode::INVALID_PARAMS);
            continue;
        }
        Replica::Descriptor replica;
        ErrorCode err = FindFirstCompleteReplica(replica_list, replica);
        if (err != ErrorCode::OK) {
            if (err == ErrorCode::INVALID_REPLICA) {
                LOG(ERROR) << "no_complete_replicas_found key=" << key;
            }
            results[i] = tl::unexpected(err);
            continue;
        }
        if (!replica.is_memory_replica()) {
            results[i] = tl::unexpected(ErrorCode::INVALID_REPLICA);
            continue;
        }
        auto& memory_descriptor = replica.get_memory_descriptor();
        if (memory_descriptor.buffer_descriptors.empty()) {
            results[i] = tl::unexpected(ErrorCode::INVALID_REPLICA);
            continue;
        }
        auto& buffer_descriptor = memory_descriptor.buffer_descriptors[0];
        auto seg = buffer_descriptor.transport_endpoint_;
        auto& op = seg_to_op_map[seg];
        op.replicas.emplace_back(replica);
        op.batched_slices.emplace_back(slices_it->second);
        op.key_indexes.emplace_back(i);
    }
    for (auto& seg_to_op : seg_to_op_map) {
        auto& op = seg_to_op.second;
        auto future = transfer_submitter_->submit_batch(
            op.replicas, op.batched_slices, TransferRequest::READ);
        if (!future) {
            for (auto index : op.key_indexes) {
                results[index] = tl::unexpected(ErrorCode::TRANSFER_FAIL);
                LOG(ERROR) << "Failed to submit transfer operation for key: "
                           << object_keys[index];
            }
            continue;
        }
        op.futures.emplace_back(std::move(*future));
    }
    for (auto& seg_to_op : seg_to_op_map) {
        auto& op = seg_to_op.second;
        if (op.futures.empty()) {
            continue;
        }
        ErrorCode result = op.futures[0].get();
        if (result != ErrorCode::OK) {
            for (auto index : op.key_indexes) {
                results[index] = tl::unexpected(ErrorCode::TRANSFER_FAIL);
                LOG(ERROR) << "Failed to submit transfer operation for key: "
                           << object_keys[index];
            }
        } else {
            for (auto index : op.key_indexes) {
                VLOG(1) << "Transfer completed successfully for key: "
                        << object_keys[index];
                results[index] = {};
            }
        }
    }
    return results;
}

std::vector<tl::expected<void, ErrorCode>> Client::BatchGet(
    const std::vector<std::string>& object_keys,
    const std::vector<QueryResult>& query_results,
    std::unordered_map<std::string, std::vector<Slice>>& slices,
    bool prefer_alloc_in_same_node) {
    if (!transfer_submitter_) {
        LOG(ERROR) << "TransferSubmitter not initialized";
        std::vector<tl::expected<void, ErrorCode>> results;
        results.reserve(object_keys.size());
        for (size_t i = 0; i < object_keys.size(); ++i) {
            results.emplace_back(tl::unexpected(ErrorCode::INVALID_PARAMS));
        }
        return results;
    }

    // Validate input size consistency
    if (query_results.size() != object_keys.size()) {
        LOG(ERROR) << "Query results size (" << query_results.size()
                   << ") doesn't match object keys size (" << object_keys.size()
                   << ")";
        std::vector<tl::expected<void, ErrorCode>> results;
        results.reserve(object_keys.size());
        for (size_t i = 0; i < object_keys.size(); ++i) {
            results.emplace_back(tl::unexpected(ErrorCode::INVALID_PARAMS));
        }
        return results;
    }
    if (prefer_alloc_in_same_node) {
        return BatchGetWhenPreferSameNode(object_keys, query_results, slices);
    }

    // Collect all transfer operations for parallel execution
    std::vector<std::tuple<size_t, std::string, TransferFuture, Replica::Descriptor>> pending_transfers;
    std::vector<tl::expected<void, ErrorCode>> results(object_keys.size());
    // Record batch get transfer latency (Submit + Wait)
    auto t0_batch_get = std::chrono::steady_clock::now();

    // Collect cache hit statistics for the entire batch
    size_t total_cache_hits = 0;
    size_t total_blocks = 0;

    // Submit all transfers in parallel
    for (size_t i = 0; i < object_keys.size(); ++i) {
        const auto& key = object_keys[i];
        const auto& query_result = query_results[i];

        auto slices_it = slices.find(key);
        if (slices_it == slices.end()) {
            LOG(ERROR) << "Slices not found for key: " << key;
            results[i] = tl::unexpected(ErrorCode::INVALID_PARAMS);
            continue;
        }

        // Find the first complete replica for this key
        Replica::Descriptor replica;
        ErrorCode err =
            FindFirstCompleteReplica(query_result.replicas, replica);
        if (err != ErrorCode::OK) {
            if (err == ErrorCode::INVALID_REPLICA) {
                LOG(ERROR) << "no_complete_replicas_found key=" << key;
            }
            results[i] = tl::unexpected(err);
            continue;
        }

        if (hot_cache_ && replica.is_memory_replica()) {
            auto& mem_desc = replica.get_memory_descriptor();
            size_t key_blocks = mem_desc.buffer_descriptors.size();
            size_t key_cache_hits = updateReplicaDescriptorFromCache(key, replica);
            total_blocks += key_blocks;
            total_cache_hits += key_cache_hits;
        }

        // Submit transfer operation asynchronously
        auto future = transfer_submitter_->submit(replica, slices_it->second,
                                                  TransferRequest::READ);
        if (!future) {
            LOG(ERROR) << "Failed to submit transfer operation for key: "
                       << key;
            results[i] = tl::unexpected(ErrorCode::TRANSFER_FAIL);
            continue;
        }

        VLOG(1) << "Submitted transfer for key " << key
                << " using strategy: " << static_cast<int>(future->strategy());

        pending_transfers.emplace_back(i, key, std::move(*future), replica);
    }

    // Wait for all transfers to complete
    for (auto& [index, key, future, stored_replica] : pending_transfers) {
        ErrorCode result = future.get();
        if (result != ErrorCode::OK) {
            LOG(ERROR) << "Transfer failed for key: " << key
                       << " with error: " << static_cast<int>(result);
            results[index] = tl::unexpected(result);
        } else {
            VLOG(1) << "Transfer completed successfully for key: " << key;
            results[index] = {};

            // asynchronously update local hot cache with TE transfer slices
            if (hot_cache_ && hot_cache_handler_) {
                auto slices_it = slices.find(key);
                if (slices_it != slices.end()) {
                    ProcessSlicesAsync(key, slices_it->second, stored_replica);
                }
            }
        }
    }

    // As lease expired is a rare case, we check all the results with the same
    // time_point to avoid too many syscalls
    std::chrono::steady_clock::time_point now =
        std::chrono::steady_clock::now();
    for (size_t i = 0; i < object_keys.size(); ++i) {
        if (results[i].has_value() && query_results[i].IsLeaseExpired(now)) {
            LOG(WARNING) << "lease_expired_before_data_transfer_completed key="
                         << object_keys[i];
            results[i] = tl::unexpected(ErrorCode::LEASE_EXPIRED);
        }
    }

    auto us_batch_get = std::chrono::duration_cast<std::chrono::microseconds>(
                            std::chrono::steady_clock::now() - t0_batch_get)
                            .count();
    if (metrics_) {
        metrics_->transfer_metric.batch_get_latency_us.observe(us_batch_get);
    }

    // Log overall cache hit statistics for the entire batch
    if (hot_cache_ && total_blocks > 0) {
        VLOG(1) << "BatchGet completed: num_keys=" << object_keys.size()
                << " total_cache_hits=" << total_cache_hits
                << " total_blocks=" << total_blocks;
    } else {
        VLOG(1) << "BatchGet completed for " << object_keys.size() << " keys";
    }
    return results;
}


size_t Client::updateReplicaDescriptorFromCache(const std::string& key,
    Replica::Descriptor& replica) {
    if (!replica.is_memory_replica() || !hot_cache_) {
        return 0;
    }

    size_t cache_hits = 0;
    auto& mem_desc = replica.get_memory_descriptor();
    for (size_t i = 0; i < mem_desc.buffer_descriptors.size(); i++) {
        std::string composite_key = key + "_" + std::to_string(i);
        HotMemBlock* blk = hot_cache_->GetHotSlice(composite_key);
        if (blk != nullptr) {
            cache_hits++;
            mem_desc.buffer_descriptors[i].transport_endpoint_ = local_hostname_;
            mem_desc.buffer_descriptors[i].buffer_address_ =
                reinterpret_cast<uintptr_t>(blk->addr);
            if (mem_desc.buffer_descriptors[i].size_ != blk->size) {
                LOG(WARNING) << "Cache hit but size mismatch for key: " << composite_key;
            }
        }
    }
    return cache_hits;
}

tl::expected<void, ErrorCode> Client::Put(const ObjectKey& key,
                                          std::vector<Slice>& slices,
                                          const ReplicateConfig& config) {
    // Prepare slice lengths
    std::vector<size_t> slice_lengths;
    for (size_t i = 0; i < slices.size(); ++i) {
        slice_lengths.emplace_back(slices[i].size);
    }

    // Start put operation
    auto start_result = master_client_.PutStart(key, slice_lengths, config);
    if (!start_result) {
        ErrorCode err = start_result.error();
        if (err == ErrorCode::OBJECT_ALREADY_EXISTS) {
            VLOG(1) << "object_already_exists key=" << key;
            return {};
        }
        if (err == ErrorCode::NO_AVAILABLE_HANDLE) {
            LOG(WARNING) << "Failed to start put operation for key=" << key
                         << PUT_NO_SPACE_HELPER_STR;
        } else {
            LOG(ERROR) << "Failed to start put operation for key=" << key
                       << ": " << toString(err);
        }
        return tl::unexpected(err);
    }

    // Record Put transfer latency (all replicas)
    auto t0_put = std::chrono::steady_clock::now();

    // We must deal with disk replica first, then the disk putrevoke/putend can
    // be called surely
    if (storage_backend_) {
        for (auto it = start_result.value().rbegin();
             it != start_result.value().rend(); ++it) {
            const auto& replica = *it;
            if (replica.is_disk_replica()) {
                // Store to local file if storage backend is available
                auto disk_descriptor = replica.get_disk_descriptor();
                PutToLocalFile(key, slices, disk_descriptor);
                break;  // Only one disk replica is needed
            }
        }
    }

    for (const auto& replica : start_result.value()) {
        if (replica.is_memory_replica()) {
            // Transfer data using allocated handles from all replicas
            ErrorCode transfer_err = TransferWrite(replica, slices);
            if (transfer_err != ErrorCode::OK) {
                // Revoke put operation
                auto revoke_result =
                    master_client_.PutRevoke(key, ReplicaType::MEMORY);
                if (!revoke_result) {
                    LOG(ERROR) << "Failed to revoke put operation";
                    return tl::unexpected(revoke_result.error());
                }
                return tl::unexpected(transfer_err);
            }
        }
    }

    auto us_put = std::chrono::duration_cast<std::chrono::microseconds>(
                      std::chrono::steady_clock::now() - t0_put)
                      .count();
    if (metrics_) {
        metrics_->transfer_metric.put_latency_us.observe(us_put);
    }

    // End put operation
    auto end_result = master_client_.PutEnd(key, ReplicaType::MEMORY);
    if (!end_result) {
        ErrorCode err = end_result.error();
        LOG(ERROR) << "Failed to end put operation: " << err;
        return tl::unexpected(err);
    }

    return {};
}

// TODO: `client.cpp` is too long, consider split it into multiple files
enum class PutOperationState {
    PENDING,
    MASTER_FAILED,
    TRANSFER_FAILED,
    FINALIZE_FAILED,
    SUCCESS
};

class PutOperation {
   public:
    PutOperation(std::string_view k, const std::vector<Slice>& s)
        : key(k), slices(s) {
        value_length = CalculateSliceSize(slices);
        // Initialize with a pending error state to ensure result is always set
        result = tl::unexpected(ErrorCode::INTERNAL_ERROR);
    }

    std::string key;
    std::vector<Slice> slices;
    size_t value_length;
    std::vector<std::vector<Slice>> batched_slices;

    // Enhanced state tracking
    PutOperationState state = PutOperationState::PENDING;
    tl::expected<void, ErrorCode> result;
    std::vector<Replica::Descriptor> replicas;
    std::vector<TransferFuture> pending_transfers;

    // Error context for debugging
    std::optional<std::string> failure_context;

    // Helper methods for robust state management
    void SetSuccess() {
        state = PutOperationState::SUCCESS;
        result = {};
        failure_context.reset();
    }

    void SetError(ErrorCode error, const std::string& context = "") {
        result = tl::unexpected(error);
        if (!context.empty()) {
            failure_context = toString(error) + ": " + context + "; " +
                              failure_context.value_or("");
        }

        // Update state based on current processing stage
        if (replicas.empty()) {
            state = PutOperationState::MASTER_FAILED;
        } else if (pending_transfers.empty()) {
            state = PutOperationState::TRANSFER_FAILED;
        } else {
            state = PutOperationState::FINALIZE_FAILED;
        }
    }

    bool IsResolved() const { return state != PutOperationState::PENDING; }

    bool IsSuccessful() const {
        return state == PutOperationState::SUCCESS && result.has_value();
    }
};

std::vector<PutOperation> Client::CreatePutOperations(
    const std::vector<ObjectKey>& keys,
    const std::vector<std::vector<Slice>>& batched_slices) {
    std::vector<PutOperation> ops;
    ops.reserve(keys.size());
    for (size_t i = 0; i < keys.size(); ++i) {
        ops.emplace_back(keys[i], batched_slices[i]);
    }
    return ops;
}

void Client::StartBatchPut(std::vector<PutOperation>& ops,
                           const ReplicateConfig& config) {
    std::vector<std::string> keys;
    std::vector<std::vector<uint64_t>> slice_lengths;

    keys.reserve(ops.size());
    slice_lengths.reserve(ops.size());

    for (const auto& op : ops) {
        keys.emplace_back(op.key);

        std::vector<uint64_t> slice_sizes;
        slice_sizes.reserve(op.slices.size());
        for (const auto& slice : op.slices) {
            slice_sizes.emplace_back(slice.size);
        }
        slice_lengths.emplace_back(std::move(slice_sizes));
    }

    auto start_responses =
        master_client_.BatchPutStart(keys, slice_lengths, config);

    // Ensure response size matches request size
    if (start_responses.size() != ops.size()) {
        LOG(ERROR) << "BatchPutStart response size mismatch: expected "
                   << ops.size() << ", got " << start_responses.size();
        for (auto& op : ops) {
            op.SetError(ErrorCode::RPC_FAIL,
                        "BatchPutStart response size mismatch");
        }
        return;
    }

    // Process individual responses with robust error handling
    for (size_t i = 0; i < ops.size(); ++i) {
        if (!start_responses[i]) {
            ops[i].SetError(start_responses[i].error(),
                            "Master failed to start put operation");
        } else {
            ops[i].replicas = start_responses[i].value();
            // Operation continues to next stage - result remains INTERNAL_ERROR
            // until fully successful
            VLOG(1) << "Successfully started put for key " << ops[i].key
                    << " with " << ops[i].replicas.size() << " replicas";
        }
    }
}

void Client::SubmitTransfers(std::vector<PutOperation>& ops) {
    if (!transfer_submitter_) {
        LOG(ERROR) << "TransferSubmitter not initialized";
        for (auto& op : ops) {
            op.SetError(ErrorCode::INVALID_PARAMS,
                        "TransferSubmitter not initialized");
        }
        return;
    }

    for (auto& op : ops) {
        // Skip operations that already failed in previous stages
        if (op.IsResolved()) {
            continue;
        }

        // Skip operations that don't have replicas (failed in StartBatchPut)
        if (op.replicas.empty()) {
            op.SetError(ErrorCode::INTERNAL_ERROR,
                        "No replicas available for transfer");
            continue;
        }

        bool all_transfers_submitted = true;
        std::string failure_context;

        // We must deal with disk replica first, then the disk putrevoke/putend
        // can be called surely
        if (storage_backend_) {
            for (auto it = op.replicas.rbegin(); it != op.replicas.rend();
                 ++it) {
                const auto& replica = *it;
                if (replica.is_disk_replica()) {
                    auto disk_descriptor = replica.get_disk_descriptor();
                    PutToLocalFile(op.key, op.slices, disk_descriptor);
                    break;  // Only one disk replica is needed
                }
            }
        }

        for (size_t replica_idx = 0; replica_idx < op.replicas.size();
             ++replica_idx) {
            const auto& replica = op.replicas[replica_idx];
            if (replica.is_memory_replica()) {
                auto submit_result = transfer_submitter_->submit(
                    replica, op.slices, TransferRequest::WRITE);

                if (!submit_result) {
                    failure_context = "Failed to submit transfer for replica " +
                                      std::to_string(replica_idx);
                    all_transfers_submitted = false;
                    break;
                }

                op.pending_transfers.emplace_back(
                    std::move(submit_result.value()));
            }
        }

        if (!all_transfers_submitted) {
            LOG(ERROR) << "Transfer submission failed for key " << op.key
                       << ": " << failure_context;
            op.SetError(ErrorCode::TRANSFER_FAIL, failure_context);
            op.pending_transfers.clear();
        } else {
            VLOG(1) << "Successfully submitted " << op.pending_transfers.size()
                    << " transfers for key " << op.key;
        }
    }
}

void Client::WaitForTransfers(std::vector<PutOperation>& ops) {
    for (auto& op : ops) {
        // Skip operations that already failed or completed
        if (op.IsResolved()) {
            continue;
        }

        // Skip operations with no pending transfers (failed in SubmitTransfers)
        if (op.pending_transfers.empty()) {
            op.SetError(ErrorCode::INTERNAL_ERROR,
                        "No pending transfers to wait for");
            continue;
        }

        bool all_transfers_succeeded = true;
        ErrorCode first_error = ErrorCode::OK;
        size_t failed_transfer_idx = 0;

        for (size_t i = 0; i < op.pending_transfers.size(); ++i) {
            ErrorCode transfer_result = op.pending_transfers[i].get();
            if (transfer_result != ErrorCode::OK) {
                if (all_transfers_succeeded) {
                    // Record the first error for reporting
                    first_error = transfer_result;
                    failed_transfer_idx = i;
                    all_transfers_succeeded = false;
                }
                // Continue waiting for all transfers to avoid resource leaks
            }
        }

        if (all_transfers_succeeded) {
            VLOG(1) << "All transfers completed successfully for key "
                    << op.key;
            // Transfer phase successful - continue to finalization
            // Note: Don't mark as SUCCESS yet, need to complete finalization
        } else {
            std::string error_context =
                "Transfer " + std::to_string(failed_transfer_idx) + " failed";
            LOG(ERROR) << "Transfer failed for key " << op.key << ": "
                       << toString(first_error) << " (" << error_context << ")";
            op.SetError(first_error, error_context);
        }
    }
}

void Client::FinalizeBatchPut(std::vector<PutOperation>& ops) {
    // For each operation,
    // If transfers completed successfully, we need to call BatchPutEnd
    // If the operation failed but has allocated replicas, we need to call
    // BatchPutRevoke

    std::vector<std::string> successful_keys;
    std::vector<size_t> successful_indices;
    std::vector<std::string> failed_keys;
    std::vector<size_t> failed_indices;

    // Reserve space to avoid reallocations
    successful_keys.reserve(ops.size());
    successful_indices.reserve(ops.size());
    failed_keys.reserve(ops.size());
    failed_indices.reserve(ops.size());

    for (size_t i = 0; i < ops.size(); ++i) {
        auto& op = ops[i];

        // Check if operation completed transfers successfully and needs
        // finalization
        if (!op.IsResolved() && !op.replicas.empty() &&
            !op.pending_transfers.empty()) {
            // Transfers completed, needs BatchPutEnd
            successful_keys.emplace_back(op.key);
            successful_indices.emplace_back(i);
        } else if (op.state != PutOperationState::PENDING &&
                   !op.replicas.empty()) {
            // Operation failed but has allocated replicas, needs BatchPutRevoke
            failed_keys.emplace_back(op.key);
            failed_indices.emplace_back(i);
        }
        // Operations without replicas (early failures) don't need finalization
    }

    // Process successful operations
    if (!successful_keys.empty()) {
        auto end_responses = master_client_.BatchPutEnd(successful_keys);
        if (end_responses.size() != successful_keys.size()) {
            LOG(ERROR) << "BatchPutEnd response size mismatch: expected "
                       << successful_keys.size() << ", got "
                       << end_responses.size();
            for (size_t idx : successful_indices) {
                ops[idx].SetError(ErrorCode::RPC_FAIL,
                                  "BatchPutEnd response size mismatch");
            }
        } else {
            // Process individual responses
            for (size_t i = 0; i < end_responses.size(); ++i) {
                const size_t op_idx = successful_indices[i];
                if (!end_responses[i]) {
                    LOG(ERROR) << "Failed to finalize put for key "
                               << successful_keys[i] << ": "
                               << toString(end_responses[i].error());
                    ops[op_idx].SetError(end_responses[i].error(),
                                         "BatchPutEnd failed");
                } else {
                    // Operation fully successful
                    ops[op_idx].SetSuccess();
                    VLOG(1) << "Successfully completed put for key "
                            << successful_keys[i];
                }
            }
        }
    }

    // Process failed operations that need cleanup
    if (!failed_keys.empty()) {
        auto revoke_responses = master_client_.BatchPutRevoke(failed_keys);
        if (revoke_responses.size() != failed_keys.size()) {
            LOG(ERROR) << "BatchPutRevoke response size mismatch: expected "
                       << failed_keys.size() << ", got "
                       << revoke_responses.size();
            // Mark all failed operations with revoke RPC failure
            for (size_t idx : failed_indices) {
                ops[idx].SetError(ErrorCode::RPC_FAIL,
                                  "BatchPutRevoke response size mismatch");
            }
        } else {
            // Process individual revoke responses
            for (size_t i = 0; i < revoke_responses.size(); ++i) {
                const size_t op_idx = failed_indices[i];
                if (!revoke_responses[i]) {
                    LOG(ERROR)
                        << "Failed to revoke put for key " << failed_keys[i]
                        << ": " << toString(revoke_responses[i].error());
                    // Preserve original error but note revoke failure in
                    // context
                    std::string original_context =
                        ops[op_idx].failure_context.value_or("unknown error");
                    ops[op_idx].failure_context =
                        original_context + "; revoke also failed";
                } else {
                    LOG(INFO) << "Successfully revoked failed put for key "
                              << failed_keys[i];
                }
            }
        }
    }

    // Ensure all operations have definitive results
    for (auto& op : ops) {
        if (!op.IsResolved()) {
            op.SetError(ErrorCode::INTERNAL_ERROR,
                        "Operation not resolved after finalization");
            LOG(ERROR) << "Operation for key " << op.key
                       << " was not properly resolved";
        }
    }
}

std::vector<tl::expected<void, ErrorCode>> Client::CollectResults(
    const std::vector<PutOperation>& ops) {
    std::vector<tl::expected<void, ErrorCode>> results;
    results.reserve(ops.size());

    int no_available_handle_count = 0;
    for (const auto& op : ops) {
        // With the new structure, result is always set (never nullopt)
        results.emplace_back(op.result);

        // Additional validation and logging for debugging
        if (!op.result.has_value()) {
            // if error == object already exist, consider as ok
            if (op.result.error() == ErrorCode::OBJECT_ALREADY_EXISTS) {
                results.back() = {};
                continue;
            }
            if (op.result.error() == ErrorCode::NO_AVAILABLE_HANDLE) {
                no_available_handle_count++;
            } else {
                LOG(ERROR) << "Operation for key " << op.key
                           << " failed: " << toString(op.result.error())
                           << (op.failure_context
                                   ? (" (" + *op.failure_context + ")")
                                   : "");
            }
        } else {
            VLOG(1) << "Operation for key " << op.key
                    << " completed successfully";
        }
    }
    if (no_available_handle_count > 0) {
        LOG(WARNING) << "BatchPut failed for " << no_available_handle_count
                     << " keys" << PUT_NO_SPACE_HELPER_STR;
    }

    return results;
}

std::vector<tl::expected<void, ErrorCode>> Client::BatchPutWhenPreferSameNode(
    std::vector<PutOperation>& ops) {
    auto t0 = std::chrono::steady_clock::now();
    std::unordered_map<std::string, PutOperation> seg_to_ops{};
    for (auto& op : ops) {
        if (op.IsResolved()) {
            continue;
        }
        if (op.replicas.empty()) {
            op.SetError(ErrorCode::INTERNAL_ERROR,
                        "No replicas available for transfer");
            continue;
        }
        auto replica = op.replicas[0];
        if (!replica.is_memory_replica()) {
            op.SetError(ErrorCode::INVALID_PARAMS, "only memory is supported.");
            continue;
        }
        auto& memory_descriptor = replica.get_memory_descriptor();
        if (memory_descriptor.buffer_descriptors.empty()) {
            op.SetError(ErrorCode::INVALID_PARAMS,
                        "buffer descriptors is empty.");
            continue;
        }
        auto& buffer_descriptor = memory_descriptor.buffer_descriptors[0];
        auto seg = buffer_descriptor.transport_endpoint_;
        if (seg_to_ops.find(seg) == seg_to_ops.end()) {
            seg_to_ops.emplace(seg, PutOperation(op.key, op.slices));
        }
        // multiple replica-slices
        seg_to_ops.at(seg).batched_slices.emplace_back(op.slices);
        seg_to_ops.at(seg).replicas.emplace_back(replica);
    }
    std::vector<PutOperation> merged_ops;
    merged_ops.reserve(seg_to_ops.size());
    for (auto& seg_to_op : seg_to_ops) {
        auto& op = seg_to_op.second;
        bool all_transfers_submitted = true;
        std::string failure_context;
        merged_ops.emplace_back(op.key, op.slices);
        auto& merged_op = merged_ops.back();
        merged_op.replicas = op.replicas;
        auto submit_result = transfer_submitter_->submit_batch(
            op.replicas, op.batched_slices, TransferRequest::WRITE);
        if (!submit_result) {
            failure_context = "Failed to submit batch transfer";
            all_transfers_submitted = false;
        } else {
            merged_op.pending_transfers.emplace_back(
                std::move(submit_result.value()));
        }
        if (!all_transfers_submitted) {
            LOG(ERROR) << "Transfer submission failed for key " << op.key
                       << ": " << failure_context;
            merged_op.SetError(ErrorCode::TRANSFER_FAIL, failure_context);
            merged_op.pending_transfers.clear();
        } else {
            VLOG(1) << "Successfully submitted "
                    << merged_op.pending_transfers.size()
                    << " transfers for key " << merged_ops.back().key;
        }
    }
    WaitForTransfers(merged_ops);
    for (auto& op : merged_ops) {
        auto& memory_descriptor = op.replicas[0].get_memory_descriptor();
        auto& buffer_descriptor = memory_descriptor.buffer_descriptors[0];
        auto seg = buffer_descriptor.transport_endpoint_;
        seg_to_ops.at(seg).state = op.state;
    }
    for (auto& op : ops) {
        if (op.IsResolved()) {
            continue;
        }
        auto& memory_descriptor = op.replicas[0].get_memory_descriptor();
        auto& buffer_descriptor = memory_descriptor.buffer_descriptors[0];
        auto seg = buffer_descriptor.transport_endpoint_;
        op.state = seg_to_ops.at(seg).state;
        auto state = std::make_shared<EmptyOperationState>();
        auto future = TransferFuture(state);
        op.pending_transfers.emplace_back(std::move(future));
    }
    auto us = std::chrono::duration_cast<std::chrono::microseconds>(
                  std::chrono::steady_clock::now() - t0)
                  .count();
    if (metrics_) {
        metrics_->transfer_metric.batch_put_latency_us.observe(us);
    }
    FinalizeBatchPut(ops);
    return CollectResults(ops);
}

std::vector<tl::expected<void, ErrorCode>> Client::BatchPut(
    const std::vector<ObjectKey>& keys,
    std::vector<std::vector<Slice>>& batched_slices,
    const ReplicateConfig& config) {
    std::vector<PutOperation> ops = CreatePutOperations(keys, batched_slices);
    if (config.prefer_alloc_in_same_node) {
        if (config.replica_num != 1) {
            LOG(ERROR) << "prefer_alloc_in_same_node is not supported with "
                          "replica_num != 1";
            return std::vector<tl::expected<void, ErrorCode>>(
                keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS));
        }
        StartBatchPut(ops, config);
        return BatchPutWhenPreferSameNode(ops);
    }
    StartBatchPut(ops, config);

    auto t0 = std::chrono::steady_clock::now();
    SubmitTransfers(ops);
    WaitForTransfers(ops);
    auto us = std::chrono::duration_cast<std::chrono::microseconds>(
                  std::chrono::steady_clock::now() - t0)
                  .count();
    if (metrics_) {
        metrics_->transfer_metric.batch_put_latency_us.observe(us);
    }

    FinalizeBatchPut(ops);
    return CollectResults(ops);
}

tl::expected<void, ErrorCode> Client::Remove(const ObjectKey& key) {
    auto result = master_client_.Remove(key);
    // if (storage_backend_) {
    //     storage_backend_->RemoveFile(key);
    // }
    if (!result) {
        return tl::unexpected(result.error());
    }
    return {};
}

tl::expected<long, ErrorCode> Client::RemoveByRegex(const ObjectKey& str) {
    auto result = master_client_.RemoveByRegex(str);
    // if (storage_backend_) {
    //     storage_backend_->RemoveByRegex(str);
    // }
    if (!result) {
        return tl::unexpected(result.error());
    }
    return result.value();
}

tl::expected<long, ErrorCode> Client::RemoveAll() {
    // if (storage_backend_) {
    //     storage_backend_->RemoveAll();
    // }
    return master_client_.RemoveAll();
}

tl::expected<void, ErrorCode> Client::MountSegment(const void* buffer,
                                                   size_t size) {
    auto check_result = CheckRegisterMemoryParams(buffer, size);
    if (!check_result) {
        return tl::unexpected(check_result.error());
    }

    std::lock_guard<std::mutex> lock(mounted_segments_mutex_);

    // Check if the segment overlaps with any existing segment
    for (auto& it : mounted_segments_) {
        auto& mtseg = it.second;
        uintptr_t l1 = reinterpret_cast<uintptr_t>(mtseg.base);
        uintptr_t r1 = reinterpret_cast<uintptr_t>(mtseg.size) + l1;
        uintptr_t l2 = reinterpret_cast<uintptr_t>(buffer);
        uintptr_t r2 = reinterpret_cast<uintptr_t>(size) + l2;
        if (std::max(l1, l2) < std::min(r1, r2)) {
            LOG(ERROR) << "segment_overlaps base1=" << mtseg.base
                       << " size1=" << mtseg.size << " base2=" << buffer
                       << " size2=" << size;
            return tl::unexpected(ErrorCode::INVALID_PARAMS);
        }
    }
#ifndef USE_UB
    int rc = transfer_engine_->registerLocalMemory(
        (void*)buffer, size, kWildcardLocation, true, true);
    if (rc != 0) {
        LOG(ERROR) << "register_local_memory_failed base=" << buffer
                   << " size=" << size << ", error=" << rc;
        return tl::unexpected(ErrorCode::INVALID_PARAMS);
    }
#endif
    // Build segment with logical name; attach TE endpoint for transport
    Segment segment;
    segment.id = generate_uuid();
    segment.name = local_hostname_;
    segment.base = reinterpret_cast<uintptr_t>(buffer);
    segment.size = size;
    // For P2P handshake mode, publish the actual transport endpoint that was
    // negotiated by the transfer engine. Otherwise, keep the logical hostname
    // so metadata backends (HTTP/etcd/redis) can resolve the segment by name.
    if (metadata_connstring_ == P2PHANDSHAKE) {
        segment.te_endpoint = transfer_engine_->getLocalIpAndPort();
    } else {
        segment.te_endpoint = local_hostname_;
    }

    auto mount_result = master_client_.MountSegment(segment, client_id_);
    if (!mount_result) {
        ErrorCode err = mount_result.error();
        LOG(ERROR) << "mount_segment_to_master_failed base=" << buffer
                   << " size=" << size << ", error=" << err;
        return tl::unexpected(err);
    }

    mounted_segments_[segment.id] = segment;
    return {};
}

tl::expected<void, ErrorCode> Client::UnmountSegment(const void* buffer,
                                                     size_t size) {
    std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
    auto segment = mounted_segments_.end();

    for (auto it = mounted_segments_.begin(); it != mounted_segments_.end();
         ++it) {
        if (it->second.base == reinterpret_cast<uintptr_t>(buffer) &&
            it->second.size == size) {
            segment = it;
            break;
        }
    }
    if (segment == mounted_segments_.end()) {
        LOG(ERROR) << "segment_not_found base=" << buffer << " size=" << size;
        return tl::unexpected(ErrorCode::INVALID_PARAMS);
    }

    auto unmount_result =
        master_client_.UnmountSegment(segment->second.id, client_id_);
    if (!unmount_result) {
        ErrorCode err = unmount_result.error();
        LOG(ERROR) << "Failed to unmount segment from master: "
                   << toString(err);
        return tl::unexpected(err);
    }

    int rc = transfer_engine_->unregisterLocalMemory(
        reinterpret_cast<void*>(segment->second.base));
    if (rc != 0) {
        LOG(ERROR) << "Failed to unregister transfer buffer with transfer "
                      "engine ret is "
                   << rc;
        if (rc != ERR_ADDRESS_NOT_REGISTERED) {
            return tl::unexpected(ErrorCode::INTERNAL_ERROR);
        }
        // Otherwise, the segment is already unregistered from transfer
        // engine, we can continue
    }

    mounted_segments_.erase(segment);
    return {};
}

tl::expected<void, ErrorCode> Client::RegisterLocalMemory(
    void* addr, size_t length, const std::string& location,
    bool remote_accessible, bool update_metadata) {
    auto check_result = CheckRegisterMemoryParams(addr, length);
    if (!check_result) {
        return tl::unexpected(check_result.error());
    }
    if (this->transfer_engine_->registerLocalMemory(
            addr, length, location, remote_accessible, update_metadata) != 0) {
        return tl::unexpected(ErrorCode::INVALID_PARAMS);
    }
    return {};
}

tl::expected<void, ErrorCode> Client::unregisterLocalMemory(
    void* addr, bool update_metadata) {
    if (this->transfer_engine_->unregisterLocalMemory(addr, update_metadata) !=
        0) {
        return tl::unexpected(ErrorCode::INVALID_PARAMS);
    }
    return {};
}

tl::expected<bool, ErrorCode> Client::IsExist(const std::string& key) {
    auto result = master_client_.ExistKey(key);
    return result;
}

std::vector<tl::expected<bool, ErrorCode>> Client::BatchIsExist(
    const std::vector<std::string>& keys) {
    auto response = master_client_.BatchExistKey(keys);

    // Check if we got the expected number of responses
    if (response.size() != keys.size()) {
        LOG(ERROR) << "BatchExistKey response size mismatch. Expected: "
                   << keys.size() << ", Got: " << response.size();
        // Return vector of RPC_FAIL errors
        std::vector<tl::expected<bool, ErrorCode>> results;
        results.reserve(keys.size());
        for (size_t i = 0; i < keys.size(); ++i) {
            results.emplace_back(tl::unexpected(ErrorCode::RPC_FAIL));
        }
        return results;
    }

    // Return the response directly as it's already in the correct
    // format
    return response;
}

void Client::PrepareStorageBackend(const std::string& storage_root_dir,
                                   const std::string& fsdir) {
    // Initialize storage backend
    storage_backend_ = StorageBackend::Create(storage_root_dir, fsdir);
    if (!storage_backend_) {
        LOG(INFO) << "Failed to initialize storage backend";
    }
}

void Client::PutToLocalFile(const std::string& key,
                            const std::vector<Slice>& slices,
                            const DiskDescriptor& disk_descriptor) {
    if (!storage_backend_) return;

    size_t total_size = 0;
    for (const auto& slice : slices) {
        total_size += slice.size;
    }

    std::string path = disk_descriptor.file_path;
    // Currently, persistence is achieved through asynchronous writes, but
    // before asynchronous writing in 3FS, significant performance degradation
    // may occur due to data copying. Profiling reveals that the number of page
    // faults triggered in this scenario is nearly double the normal count.
    // Future plans include introducing a reuse buffer list to address this
    // performance degradation issue.

    std::string value;
    value.reserve(total_size);
    for (const auto& slice : slices) {
        value.append(static_cast<char*>(slice.ptr), slice.size);
    }

    write_thread_pool_.enqueue([this, backend = storage_backend_, key,
                                value = std::move(value), path] {
        // Store the object
        auto store_result = backend->StoreObject(path, value);
        ReplicaType replica_type = ReplicaType::DISK;

        if (!store_result) {
            // If storage failed, revoke the put operation
            LOG(ERROR) << "Failed to store object for key: " << key;
            auto revoke_result = master_client_.PutRevoke(key, replica_type);
            if (!revoke_result) {
                LOG(ERROR) << "Failed to revoke put operation for key: " << key;
            }
            return;
        }

        // If storage succeeded, end the put operation
        auto end_result = master_client_.PutEnd(key, replica_type);
        if (!end_result) {
            LOG(ERROR) << "Failed to end put operation for key: " << key;
        }
    });
}

ErrorCode Client::TransferData(const Replica::Descriptor& replica_descriptor,
                               std::vector<Slice>& slices,
                               TransferRequest::OpCode op_code) {
    if (!transfer_submitter_) {
        LOG(ERROR) << "TransferSubmitter not initialized";
        return ErrorCode::INVALID_PARAMS;
    }

    auto future =
        transfer_submitter_->submit(replica_descriptor, slices, op_code);
    if (!future) {
        LOG(ERROR) << "Failed to submit transfer operation";
        return ErrorCode::TRANSFER_FAIL;
    }

    VLOG(1) << "Using transfer strategy: " << future->strategy();

    return future->get();
}

ErrorCode Client::TransferWrite(const Replica::Descriptor& replica_descriptor,
                                std::vector<Slice>& slices) {
    return TransferData(replica_descriptor, slices, TransferRequest::WRITE);
}

ErrorCode Client::TransferRead(const Replica::Descriptor& replica_descriptor,
                               std::vector<Slice>& slices) {
    size_t total_size = 0;
    if (replica_descriptor.is_memory_replica()) {
        auto& mem_desc = replica_descriptor.get_memory_descriptor();
        for (const auto& handle : mem_desc.buffer_descriptors) {
            total_size += handle.size_;
        }
    } else {
        auto& disk_desc = replica_descriptor.get_disk_descriptor();
        total_size = disk_desc.object_size;
    }

    size_t slices_size = CalculateSliceSize(slices);
    if (slices_size < total_size) {
        LOG(ERROR) << "Slice size " << slices_size << " is smaller than total "
                   << "size " << total_size;
        return ErrorCode::INVALID_PARAMS;
    }

    return TransferData(replica_descriptor, slices, TransferRequest::READ);
}

void Client::PingThreadMain(bool is_ha_mode,
                            std::string current_master_address) {
    // How many failed pings before getting latest master view from etcd
    const int max_ping_fail_count = 3;
    // How long to wait for next ping after success
    const int success_ping_interval_ms = 1000;
    // How long to wait for next ping after failure
    const int fail_ping_interval_ms = 1000;
    // Increment after a ping failure, reset after a ping success
    int ping_fail_count = 0;

    auto remount_segment = [this]() {
        // This lock must be held until the remount rpc is finished,
        // otherwise there will be corner cases, e.g., a segment is
        // unmounted successfully first, and then remounted again in
        // this thread.
        std::lock_guard<std::mutex> lock(mounted_segments_mutex_);
        std::vector<Segment> segments;
        for (auto it : mounted_segments_) {
            auto& segment = it.second;
            segments.emplace_back(segment);
        }
        auto remount_result =
            master_client_.ReMountSegment(segments, client_id_);
        if (!remount_result) {
            ErrorCode err = remount_result.error();
            LOG(ERROR) << "Failed to remount segments: " << err;
        }
    };
    // Use another thread to remount segments to avoid blocking the ping
    // thread
    std::future<void> remount_segment_future;

    while (ping_running_) {
        // Join the remount segment thread if it is ready
        if (remount_segment_future.valid() &&
            remount_segment_future.wait_for(std::chrono::seconds(0)) ==
                std::future_status::ready) {
            remount_segment_future = std::future<void>();
        }

        // Ping master
        auto ping_result = master_client_.Ping(client_id_);
        if (ping_result) {
            // Reset ping failure count
            ping_fail_count = 0;
            auto& ping_response = ping_result.value();
            if (ping_response.client_status == ClientStatus::NEED_REMOUNT &&
                !remount_segment_future.valid()) {
                // Ensure at most one remount segment thread is running
                remount_segment_future =
                    std::async(std::launch::async, remount_segment);
            }
            std::this_thread::sleep_for(
                std::chrono::milliseconds(success_ping_interval_ms));
            continue;
        }

        ping_fail_count++;
        if (ping_fail_count < max_ping_fail_count) {
            LOG(ERROR) << "Failed to ping master";
            std::this_thread::sleep_for(
                std::chrono::milliseconds(fail_ping_interval_ms));
            continue;
        }

        // Exceeded ping failure threshold. Reconnect based on mode.
        if (is_ha_mode) {
            LOG(ERROR)
                << "Failed to ping master for " << ping_fail_count
                << " times; fetching latest master view and reconnecting";
            std::string master_address;
            ViewVersionId next_version = 0;
            auto err =
                master_view_helper_.GetMasterView(master_address, next_version);
            if (err != ErrorCode::OK) {
                LOG(ERROR) << "Failed to get new master view: "
                           << toString(err);
                std::this_thread::sleep_for(
                    std::chrono::milliseconds(fail_ping_interval_ms));
                continue;
            }

            err = master_client_.Connect(master_address);
            if (err != ErrorCode::OK) {
                LOG(ERROR) << "Failed to connect to master " << master_address
                           << ": " << toString(err);
                std::this_thread::sleep_for(
                    std::chrono::milliseconds(fail_ping_interval_ms));
                continue;
            }

            current_master_address = master_address;
            LOG(INFO) << "Reconnected to master " << master_address;
            ping_fail_count = 0;
        } else {
            LOG(ERROR) << "Failed to ping master for " << ping_fail_count
                       << " times (non-HA); reconnecting to "
                       << current_master_address;
            auto err = master_client_.Connect(current_master_address);
            if (err != ErrorCode::OK) {
                LOG(ERROR) << "Reconnect failed to " << current_master_address
                           << ": " << toString(err);
                std::this_thread::sleep_for(
                    std::chrono::milliseconds(fail_ping_interval_ms));
                continue;
            }
            LOG(INFO) << "Reconnected to master " << current_master_address;
            ping_fail_count = 0;
        }
    }
    // Explicitly wait for the remount segment thread to finish
    if (remount_segment_future.valid()) {
        remount_segment_future.wait();
    }
}

ErrorCode Client::FindFirstCompleteReplica(
    const std::vector<Replica::Descriptor>& replica_list,
    Replica::Descriptor& replica) {
    // Find the first complete replica
    for (size_t i = 0; i < replica_list.size(); ++i) {
        if (replica_list[i].status == ReplicaStatus::COMPLETE) {
            replica = replica_list[i];
            return ErrorCode::OK;
        }
    }

    // No complete replica found
    return ErrorCode::INVALID_REPLICA;
}

ErrorCode Client::InitLocalHotCache() {
    // Defaults: enable hot cache with 1GB by default
    size_t total_cache = 1ull * 1024 * 1024 * 1024; // 1GB default total size
    bool enable_hot_cache = false;

    // Read from environment
    if (const char* ev_size = std::getenv("LOCAL_HOT_CACHE_SIZE")) {
        std::string ev_size_str(ev_size);
        // Check for negative values
        if (!ev_size_str.empty() && ev_size_str[0] == '-') {
            LOG(WARNING) << "Invalid LOCAL_HOT_CACHE_SIZE='" << ev_size_str << "', disable local hot cache";
            return ErrorCode::INVALID_PARAMS;
        }
        try {
            unsigned long long v = std::stoull(ev_size_str, nullptr, 10);
            if (v > 0) {
                enable_hot_cache = true;
                total_cache = static_cast<size_t>(v);
            } else {
                LOG(WARNING) << "Invalid LOCAL_HOT_CACHE_SIZE='" << ev_size_str << "', disable local hot cache";
                return ErrorCode::INVALID_PARAMS;
            }
        } catch (const std::exception&) {
            LOG(WARNING) << "Invalid LOCAL_HOT_CACHE_SIZE='" << ev_size_str << "', disable local hot cache";
            return ErrorCode::INVALID_PARAMS;
        }
    }

    if (enable_hot_cache) {
        hot_cache_ = std::make_shared<LocalHotCache>(total_cache);
        // Check if cache initialization was successful
        if (hot_cache_->GetCacheSize() == 0) {
            LOG(ERROR) << "Local hot cache creation failed: no blocks allocated. "
                       << "total_cache=" << total_cache;
            hot_cache_.reset();
            return ErrorCode::INVALID_PARAMS;
        }
        LOG(INFO) << "Local hot cache enabled with cache size=" << total_cache
                  << ", block amount=" << hot_cache_->GetCacheSize();
        // Create async handler with 2 worker threads
        hot_cache_handler_ = std::make_unique<LocalHotCacheHandler>(hot_cache_, 2);
    } else {
        hot_cache_.reset();
        hot_cache_handler_.reset();
        LOG(INFO) << "Local hot cache disabled";
    }
    return ErrorCode::OK;
}

void Client::ProcessSlicesAsync(
    const std::string& key,
    const std::vector<Slice>& slices,
    const Replica::Descriptor& replica) {
    if (!(hot_cache_ && hot_cache_handler_ && replica.is_memory_replica())) {
        return;
    }

    const auto& mem_desc = replica.get_memory_descriptor();

    // Identify TE transfer slices (non-local) and submit async put tasks
    for (size_t i = 0; i < slices.size(); ++i) {
        // Only cache slices that came from TE transfer (not local)
        if (mem_desc.buffer_descriptors[i].transport_endpoint_ != local_hostname_) {
            std::string composite_key = key + "_" + std::to_string(i);
            hot_cache_handler_->SubmitPutTask(composite_key, slices[i]);
        }
    }
}

}  // namespace mooncake