// 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 <folly/FBVector.h>
#include <gtest/gtest.h>
#include <cstdint>
#include <cstdlib>
#include <filesystem>
#include <iostream>
#include <memory>
#include <numeric>
#include <string>
#include <system_error>
#include <utility>
#include <vector>

#include "bitset/bitset.h"
#include "common/Types.h"
#include "common/protobuf_utils.h"
#include "gtest/gtest.h"
#include "index/Index.h"
#include "index/IndexFactory.h"
#include "index/IndexInfo.h"
#include "index/Meta.h"
#include "index/ScalarIndex.h"
#include "index/StringIndex.h"
#include "index/StringIndexMarisa.h"
#include "knowhere/dataset.h"
#include "pb/common.pb.h"
#include "pb/schema.pb.h"
#include "storage/ChunkManager.h"
#include "storage/LocalChunkManagerSingleton.h"
#include "storage/Types.h"
#include "storage/Util.h"
#include "test_utils/AssertUtils.h"
#include "test_utils/Constants.h"
#include "test_utils/indexbuilder_test_utils.h"

constexpr int64_t nb = 100;
namespace schemapb = milvus::proto::schema;

namespace milvus {
namespace index {

static storage::FileManagerContext
CreateStringTestFileManagerContext(const std::string& root_path) {
    std::filesystem::create_directories(root_path);
    storage::StorageConfig storage_config;
    storage_config.storage_type = "local";
    storage_config.root_path = root_path;
    auto chunk_manager = storage::CreateChunkManager(storage_config);
    auto fs = storage::InitArrowFileSystem(storage_config);
    storage::FieldDataMeta field_meta{1, 2, 3, 101};
    field_meta.field_schema.set_data_type(proto::schema::DataType::VarChar);
    storage::IndexMeta index_meta{3, 101, 1000, 10000};
    storage::FileManagerContext ctx(field_meta, index_meta, chunk_manager, fs);
    return ctx;
}

static storage::FileManagerContext
CreateStringTestFileManagerContext() {
    return CreateStringTestFileManagerContext(TestLocalPath);
}

class StringIndexBaseTest : public ::testing::Test {
 protected:
    void
    SetUp() override {
        strs = GenStrArr(nb);
        *str_arr.mutable_data() = {strs.begin(), strs.end()};
    }

 protected:
    std::vector<std::string> strs;
    schemapb::StringArray str_arr;
};

class StringIndexMarisaTest : public StringIndexBaseTest {};

TEST_F(StringIndexMarisaTest, Constructor) {
    auto index = milvus::index::CreateStringIndexMarisa();
}

TEST_F(StringIndexMarisaTest, Build) {
    auto index = milvus::index::CreateStringIndexMarisa();
    index->Build(strs.size(), strs.data());
}

TEST_F(StringIndexMarisaTest, HasRawData) {
    auto index = milvus::index::CreateStringIndexMarisa();
    index->Build(nb, strs.data());
    ASSERT_TRUE(index->HasRawData());
}

TEST_F(StringIndexMarisaTest, Count) {
    auto index = milvus::index::CreateStringIndexMarisa();
    index->Build(nb, strs.data());
    ASSERT_EQ(strs.size(), index->Count());
}

TEST_F(StringIndexMarisaTest, In) {
    auto index = milvus::index::CreateStringIndexMarisa();
    index->Build(nb, strs.data());
    auto bitset = index->In(strs.size(), strs.data());
    ASSERT_EQ(bitset.size(), strs.size());
    ASSERT_TRUE(bitset.any());
}

TEST_F(StringIndexMarisaTest, InHasNull) {
    auto index = milvus::index::CreateStringIndexMarisa();
    FixedVector<bool> is_null(nb);
    std::vector<std::string> in_strs(nb / 2);
    int j = 0;
    for (int i = 0; i < nb; i++) {
        is_null[i] = i % 2 == 0;
        if (i % 2 == 0) {
            in_strs[j] = strs[i];
            j++;
        }
    }

    index->Build(nb, strs.data(), is_null.data());
    auto bitset = index->In(in_strs.size(), in_strs.data());
    ASSERT_EQ(bitset.size(), strs.size());
    ASSERT_TRUE(bitset.count() == (nb / 2))
        << "count: " << bitset.count() << " nb: " << nb;
}

TEST_F(StringIndexMarisaTest, NotIn) {
    auto index = milvus::index::CreateStringIndexMarisa();
    index->Build(nb, strs.data());
    auto bitset = index->NotIn(strs.size(), strs.data());
    ASSERT_EQ(bitset.size(), strs.size());
    ASSERT_TRUE(bitset.none());
}

TEST_F(StringIndexMarisaTest, NotInHasNull) {
    auto index = milvus::index::CreateStringIndexMarisa();
    FixedVector<bool> is_null(nb);
    std::vector<std::string> in_strs(nb / 2);
    int j = 0;
    for (int i = 0; i < nb; i++) {
        is_null[i] = i % 2 == 0;
        if (i % 2 == 0) {
            in_strs[j] = strs[i];
            j++;
        }
    }
    index->Build(nb, strs.data(), is_null.data());
    auto bitset = index->NotIn(in_strs.size(), in_strs.data());
    ASSERT_EQ(bitset.size(), strs.size());
    std::cout << "bitset: " << bitset.count() << std::endl;
    ASSERT_TRUE(bitset.count() == 0);
}

TEST_F(StringIndexMarisaTest, Range) {
    auto index = milvus::index::CreateStringIndexMarisa();
    std::vector<std::string> strings(nb);
    std::vector<int> counts(10);
    for (int i = 0; i < nb; ++i) {
        int val = std::rand() % 10;
        counts[val]++;
        strings[i] = std::to_string(val);
    }
    index->Build(nb, strings.data());

    {
        // [0...9]
        auto bitset = index->Range("0", milvus::OpType::GreaterEqual);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        // [5...9]
        int expect = std::accumulate(counts.begin() + 5, counts.end(), 0);
        auto bitset = index->Range("5", milvus::OpType::GreaterEqual);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), expect);
    }

