// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed 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 <fcntl.h>
#include <gtest/gtest.h>
#include <unistd.h>

#include <algorithm>
#include <atomic>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <fstream>
#include <future>
#include <limits>
#include <memory>
#include <stdexcept>
#include <string>
#include <thread>
#include <utility>
#include <vector>

#include "folly/CancellationToken.h"
#include "folly/ScopeGuard.h"
#include "common/Common.h"
#include "common/EasyAssert.h"
#include "filemanager/InputStream.h"
#include "test_utils/Constants.h"
#include "milvus-storage/filesystem/fs.h"
#include "storage/IndexEntryDirectStreamWriter.h"
#include "storage/IndexEntryEncryptedLocalWriter.h"
#include "storage/IndexEntryReader.h"
#include "storage/EntryStreamUtils.h"
#include "storage/PluginLoader.h"
#include "storage/RemoteInputStream.h"
#include "storage/RemoteOutputStream.h"
#include "storage/ThreadPools.h"

using namespace milvus::storage;

namespace {

class IndexEntryStreamConfigGuard {
 public:
    IndexEntryStreamConfigGuard()
        : budget_(TransientMemoryBudget::GetLoadTransientBudget()),
          capacity_bytes_(budget_.CapacityBytes()) {
    }

    ~IndexEntryStreamConfigGuard() {
        budget_.SetCapacityBytes(capacity_bytes_);
    }

 private:
    TransientMemoryBudget& budget_;
    size_t capacity_bytes_;
};

// Simple XOR-based mock cipher for testing (NOT for production use!)
class MockEncryptor : public plugin::IEncryptor {
 public:
    explicit MockEncryptor(uint8_t key) : key_(key) {
    }

    std::string
    Encrypt(const std::string& plaintext) const override {
        return Encrypt(plaintext.data(), plaintext.size());
    }

    std::string
    Encrypt(std::string_view plaintext) const override {
        return Encrypt(plaintext.data(), plaintext.size());
    }

    std::string
    Encrypt(const void* data, size_t len) const override {
        // Format: [1-byte key][XOR'd data]
        std::string result;
        result.reserve(1 + len);
        result.push_back(static_cast<char>(key_));
        const auto* src = static_cast<const uint8_t*>(data);
        for (size_t i = 0; i < len; i++) {
            result.push_back(static_cast<char>(src[i] ^ key_));
        }
        return result;
    }

    std::string
    GetKey() const override {
        return std::string(1, static_cast<char>(key_));
    }

 private:
    uint8_t key_;
};

class MockDecryptor : public plugin::IDecryptor {
 public:
    std::string
    Decrypt(const std::string& ciphertext) const override {
        return Decrypt(ciphertext.data(), ciphertext.size());
    }

    std::string
    Decrypt(std::string_view ciphertext) const override {
        return Decrypt(ciphertext.data(), ciphertext.size());
    }

    std::string
    Decrypt(const void* data, size_t len) const override {
        if (len < 1) {
            return "";
        }
        const auto* src = static_cast<const uint8_t*>(data);
        uint8_t key = src[0];
        std::string result;
        result.reserve(len - 1);
        for (size_t i = 1; i < len; i++) {
            result.push_back(static_cast<char>(src[i] ^ key));
        }
        return result;
    }

    std::string
    GetKey() const override {
        return "";
    }
};

class MockCipherPlugin : public plugin::ICipherPlugin {
 public:
    std::string
    getPluginName() const override {
        return "CipherPlugin";
    }

    void
    Update(int64_t, int64_t, const std::string&) override {
    }

    std::pair<std::shared_ptr<plugin::IEncryptor>, std::string>
    GetEncryptor(int64_t, int64_t) const override {
        return {std::make_shared<MockEncryptor>(0x5A), "mock_edek"};
    }

    std::shared_ptr<plugin::IDecryptor>
    GetDecryptor(int64_t, int64_t, const std::string&) const override {
        return std::make_shared<MockDecryptor>();
    }
};

class DelayedFailingInputStream : public milvus::InputStream {
 public:
    struct Rule {
        size_t offset;
        std::chrono::milliseconds delay;
        bool fail;
    };

    DelayedFailingInputStream(std::shared_ptr<milvus::InputStream> base,
                              std::vector<Rule> rules)
        : base_(std::move(base)), rules_(std::move(rules)) {
    }

    size_t
    Size() const override {
        return base_->Size();
    }

    bool
    Seek(int64_t offset) override {
        return base_->Seek(offset);
    }

    size_t
    Tell() const override {
        return base_->Tell();
    }

    bool
    Eof() const override {
        return base_->Eof();
    }

    size_t
    Read(void* ptr, size_t size) override {
        return base_->Read(ptr, size);
    }

    size_t
    ReadAt(void* ptr, size_t offset, size_t size) override {
        for (const auto& rule : rules_) {
            if (rule.offset == offset) {
                std::this_thread::sleep_for(rule.delay);
                if (rule.fail) {
                    return 0;
                }
                break;
            }
        }
        return base_->ReadAt(ptr, offset, size);
    }

    size_t
    Read(int fd, size_t size) override {
        return base_->Read(fd, size);
    }

 private:
    std::shared_ptr<milvus::InputStream> base_;
    std::vector<Rule> rules_;
};

class TrackingDelayedInputStream : public milvus::InputStream {
 public:
    TrackingDelayedInputStream(std::shared_ptr<milvus::InputStream> base,
                               size_t tracked_read_size,
                               std::chrono::milliseconds delay)
        : base_(std::move(base)),
          tracked_read_size_(tracked_read_size),
          delay_(delay) {
    }

    void
    EnableTracking() {
        tracking_enabled_.store(true);
    }

    size_t
    MaxActiveReads() const {
        return max_active_reads_.load();
    }

    size_t
    Size() const override {
        return base_->Size();
    }

    bool
    Seek(int64_t offset) override {
        return base_->Seek(offset);
    }

    size_t
    Tell() const override {
        return base_->Tell();
    }

    bool
    Eof() const override {
        return base_->Eof();
    }

    size_t
    Read(void* ptr, size_t size) override {
        return base_->Read(ptr, size);
    }

    size_t
    ReadAt(void* ptr, size_t offset, size_t size) override {
        if (!tracking_enabled_.load() || size != tracked_read_size_) {
            return base_->ReadAt(ptr, offset, size);
        }

        auto active = active_reads_.fetch_add(1) + 1;
        UpdateMaxActiveReads(active);
        std::this_thread::sleep_for(delay_);
        try {
            auto bytes_read = base_->ReadAt(ptr, offset, size);
            active_reads_.fetch_sub(1);
            return bytes_read;
        } catch (...) {
            active_reads_.fetch_sub(1);
            throw;
        }
    }

    size_t
    Read(int fd, size_t size) override {
        return base_->Read(fd, size);
    }

 private:
    void
    UpdateMaxActiveReads(size_t active) {
        auto current = max_active_reads_.load();
        while (active > current &&
               !max_active_reads_.compare_exchange_weak(current, active)) {
        }
    }

    std::shared_ptr<milvus::InputStream> base_;
    size_t tracked_read_size_;
    std::chrono::milliseconds delay_;
    std::atomic<bool> tracking_enabled_{false};
    std::atomic<size_t> active_reads_{0};
    std::atomic<size_t> max_active_reads_{0};
};

}  // namespace

namespace {

std::string
GetRootPath() {
    return TestLocalPath + "index_writer_test";
}
const std::string kV3FilePath = "test_v3_index";

std::vector<uint8_t>
GeneratePattern(size_t size) {
    std::vector<uint8_t> data(size);
    for (size_t i = 0; i < size; ++i) {
        data[i] = static_cast<uint8_t>(i % 256);
    }
    return data;
}

void
VerifyPattern(const std::vector<uint8_t>& data, size_t expected_size) {
    ASSERT_EQ(data.size(), expected_size);
    for (size_t i = 0; i < expected_size; ++i) {
        ASSERT_EQ(data[i], static_cast<uint8_t>(i % 256))
            << "Mismatch at byte " << i;
    }
}

}  // namespace

class IndexEntryWriterV3Test : public testing::Test {
 public:
    void
    SetUp() override {
        auto conf = milvus_storage::ArrowFileSystemConfig();
        conf.storage_type = "local";
        conf.root_path = GetRootPath();
        auto result = milvus_storage::CreateArrowFileSystem(conf);
        ASSERT_TRUE(result.ok()) << result.status().ToString();
        fs_ = result.ValueOrDie();
    }

    void
    TearDown() override {
        if (fs_) {
            fs_->DeleteDirContents("");
        }
    }

 protected:
    std::shared_ptr<milvus::OutputStream>
    CreateOutputStream(const std::string& path) {
        auto result = fs_->OpenOutputStream(path);
        EXPECT_TRUE(result.ok()) << result.status().ToString();
        return std::make_shared<RemoteOutputStream>(
            std::move(result.ValueOrDie()));
    }

    std::shared_ptr<milvus::InputStream>
    CreateInputStream(const std::string& path) {
        auto result = fs_->OpenInputFile(path);
        EXPECT_TRUE(result.ok()) << result.status().ToString();
        return std::make_shared<RemoteInputStream>(
            std::move(result.ValueOrDie()));
    }

    int64_t
    GetFileSize(const std::string& path) {
        auto info = fs_->GetFileInfo(path);
        EXPECT_TRUE(info.ok()) << info.status().ToString();
        return info.ValueOrDie().size();
    }

    milvus_storage::ArrowFileSystemPtr fs_;
};