    {
        // [6...9]
        int expect = std::accumulate(counts.begin() + 6, counts.end(), 0);
        auto bitset = index->Range("5", milvus::OpType::GreaterThan);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), expect);
    }

    {
        // [0...3]
        int expect = std::accumulate(counts.begin(), counts.begin() + 4, 0);
        auto bitset = index->Range("4", milvus::OpType::LessThan);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), expect);
    }

    {
        // [0...4]
        int expect = std::accumulate(counts.begin(), counts.begin() + 5, 0);
        auto bitset = index->Range("4", milvus::OpType::LessEqual);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), expect);
    }

    {
        // [2...8]
        int expect = std::accumulate(counts.begin() + 2, counts.begin() + 9, 0);
        auto bitset = index->Range("2", true, "8", true);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), expect);
    }

    {
        // [0...8]
        int expect = std::accumulate(counts.begin(), counts.begin() + 9, 0);
        auto bitset = index->Range("0", true, "9", false);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), expect);
    }
}

TEST_F(StringIndexMarisaTest, Reverse) {
    auto index_types = GetIndexTypes<std::string>();
    for (const auto& index_type : index_types) {
        CreateIndexInfo create_index_info;
        create_index_info.index_type = index_type;
        auto index =
            milvus::index::IndexFactory::GetInstance()
                .CreatePrimitiveScalarIndex<std::string>(create_index_info);
        index->Build(nb, strs.data());
        assert_reverse<std::string>(index.get(), strs);
    }
}

TEST_F(StringIndexMarisaTest, PrefixMatch) {
    auto index = milvus::index::CreateStringIndexMarisa();
    index->Build(nb, strs.data());

    for (size_t i = 0; i < strs.size(); i++) {
        auto str = strs[i];
        auto bitset = index->PrefixMatch(str);
        ASSERT_EQ(bitset.size(), strs.size());
        ASSERT_TRUE(bitset[i]);
    }
}

TEST_F(StringIndexMarisaTest, IsNull) {
    auto index = milvus::index::CreateStringIndexMarisa();
    index->Build(nb, strs.data());
    auto bitset = index->IsNull();
    ASSERT_EQ(bitset.size(), strs.size());
    ASSERT_TRUE(bitset.count() == 0);
}

TEST_F(StringIndexMarisaTest, IsNullHasNull) {
    auto index = milvus::index::CreateStringIndexMarisa();
    FixedVector<bool> is_null(nb);
    for (int i = 0; i < nb; i++) {
        is_null[i] = i % 2 == 0;
    }
    index->Build(nb, strs.data(), is_null.data());
    auto bitset = index->IsNull();
    ASSERT_EQ(bitset.size(), strs.size());
    ASSERT_TRUE(bitset.count() == (nb / 2))
        << "count: " << bitset.count() << " nb: " << nb;
}

TEST_F(StringIndexMarisaTest, IsNotNull) {
    auto index = milvus::index::CreateStringIndexMarisa();
    index->Build(nb, strs.data());
    auto bitset = index->IsNotNull();
    ASSERT_EQ(bitset.size(), strs.size());
    ASSERT_TRUE(bitset.count() == strs.size());
}