TEST_F(IndexEntryWriterV3Test, SmallMemoryEntryRoundtrip) {
    const std::string file_path = kV3FilePath + "_small";
    const size_t entry_size =
        512 * 1024;  // 512 KB, well under 1MB cache threshold
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("small_entry", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    auto entry = reader->ReadEntry("small_entry");
    VerifyPattern(entry.data, entry_size);
}

TEST_F(IndexEntryWriterV3Test, LargeMemoryEntryRoundtrip) {
    const std::string file_path = kV3FilePath + "_large";
    const size_t entry_size =
        17 * 1024 * 1024;  // 17 MB, exceeds 16MB threshold
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("large_entry", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    auto entry = reader->ReadEntry("large_entry");
    VerifyPattern(entry.data, entry_size);
}

TEST_F(IndexEntryWriterV3Test, FdEntryRoundtrip) {
    const std::string file_path = kV3FilePath + "_fd";
    const size_t entry_size = 2 * 1024 * 1024;  // 2 MB
    auto data = GeneratePattern(entry_size);

    std::string tmp_relative = "fd_source.bin";
    std::string tmp_absolute = GetRootPath() + "/" + tmp_relative;
    {
        auto tmp_out = CreateOutputStream(tmp_relative);
        tmp_out->Write(data.data(), data.size());
        tmp_out->Close();
    }

    int fd = ::open(tmp_absolute.c_str(), O_RDONLY);
    ASSERT_NE(fd, -1) << "Failed to open temp file: " << strerror(errno);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("fd_entry", fd, entry_size);
        writer.Finish();
    }
    ::close(fd);

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    auto entry = reader->ReadEntry("fd_entry");
    VerifyPattern(entry.data, entry_size);
}

TEST_F(IndexEntryWriterV3Test, MultipleEntriesRoundtrip) {
    const std::string file_path = kV3FilePath + "_multi";
    const size_t meta_size = 128;
    const size_t data_size = 4 * 1024 * 1024;  // 4 MB

    auto meta_data = GeneratePattern(meta_size);
    auto index_data = GeneratePattern(data_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("META", meta_data.data(), meta_data.size());
        writer.WriteEntry("DATA", index_data.data(), index_data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    auto meta_entry = reader->ReadEntry("META");
    VerifyPattern(meta_entry.data, meta_size);

    auto data_entry = reader->ReadEntry("DATA");
    VerifyPattern(data_entry.data, data_size);
}

TEST_F(IndexEntryWriterV3Test, DuplicateNameThrows) {
    const std::string file_path = kV3FilePath + "_dup";
    auto data = GeneratePattern(64);

    auto output = CreateOutputStream(file_path);
    IndexEntryDirectStreamWriter writer(output);
    writer.WriteEntry("entry", data.data(), data.size());

    EXPECT_THROW(writer.WriteEntry("entry", data.data(), data.size()),
                 milvus::SegcoreError);
}

TEST_F(IndexEntryWriterV3Test, ReadEntryToFile) {
    const std::string file_path = kV3FilePath + "_tofile";
    const size_t entry_size = 1024 * 1024;  // 1 MB
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("file_entry", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    std::string local_file = GetRootPath() + "/read_entry_output.bin";
    reader->ReadEntryToFile("file_entry", local_file);

    std::ifstream ifs(local_file, std::ios::binary | std::ios::ate);
    ASSERT_TRUE(ifs.is_open());
    size_t read_size = ifs.tellg();
    ASSERT_EQ(read_size, entry_size);

    ifs.seekg(0);
    std::vector<uint8_t> read_data(read_size);
    ifs.read(reinterpret_cast<char*>(read_data.data()), read_size);
    VerifyPattern(read_data, entry_size);

    ::unlink(local_file.c_str());
}

TEST_F(IndexEntryWriterV3Test, ReadEntriesToFiles) {
    const std::string file_path = kV3FilePath + "_tofiles";
    const size_t size_a = 256 * 1024;
    const size_t size_b = 512 * 1024;

    auto data_a = GeneratePattern(size_a);
    auto data_b = GeneratePattern(size_b);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("entry_a", data_a.data(), data_a.size());
        writer.WriteEntry("entry_b", data_b.data(), data_b.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    std::string file_a = GetRootPath() + "/output_a.bin";
    std::string file_b = GetRootPath() + "/output_b.bin";

    std::vector<std::pair<std::string, std::string>> pairs = {
        {"entry_a", file_a}, {"entry_b", file_b}};
    reader->ReadEntriesToFiles(pairs);

    auto verify_file = [](const std::string& path, size_t expected_size) {
        std::ifstream ifs(path, std::ios::binary | std::ios::ate);
        ASSERT_TRUE(ifs.is_open()) << "Failed to open: " << path;
        size_t read_size = ifs.tellg();
        ASSERT_EQ(read_size, expected_size);
        ifs.seekg(0);
        std::vector<uint8_t> buf(read_size);
        ifs.read(reinterpret_cast<char*>(buf.data()), read_size);
        VerifyPattern(buf, expected_size);
    };

    verify_file(file_a, size_a);
    verify_file(file_b, size_b);

    ::unlink(file_a.c_str());
    ::unlink(file_b.c_str());
}

TEST_F(IndexEntryWriterV3Test, GetEntryNames) {
    const std::string file_path = kV3FilePath + "_names";
    auto data = GeneratePattern(64);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("alpha", data.data(), data.size());
        writer.WriteEntry("beta", data.data(), data.size());
        writer.WriteEntry("gamma", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    auto names = reader->GetEntryNames();
    ASSERT_EQ(names.size(), 4);
    EXPECT_EQ(names[0], "alpha");
    EXPECT_EQ(names[1], "beta");
    EXPECT_EQ(names[2], "gamma");
    EXPECT_EQ(names[3], "__meta__");
}

TEST_F(IndexEntryWriterV3Test, EntryNotFound) {
    const std::string file_path = kV3FilePath + "_notfound";
    auto data = GeneratePattern(64);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("exists", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    EXPECT_THROW(reader->ReadEntry("does_not_exist"), milvus::SegcoreError);
}

TEST_F(IndexEntryWriterV3Test, SmallEntryCache) {
    const std::string file_path = kV3FilePath + "_cache";
    const size_t entry_size = 256;  // well under 1MB cache threshold
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("cached_entry", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    auto entry1 = reader->ReadEntry("cached_entry");
    auto entry2 = reader->ReadEntry("cached_entry");

    ASSERT_EQ(entry1.data.size(), entry2.data.size());
    EXPECT_EQ(
        std::memcmp(entry1.data.data(), entry2.data.data(), entry1.data.size()),
        0);
    VerifyPattern(entry1.data, entry_size);
}

TEST_F(IndexEntryWriterV3Test, GetTotalBytesWritten) {
    const std::string file_path = kV3FilePath + "_totalbytes";
    const size_t entry_size = 1024;
    auto data = GeneratePattern(entry_size);

    auto output = CreateOutputStream(file_path);
    IndexEntryDirectStreamWriter writer(output);
    writer.WriteEntry("entry", data.data(), data.size());
    writer.Finish();

    size_t total = writer.GetTotalBytesWritten();
    EXPECT_GT(total, entry_size);

    int64_t actual_file_size = GetFileSize(file_path);
    EXPECT_EQ(total, static_cast<size_t>(actual_file_size));
}

TEST_F(IndexEntryWriterV3Test, WriteAfterFinishThrows) {
    const std::string file_path = kV3FilePath + "_afterfinish";
    auto data = GeneratePattern(64);

    auto output = CreateOutputStream(file_path);
    IndexEntryDirectStreamWriter writer(output);
    writer.WriteEntry("entry", data.data(), data.size());
    writer.Finish();

    EXPECT_THROW(writer.WriteEntry("another", data.data(), data.size()),
                 milvus::SegcoreError);
}

TEST_F(IndexEntryWriterV3Test, MixedSizeEntries) {
    const std::string file_path = kV3FilePath + "_mixed";
    const size_t tiny_size = 48;
    const size_t large_size = 17 * 1024 * 1024;  // 17 MB

    auto tiny_data = GeneratePattern(tiny_size);
    auto large_data = GeneratePattern(large_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("tiny_meta", tiny_data.data(), tiny_data.size());
        writer.WriteEntry("large_data", large_data.data(), large_data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    auto names = reader->GetEntryNames();
    ASSERT_EQ(names.size(), 3);
    EXPECT_EQ(names[0], "tiny_meta");
    EXPECT_EQ(names[1], "large_data");
    EXPECT_EQ(names[2], "__meta__");

    auto tiny_entry = reader->ReadEntry("tiny_meta");
    VerifyPattern(tiny_entry.data, tiny_size);

    auto large_entry = reader->ReadEntry("large_data");
    VerifyPattern(large_entry.data, large_size);
}

TEST_F(IndexEntryWriterV3Test, SetMetaRoundtrip) {
    const std::string file_path = kV3FilePath + "_meta";
    auto data = GeneratePattern(256);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());
        writer.PutMeta("index_type", std::string("bitmap"));
        writer.PutMeta("build_id", 42);
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    // Verify __meta__ entry is present
    auto names = reader->GetEntryNames();
    EXPECT_EQ(names.back(), "__meta__");

    // Read and verify the meta content via GetMeta
    EXPECT_EQ(reader->GetMeta<std::string>("index_type"), "bitmap");
    EXPECT_EQ(reader->GetMeta<int>("build_id"), 42);

    // Regular data entry should still be readable
    auto data_entry = reader->ReadEntry("data");
    VerifyPattern(data_entry.data, 256);
}

TEST_F(IndexEntryWriterV3Test, EmptyMetaRoundtrip) {
    const std::string file_path = kV3FilePath + "_emptymeta";
    auto data = GeneratePattern(128);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("entry", data.data(), data.size());
        // No PutMeta calls - meta_json_ defaults to empty JSON object
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    // __meta__ entry should exist with empty JSON object "{}"
    auto meta_entry = reader->ReadEntry("__meta__");
    std::string meta_str(reinterpret_cast<const char*>(meta_entry.data.data()),
                         meta_entry.data.size());
    EXPECT_EQ(meta_str, "{}");

    auto data_entry = reader->ReadEntry("entry");
    VerifyPattern(data_entry.data, 128);
}

TEST_F(IndexEntryWriterV3Test, ReadEntriesToFilesMultiRangeParallel) {
    // Test ReadEntriesToFiles with multiple large entries, each requiring
    // multiple ranges (> 16MB each). This tests the parallel download path
    // where all ranges from all entries are submitted at once.
    const std::string file_path = kV3FilePath + "_multirange";

    // Each entry > 16MB to require multiple ranges
    const size_t size_a = 35 * 1024 * 1024;  // 35 MB = 3 ranges
    const size_t size_b = 50 * 1024 * 1024;  // 50 MB = 4 ranges
    const size_t size_c = 20 * 1024 * 1024;  // 20 MB = 2 ranges
    // Total: 9 parallel range tasks across 3 entries

    auto data_a = GeneratePattern(size_a);
    auto data_b = GeneratePattern(size_b);
    auto data_c = GeneratePattern(size_c);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("entry_a", data_a.data(), data_a.size());
        writer.WriteEntry("entry_b", data_b.data(), data_b.size());
        writer.WriteEntry("entry_c", data_c.data(), data_c.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    std::string file_a = GetRootPath() + "/multirange_a.bin";
    std::string file_b = GetRootPath() + "/multirange_b.bin";
    std::string file_c = GetRootPath() + "/multirange_c.bin";

    std::vector<std::pair<std::string, std::string>> pairs = {
        {"entry_a", file_a}, {"entry_b", file_b}, {"entry_c", file_c}};
    reader->ReadEntriesToFiles(pairs);

    auto verify_file = [](const std::string& path, size_t expected_size) {
        std::ifstream ifs(path, std::ios::binary | std::ios::ate);
        ASSERT_TRUE(ifs.is_open()) << "Failed to open: " << path;
        size_t read_size = ifs.tellg();
        ASSERT_EQ(read_size, expected_size);
        ifs.seekg(0);
        std::vector<uint8_t> buf(read_size);
        ifs.read(reinterpret_cast<char*>(buf.data()), read_size);
        VerifyPattern(buf, expected_size);
    };

    verify_file(file_a, size_a);
    verify_file(file_b, size_b);
    verify_file(file_c, size_c);

    ::unlink(file_a.c_str());
    ::unlink(file_b.c_str());
    ::unlink(file_c.c_str());
}

// =============================================================================
// Edge Case Tests: Directory Table and Meta Entry size boundaries
// =============================================================================

TEST_F(IndexEntryWriterV3Test, LargeDirectoryTableNeedsSecondIO) {
    // Create many entries with long names to make directory table > 64KB
    // This tests the second IO path in ReadFooterAndDirectory()
    const std::string file_path = kV3FilePath + "_largedir";
    auto small_data = GeneratePattern(64);

    // Each entry in directory table is roughly:
    // {"name":"...", "offset":N, "size":N, "crc32":"XXXXXXXX"}
    // ~60 bytes per entry + name length
    // Need ~1100 entries with 50-char names to exceed 64KB
    const size_t num_entries = 1200;
    const std::string name_prefix =
        "entry_with_a_very_long_name_to_inflate_dir_";

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        for (size_t i = 0; i < num_entries; i++) {
            std::string name = name_prefix + std::to_string(i);
            writer.WriteEntry(name, small_data.data(), small_data.size());
        }
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    auto names = reader->GetEntryNames();
    // +1 for __meta__
    ASSERT_EQ(names.size(), num_entries + 1);

    // Verify a few entries
    auto entry0 = reader->ReadEntry(name_prefix + "0");
    VerifyPattern(entry0.data, 64);

    auto entry_last = reader->ReadEntry(name_prefix + "1199");
    VerifyPattern(entry_last.data, 64);
}

TEST_F(IndexEntryWriterV3Test, DirectoryFitsButMetaDoesNotFitFirstIO) {
    // Directory table fits in first 64KB, but meta entry is large enough
    // that dir_size + meta_entry_size > 64KB - 32
    // This tests the boundary where we need second IO for meta only
    const std::string file_path = kV3FilePath + "_largemetasmalldir";

    // Create a large meta JSON
    // We need meta_entry_size to be large (> ~60KB)
    // but directory table should be small (< 64KB - 32 - meta_size)
    auto data = GeneratePattern(1024);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());

        // Add a large meta value (~70KB of repeated string)
        std::string large_value(70 * 1024, 'X');
        writer.PutMeta("large_field", large_value);
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    // Verify meta can be read
    auto meta_value = reader->GetMeta<std::string>("large_field");
    ASSERT_EQ(meta_value.size(), 70 * 1024);
    ASSERT_EQ(meta_value, std::string(70 * 1024, 'X'));

    // Verify data entry
    auto entry = reader->ReadEntry("data");
    VerifyPattern(entry.data, 1024);
}

TEST_F(IndexEntryWriterV3Test, LargeMetaEntryMultiRange) {
    // Meta entry > 16MB requires multiple range reads
    const std::string file_path = kV3FilePath + "_hugemetamultirange";

    auto data = GeneratePattern(256);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());

        // Create meta with ~20MB of data (will need 2 ranges to read)
        std::string huge_value(20 * 1024 * 1024, 'Y');
        writer.PutMeta("huge_field", huge_value);
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    // Verify large meta
    auto meta_value = reader->GetMeta<std::string>("huge_field");
    ASSERT_EQ(meta_value.size(), 20 * 1024 * 1024);
    ASSERT_EQ(meta_value[0], 'Y');
    ASSERT_EQ(meta_value[meta_value.size() - 1], 'Y');

    // Verify data entry
    auto entry = reader->ReadEntry("data");
    VerifyPattern(entry.data, 256);
}

// =============================================================================
// Encrypted Entry Tests (using MockCipherPlugin)
// =============================================================================

class IndexEntryEncryptedV3Test : public IndexEntryWriterV3Test {
 protected:
    void
    SetUp() override {
        IndexEntryWriterV3Test::SetUp();
        mock_cipher_ = std::make_shared<MockCipherPlugin>();
    }

    std::shared_ptr<MockCipherPlugin> mock_cipher_;
};

TEST_F(IndexEntryEncryptedV3Test, EncryptedSmallEntryRoundtrip) {
    const std::string file_path = kV3FilePath + "_enc_small";
    const size_t entry_size = 1024;
    auto data = GeneratePattern(entry_size);

    const size_t slice_size = kStreamSliceAlignment;

    {
        // Note: remote_path should be relative to fs root
        IndexEntryEncryptedLocalWriter writer(file_path,
                                              fs_,
                                              mock_cipher_,
                                              /*ez_id=*/1,
                                              /*collection_id=*/100,
                                              GetRootPath(),
                                              slice_size);
        writer.WriteEntry("enc_entry", data.data(), data.size());
        writer.Finish();
    }

    // Verify file was created
    auto info = fs_->GetFileInfo(file_path);
    ASSERT_TRUE(info.ok());
    EXPECT_GT(info.ValueOrDie().size(), entry_size);  // ciphertext > plaintext
}

TEST_F(IndexEntryEncryptedV3Test, EncryptedWriterRejectsUnalignedSliceSize) {
    const std::string file_path = kV3FilePath + "_enc_unaligned_slice";
    EXPECT_THROW(IndexEntryEncryptedLocalWriter writer(file_path,
                                                       fs_,
                                                       mock_cipher_,
                                                       /*ez_id=*/1,
                                                       /*collection_id=*/100,
                                                       GetRootPath(),
                                                       1024),
                 milvus::SegcoreError);
}

TEST_F(IndexEntryEncryptedV3Test, EncryptedMultiSliceEntry) {
    // Entry larger than slice_size, requiring multiple slices
    const std::string file_path = kV3FilePath + "_enc_multislice";
    const size_t slice_size = kStreamSliceAlignment;
    const size_t entry_size = 5 * slice_size + 100;
    auto data = GeneratePattern(entry_size);

    {
        IndexEntryEncryptedLocalWriter writer(file_path,
                                              fs_,
                                              mock_cipher_,
                                              /*ez_id=*/1,
                                              /*collection_id=*/100,
                                              GetRootPath(),
                                              slice_size);
        writer.WriteEntry("large_enc", data.data(), data.size());
        writer.Finish();
    }

    auto info = fs_->GetFileInfo(file_path);
    ASSERT_TRUE(info.ok());
    EXPECT_GT(info.ValueOrDie().size(), entry_size);
}

TEST_F(IndexEntryEncryptedV3Test, EncryptedMultipleEntriesMultiSlice) {
    // Multiple entries, each requiring multiple slices
    const std::string file_path = kV3FilePath + "_enc_multi_multi";
    const size_t slice_size = kStreamSliceAlignment;

    const size_t size_a = 3 * slice_size + 500;
    const size_t size_b = 2 * slice_size + 100;
    const size_t size_c = 5 * slice_size;

    auto data_a = GeneratePattern(size_a);
    auto data_b = GeneratePattern(size_b);
    auto data_c = GeneratePattern(size_c);

    {
        IndexEntryEncryptedLocalWriter writer(file_path,
                                              fs_,
                                              mock_cipher_,
                                              /*ez_id=*/1,
                                              /*collection_id=*/100,
                                              GetRootPath(),
                                              slice_size);
        writer.WriteEntry("entry_a", data_a.data(), data_a.size());
        writer.WriteEntry("entry_b", data_b.data(), data_b.size());
        writer.WriteEntry("entry_c", data_c.data(), data_c.size());
        writer.Finish();
    }

    auto info = fs_->GetFileInfo(file_path);
    ASSERT_TRUE(info.ok());
    EXPECT_GT(info.ValueOrDie().size(), size_a + size_b + size_c);
}

TEST_F(IndexEntryEncryptedV3Test, EncryptedLargeMetaMultiSlice) {
    // Large meta entry requiring multiple slices
    const std::string file_path = kV3FilePath + "_enc_largemeta";
    const size_t slice_size = kStreamSliceAlignment;

    auto data = GeneratePattern(256);

    {
        IndexEntryEncryptedLocalWriter writer(file_path,
                                              fs_,
                                              mock_cipher_,
                                              /*ez_id=*/1,
                                              /*collection_id=*/100,
                                              GetRootPath(),
                                              slice_size);
        writer.WriteEntry("data", data.data(), data.size());

        // Meta is larger than one slice.
        std::string meta_value(5000, 'M');
        writer.PutMeta("large_meta", meta_value);
        writer.Finish();
    }

    auto info = fs_->GetFileInfo(file_path);
    ASSERT_TRUE(info.ok());
}

TEST_F(IndexEntryEncryptedV3Test, EncryptedFdEntryMultiSlice) {
    // Test fd-based entry with multiple slices
    const std::string file_path = kV3FilePath + "_enc_fd";
    const size_t slice_size = kStreamSliceAlignment;
    const size_t entry_size = 4 * slice_size + 200;
    auto data = GeneratePattern(entry_size);

    // Write source file
    std::string tmp_relative = "enc_fd_source.bin";
    std::string tmp_absolute = GetRootPath() + "/" + tmp_relative;
    {
        auto tmp_out = CreateOutputStream(tmp_relative);
        tmp_out->Write(data.data(), data.size());
        tmp_out->Close();
    }

    int fd = ::open(tmp_absolute.c_str(), O_RDONLY);
    ASSERT_NE(fd, -1);

    {
        IndexEntryEncryptedLocalWriter writer(file_path,
                                              fs_,
                                              mock_cipher_,
                                              /*ez_id=*/1,
                                              /*collection_id=*/100,
                                              GetRootPath(),
                                              slice_size);
        writer.WriteEntry("fd_entry", fd, entry_size);
        writer.Finish();
    }
    ::close(fd);

    auto info = fs_->GetFileInfo(file_path);
    ASSERT_TRUE(info.ok());
    EXPECT_GT(info.ValueOrDie().size(), entry_size);

    ::unlink(tmp_absolute.c_str());
}

// ---- ReadEntryStream tests ----

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamLarge) {
    const std::string file_path = kV3FilePath + "_stream_large";
    const size_t entry_size = 4 * 1024 * 1024;  // 4MB
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("test_data", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    std::vector<uint8_t> reassembled;
    size_t slice_count = 0;
    reader->ReadEntryStream(
        "test_data",
        [&](const uint8_t* d, size_t len) {
            reassembled.insert(reassembled.end(), d, d + len);
            slice_count++;
        },
        1 * 1024 * 1024);  // 1MB slices

    ASSERT_EQ(reassembled.size(), entry_size);
    VerifyPattern(reassembled, entry_size);
    ASSERT_EQ(slice_count, 4);  // 4MB / 1MB = 4 slices
}

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamSmall) {
    const std::string file_path = kV3FilePath + "_stream_small";
    const size_t entry_size = 1024;
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("small_data", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    std::vector<uint8_t> reassembled;
    size_t slice_count = 0;
    reader->ReadEntryStream("small_data", [&](const uint8_t* d, size_t len) {
        reassembled.insert(reassembled.end(), d, d + len);
        slice_count++;
    });

    ASSERT_EQ(reassembled.size(), entry_size);
    VerifyPattern(reassembled, entry_size);
    ASSERT_EQ(slice_count, 1);  // single slice (1KB < default slice size)
}

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamRejectsInvalidSliceSize) {
    const std::string file_path = kV3FilePath + "_stream_invalid_slice";
    const size_t entry_size = kMinStreamSliceSize;
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    EXPECT_THROW(reader->ReadEntryStream(
                     "data", [](const uint8_t*, size_t) {}, 0),
                 milvus::SegcoreError);
    EXPECT_THROW(
        reader->ReadEntryStream(
            "data", [](const uint8_t*, size_t) {}, kMinStreamSliceSize - 1),
        milvus::SegcoreError);
    EXPECT_THROW(
        reader->ReadEntryStream(
            "data", [](const uint8_t*, size_t) {}, kMinStreamSliceSize + 1),
        milvus::SegcoreError);
}

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamUsesDefaultSliceSize) {
    IndexEntryStreamConfigGuard guard;
    const size_t slice_size = DEFAULT_INDEX_FILE_SLICE_SIZE;
    auto& budget = TransientMemoryBudget::GetLoadTransientBudget();

    ASSERT_EQ(DefaultStreamSliceSize(), slice_size);
    milvus::SetLoadTransientBudgetBytes(0);
    ASSERT_EQ(budget.CapacityBytes(), 0);
    ASSERT_EQ(EntryStreamMaxTransientBytes(),
              std::numeric_limits<size_t>::max());

    const size_t configured_budget = 3 * slice_size;
    milvus::SetLoadTransientBudgetBytes(
        static_cast<int64_t>(configured_budget));
    ASSERT_EQ(budget.CapacityBytes(), configured_budget);
    ASSERT_EQ(EntryStreamMaxTransientBytes(),
              configured_budget + kTailMergeGrace);

    const std::string file_path = kV3FilePath + "_stream_configured_default";
    const size_t tail_size = kTailMergeGrace + 17;
    const size_t entry_size = 2 * slice_size + tail_size;
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    std::vector<size_t> slice_sizes;
    std::vector<uint8_t> reassembled;
    reader->ReadEntryStream("data", [&](const uint8_t* d, size_t len) {
        slice_sizes.push_back(len);
        reassembled.insert(reassembled.end(), d, d + len);
    });

    ASSERT_EQ(reassembled, data);
    ASSERT_EQ(slice_sizes,
              (std::vector<size_t>{slice_size, slice_size, tail_size}));
}

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamMergesSmallTail) {
    const size_t slice_size = DEFAULT_INDEX_FILE_SLICE_SIZE;
    const size_t tail_size = kTailMergeGrace / 2;
    ASSERT_GT(tail_size, 0);
    const std::string file_path = kV3FilePath + "_stream_merge_small_tail";
    const size_t entry_size = slice_size + tail_size;
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    std::vector<size_t> slice_sizes;
    std::vector<uint8_t> reassembled;
    reader->ReadEntryStream("data", [&](const uint8_t* d, size_t len) {
        slice_sizes.push_back(len);
        reassembled.insert(reassembled.end(), d, d + len);
    });

    ASSERT_EQ(reassembled, data);
    ASSERT_EQ(slice_sizes, (std::vector<size_t>{entry_size}));
}

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamEmptyEntryVerifiesCrc) {
    const std::string file_path = kV3FilePath + "_stream_empty";
    std::vector<uint8_t> data;

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("empty", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    size_t slice_count = 0;
    reader->ReadEntryStream("empty",
                            [&](const uint8_t*, size_t) { slice_count++; });
    ASSERT_EQ(slice_count, 0);
}

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamMatchesReadEntry) {
    const std::string file_path = kV3FilePath + "_stream_match";
    const size_t entry_size = 5 * 1024 * 1024;  // 5MB, non-aligned
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    // Read via ReadEntry
    auto entry = reader->ReadEntry("data");

    // Read via ReadEntryStream
    std::vector<uint8_t> streamed;
    reader->ReadEntryStream(
        "data",
        [&](const uint8_t* d, size_t len) {
            streamed.insert(streamed.end(), d, d + len);
        },
        512 * 1024);  // 512KB slices

    ASSERT_EQ(entry.data.size(), streamed.size());
    EXPECT_EQ(entry.data, streamed);
}

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamToFileWritesEntry) {
    const std::string file_path = kV3FilePath + "_stream_to_file";
    const size_t entry_size = 3 * 1024 * 1024 + 19;
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    std::string local_file = GetRootPath() + "/stream_entry_output.bin";
    reader->ReadEntryStreamToFile(
        "data", local_file, milvus::storage::io::Priority::HIGH);

    std::ifstream ifs(local_file, std::ios::binary | std::ios::ate);
    ASSERT_TRUE(ifs.is_open());
    size_t read_size = ifs.tellg();
    ASSERT_EQ(read_size, entry_size);

    ifs.seekg(0);
    std::vector<uint8_t> read_data(read_size);
    ifs.read(reinterpret_cast<char*>(read_data.data()), read_size);
    VerifyPattern(read_data, entry_size);

    ::unlink(local_file.c_str());
}

TEST_F(IndexEntryWriterV3Test, ReadEntriesStreamToFilesRunsFilesConcurrently) {
    const std::string file_path = kV3FilePath + "_stream_files_parallel";
    const size_t entry_size = 1024 * 1024;
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("entry_a", data.data(), data.size());
        writer.WriteEntry("entry_b", data.data(), data.size());
        writer.WriteEntry("entry_c", data.data(), data.size());
        writer.Finish();
    }

    auto input = std::make_shared<TrackingDelayedInputStream>(
        CreateInputStream(file_path),
        entry_size,
        std::chrono::milliseconds(50));
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);
    if (milvus::ThreadPools::GetThreadPool(milvus::ThreadPoolPriority::HIGH)
            .GetMaxThreadNum() <= 1) {
        GTEST_SKIP() << "HIGH load thread pool has only one worker";
    }

    std::string file_a = GetRootPath() + "/stream_parallel_a.bin";
    std::string file_b = GetRootPath() + "/stream_parallel_b.bin";
    std::string file_c = GetRootPath() + "/stream_parallel_c.bin";
    std::vector<std::pair<std::string, std::string>> pairs = {
        {"entry_a", file_a}, {"entry_b", file_b}, {"entry_c", file_c}};

    input->EnableTracking();
    reader->ReadEntriesStreamToFiles(pairs,
                                     milvus::storage::io::Priority::HIGH);

    EXPECT_GT(input->MaxActiveReads(), 1);

    auto verify_file = [](const std::string& path, size_t expected_size) {
        std::ifstream ifs(path, std::ios::binary | std::ios::ate);
        ASSERT_TRUE(ifs.is_open()) << "Failed to open: " << path;
        size_t read_size = ifs.tellg();
        ASSERT_EQ(read_size, expected_size);
        ifs.seekg(0);
        std::vector<uint8_t> read_data(read_size);
        ifs.read(reinterpret_cast<char*>(read_data.data()), read_size);
        VerifyPattern(read_data, expected_size);
    };
    verify_file(file_a, entry_size);
    verify_file(file_b, entry_size);
    verify_file(file_c, entry_size);

    ::unlink(file_a.c_str());
    ::unlink(file_b.c_str());
    ::unlink(file_c.c_str());
}

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamConsumerExceptionDoesNotLeak) {
    const std::string file_path = kV3FilePath + "_stream_consumer_throw";
    const size_t slice_size = 64 * 1024;
    const size_t entry_size = 4 * slice_size;
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    size_t slice_count = 0;
    EXPECT_THROW(reader->ReadEntryStream(
                     "data",
                     [&](const uint8_t*, size_t) {
                         slice_count++;
                         throw std::runtime_error("stop stream");
                     },
                     slice_size),
                 std::runtime_error);
    EXPECT_EQ(slice_count, 1);

    std::vector<uint8_t> streamed;
    reader->ReadEntryStream(
        "data",
        [&](const uint8_t* d, size_t len) {
            streamed.insert(streamed.end(), d, d + len);
        },
        slice_size);
    EXPECT_EQ(streamed, data);
}

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamDrainsActiveTasksAfterError) {
    const std::string file_path = kV3FilePath + "_stream_active_task_error";
    const size_t slice_size = 64 * 1024;
    const size_t entry_size = 3 * slice_size;
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());
        writer.Finish();
    }

    auto base_input = CreateInputStream(file_path);
    auto input = std::make_shared<DelayedFailingInputStream>(
        base_input,
        std::vector<DelayedFailingInputStream::Rule>{
            {MILVUS_V3_MAGIC_SIZE, std::chrono::milliseconds(80), false},
            {MILVUS_V3_MAGIC_SIZE + slice_size,
             std::chrono::milliseconds(40),
             true},
        });
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    size_t slice_count = 0;
    EXPECT_THROW(
        reader->ReadEntryStream(
            "data", [&](const uint8_t*, size_t) { slice_count++; }, slice_size),
        milvus::SegcoreError);
    EXPECT_EQ(slice_count, 1);

    auto clean_input = CreateInputStream(file_path);
    auto clean_reader = IndexEntryReader::Open(clean_input, file_size);
    std::vector<uint8_t> streamed;
    clean_reader->ReadEntryStream(
        "data",
        [&](const uint8_t* d, size_t len) {
            streamed.insert(streamed.end(), d, d + len);
        },
        slice_size);
    EXPECT_EQ(streamed, data);
}

TEST_F(IndexEntryWriterV3Test, ReadEntryStreamCancellationWhileWaitingBudget) {
    const std::string file_path = kV3FilePath + "_stream_cancel_budget_wait";
    const size_t slice_size = kMinStreamSliceSize;
    auto data = GeneratePattern(slice_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("data", data.data(), data.size());
        writer.Finish();
    }

    auto& budget = TransientMemoryBudget::GetLoadTransientBudget();
    auto old_capacity = budget.CapacityBytes();
    budget.SetCapacityBytes(slice_size);
    budget.Acquire(slice_size);
    auto cleanup = folly::makeGuard([&budget, old_capacity, slice_size]() {
        budget.Release(slice_size);
        budget.SetCapacityBytes(old_capacity);
    });

    folly::CancellationSource source;
    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(
        input, file_size, 0, milvus::HIGH, source.getToken());

    std::atomic<size_t> slice_count{0};
    auto future = std::async(std::launch::async, [&]() {
        reader->ReadEntryStream(
            "data",
            [&slice_count](const uint8_t*, size_t) {
                slice_count.fetch_add(1);
            },
            slice_size);
    });

    EXPECT_EQ(future.wait_for(std::chrono::milliseconds(50)),
              std::future_status::timeout);
    EXPECT_EQ(slice_count.load(), 0);

    source.requestCancellation();

    ASSERT_EQ(future.wait_for(std::chrono::seconds(2)),
              std::future_status::ready);
    try {
        future.get();
        FAIL() << "expected cancellation";
    } catch (const milvus::SegcoreError& e) {
        EXPECT_EQ(e.get_error_code(), milvus::ErrorCode::FollyCancel);
    }
    EXPECT_EQ(slice_count.load(), 0);
}

TEST_F(IndexEntryWriterV3Test, GetEntrySize) {
    const std::string file_path = kV3FilePath + "_entry_size";
    const size_t entry_size = 3 * 1024 * 1024;
    auto data = GeneratePattern(entry_size);

    {
        auto output = CreateOutputStream(file_path);
        IndexEntryDirectStreamWriter writer(output);
        writer.WriteEntry("sized_entry", data.data(), data.size());
        writer.Finish();
    }

    auto input = CreateInputStream(file_path);
    int64_t file_size = GetFileSize(file_path);
    auto reader = IndexEntryReader::Open(input, file_size);

    ASSERT_EQ(reader->GetEntrySize("sized_entry"), entry_size);
}