TEST_F(StringIndexMarisaTest, IsNotNullHasNull) {
    auto index = milvus::index::CreateStringIndexMarisa();
    FixedVector<bool> is_null(nb);
    for (int i = 0; i < nb; i++) {
        is_null[i] = i % 2 == 0;
    }
    index->Build(nb, strs.data(), is_null.data());
    auto bitset = index->IsNotNull();
    ASSERT_EQ(bitset.size(), strs.size());
    ASSERT_TRUE(bitset.count() == (nb / 2))
        << "count: " << bitset.count() << " nb: " << nb;
}

TEST_F(StringIndexMarisaTest, Query) {
    auto index = milvus::index::CreateStringIndexMarisa();
    index->Build(nb, strs.data());

    {
        auto ds = knowhere::GenDataSet(strs.size(), 8, strs.data());
        ds->Set<milvus::OpType>(milvus::index::OPERATOR_TYPE,
                                milvus::OpType::In);
        auto bitset = index->Query(ds);
        ASSERT_TRUE(bitset.any());
    }

    {
        auto ds = knowhere::GenDataSet(strs.size(), 8, strs.data());
        ds->Set<milvus::OpType>(milvus::index::OPERATOR_TYPE,
                                milvus::OpType::NotIn);
        auto bitset = index->Query(ds);
        ASSERT_TRUE(bitset.none());
    }

    {
        auto ds = std::make_shared<knowhere::DataSet>();
        ds->Set<milvus::OpType>(milvus::index::OPERATOR_TYPE,
                                milvus::OpType::GreaterEqual);
        ds->Set<std::string>(milvus::index::RANGE_VALUE, "0");
        auto bitset = index->Query(ds);
        ASSERT_EQ(bitset.size(), strs.size());
        ASSERT_EQ(bitset.count(), strs.size());
    }

    {
        auto ds = std::make_shared<knowhere::DataSet>();
        ds->Set<milvus::OpType>(milvus::index::OPERATOR_TYPE,
                                milvus::OpType::Range);
        ds->Set<std::string>(milvus::index::LOWER_BOUND_VALUE, "0");
        ds->Set<std::string>(milvus::index::UPPER_BOUND_VALUE, "range");
        ds->Set<bool>(milvus::index::LOWER_BOUND_INCLUSIVE, true);
        ds->Set<bool>(milvus::index::UPPER_BOUND_INCLUSIVE, true);
        auto bitset = index->Query(ds);
        ASSERT_TRUE(bitset.any());
    }

    {
        for (size_t i = 0; i < strs.size(); i++) {
            auto ds = std::make_shared<knowhere::DataSet>();
            ds->Set<milvus::OpType>(milvus::index::OPERATOR_TYPE,
                                    milvus::OpType::PrefixMatch);
            ds->Set<std::string>(milvus::index::MATCH_VALUE,
                                 std::move(strs[i]));
            auto bitset = index->Query(ds);
            ASSERT_EQ(bitset.size(), strs.size());
            ASSERT_TRUE(bitset[i]);
        }
    }
}

TEST_F(StringIndexMarisaTest, Codec) {
    auto file_manager_ctx = CreateStringTestFileManagerContext();
    auto index = milvus::index::CreateStringIndexMarisa(file_manager_ctx);
    std::vector<std::string> strings(nb);
    for (int i = 0; i < nb; ++i) {
        strings[i] = std::to_string(std::rand() % 10);
    }

    index->Build(nb, strings.data());

    std::vector<std::string> invalid_strings = {std::to_string(nb)};
    auto copy_index = milvus::index::CreateStringIndexMarisa(file_manager_ctx);

    {
        auto create_index_result = index->UploadUnified({});
        auto index_files = create_index_result->GetIndexFiles();
        Config load_config;
        load_config["index_files"] = index_files;
        load_config[milvus::LOAD_PRIORITY] =
            milvus::proto::common::LoadPriority::HIGH;
        copy_index->LoadUnified(load_config);
    }

    {
        auto bitset = copy_index->In(nb, strings.data());
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_TRUE(bitset.any());
    }

    {
        auto bitset = copy_index->In(1, invalid_strings.data());
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_TRUE(bitset.none());
    }

    {
        auto bitset = copy_index->NotIn(nb, strings.data());
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_TRUE(bitset.none());
    }

    {
        auto bitset = copy_index->Range("0", milvus::OpType::GreaterEqual);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        auto bitset = copy_index->Range("90", milvus::OpType::LessThan);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        auto bitset = copy_index->Range("9", milvus::OpType::LessEqual);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        auto bitset = copy_index->Range("0", true, "9", true);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        auto bitset = copy_index->Range("0", true, "90", false);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        for (size_t i = 0; i < nb; i++) {
            auto str = strings[i];
            auto bitset = copy_index->PrefixMatch(str);
            ASSERT_EQ(bitset.size(), nb);
            ASSERT_TRUE(bitset[i]);
        }
    }
}

TEST_F(StringIndexMarisaTest, UnifiedCodecRecreatesMissingLocalChunkDir) {
    auto file_manager_ctx = CreateStringTestFileManagerContext(
        TestRemotePath + "string_marisa_missing_local_chunk/");
    auto index = milvus::index::CreateStringIndexMarisa(file_manager_ctx);
    std::vector<std::string> strings(nb);
    for (int i = 0; i < nb; ++i) {
        strings[i] = std::to_string(std::rand() % 10);
    }

    index->Build(nb, strings.data());

    auto local_cm =
        storage::LocalChunkManagerSingleton::GetInstance().GetChunkManager();
    ASSERT_NE(local_cm, nullptr);
    auto local_root = local_cm->GetRootPath();
    ASSERT_FALSE(local_root.empty());

    std::error_code ec;
    std::filesystem::remove_all(local_root, ec);
    ASSERT_FALSE(ec) << ec.message();
    ASSERT_FALSE(std::filesystem::exists(local_root));

    auto create_index_result = index->UploadUnified({});
    ASSERT_TRUE(std::filesystem::is_directory(local_root));
    auto index_files = create_index_result->GetIndexFiles();

    std::filesystem::remove_all(local_root, ec);
    ASSERT_FALSE(ec) << ec.message();
    ASSERT_FALSE(std::filesystem::exists(local_root));

    auto copy_index = milvus::index::CreateStringIndexMarisa(file_manager_ctx);
    Config load_config;
    load_config["index_files"] = index_files;
    load_config[milvus::LOAD_PRIORITY] =
        milvus::proto::common::LoadPriority::HIGH;
    copy_index->LoadUnified(load_config);
    ASSERT_TRUE(std::filesystem::is_directory(local_root));

    auto bitset = copy_index->In(nb, strings.data());
    ASSERT_EQ(bitset.size(), nb);
    ASSERT_TRUE(bitset.any());
}

TEST_F(StringIndexMarisaTest, BaseIndexCodec) {
    auto file_manager_ctx = CreateStringTestFileManagerContext();
    milvus::index::IndexBasePtr index =
        milvus::index::CreateStringIndexMarisa(file_manager_ctx);
    std::vector<std::string> strings(nb);
    for (int i = 0; i < nb; ++i) {
        strings[i] = std::to_string(std::rand() % 10);
    }
    *str_arr.mutable_data() = {strings.begin(), strings.end()};
    std::vector<uint8_t> data(str_arr.ByteSizeLong(), 0);
    str_arr.SerializeToArray(data.data(), str_arr.ByteSizeLong());
    index->BuildWithRawDataForUT(str_arr.ByteSizeLong(), data.data());

    std::vector<std::string> invalid_strings = {std::to_string(nb)};
    auto copy_index = milvus::index::CreateStringIndexMarisa(file_manager_ctx);

    {
        auto create_index_result = index->UploadUnified({});
        auto index_files = create_index_result->GetIndexFiles();
        Config load_config;
        load_config["index_files"] = index_files;
        load_config[milvus::LOAD_PRIORITY] =
            milvus::proto::common::LoadPriority::HIGH;
        copy_index->LoadUnified(load_config);
    }

    {
        auto bitset = copy_index->In(nb, strings.data());
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_TRUE(bitset.any());
    }

    {
        auto bitset = copy_index->In(1, invalid_strings.data());
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_TRUE(bitset.none());
    }

    {
        auto bitset = copy_index->NotIn(nb, strings.data());
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_TRUE(bitset.none());
    }

    {
        auto bitset = copy_index->Range("0", milvus::OpType::GreaterEqual);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        auto bitset = copy_index->Range("90", milvus::OpType::LessThan);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        auto bitset = copy_index->Range("9", milvus::OpType::LessEqual);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        auto bitset = copy_index->Range("0", true, "9", true);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        auto bitset = copy_index->Range("0", true, "90", false);
        ASSERT_EQ(bitset.size(), nb);
        ASSERT_EQ(bitset.count(), nb);
    }

    {
        for (size_t i = 0; i < nb; i++) {
            auto str = strings[i];
            auto bitset = copy_index->PrefixMatch(str);
            ASSERT_EQ(bitset.size(), nb);
            ASSERT_TRUE(bitset[i]);
        }
    }
}
}  // namespace index
}  // namespace milvus
