#include <iostream>
#include <random>

#include "Database.h"

using namespace std;
using namespace VectorDB;

shared_ptr<Database> client = Database::Create();
const string collection_name = "TEST";
const string field_id_name = "identity";
const string field_age_name = "age";
const string field_face_name = "face";
const uint32_t dimension = 4;
const int64_t row_count = 10;
vector<int64_t> insert_ids;
vector<int8_t> insert_ages;
vector<std::vector<float>> insert_vectors;
default_random_engine ran(time(nullptr));

void CheckStatus(string&& prefix, const Status& status) {
    if (!status.IsOk()) {
        cout << prefix << " " << status.Message() << endl;
        exit(1);
    }
}

void Connect() {
    ConnectParam connect_param;
    auto status = client->Connect(connect_param);
    CheckStatus("Failed to connect milvus server:", status);
    cout << "Connect to milvus server success." << endl;
}

void DropCollection() {
    auto status = client->DropCollection(collection_name);
    CheckStatus("Failed to drop TEST collection:", status);
    cout << "Drop TEST collection success." << endl;
}

void CreateCollection() {
    // 通过 CollectionSchema 设置数据库列：identity（主键）、age、face（向量）
    CollectionSchema collection_schema(collection_name);
    collection_schema.AddField({field_id_name, DataType::INT64, "user id", true, false});
    collection_schema.AddField({field_age_name, DataType::INT8, "user age"});
    collection_schema.AddField(
        FieldSchema(field_face_name, DataType::FLOAT_VECTOR, "face signature").WithDimension(dimension));

    // 给向量列建个索引
    IndexDesc index_desc(field_face_name, "face_index", 0, IndexType::FLAT, MetricType::L2);

    auto status = client->CreateCollection(collection_schema, index_desc);
    CheckStatus("Failed to create collection:", status);
    cout << "Successfully create collection." << endl;
}

void Insert() {
    uniform_int_distribution<int> int_gen(1, 100);
    uniform_real_distribution<float> float_gen(0.0, 1.0);
    for (auto i = 0; i < row_count; ++i) {
        insert_ids.push_back(i);
        insert_ages.push_back(int_gen(ran));
        vector<float> vector(dimension);

        for (auto j = 0; j < dimension; ++j) {
            vector[j] = float_gen(ran);
        }
        insert_vectors.emplace_back(vector);
    }

    vector<FieldDataPtr> fields_data{make_shared<Int64FieldData>(field_id_name, insert_ids),
                                     make_shared<Int8FieldData>(field_age_name, insert_ages),
                                     make_shared<FloatVecFieldData>(field_face_name, insert_vectors)};
    DmlResults dml_results;
    auto status = client->Insert(collection_name, fields_data, dml_results);
    CheckStatus("Failed to insert:", status);
    cout << "Successfully insert " << dml_results.IdArray().IntIDArray().size() << " rows." << endl;
}

void Search() {
    SearchArguments arguments;
    arguments.SetCollectionName(collection_name);
    arguments.SetTopK(10);
    arguments.SetMetricType(MetricType::L2);  // 要与索引一致
    arguments.AddOutputField(field_age_name);
    arguments.SetExpression(field_age_name + " > 40");
    // set to strong guarantee so that the search is executed after the inserted data is persisted
    arguments.SetGuaranteeTimestamp(GuaranteeStrongTs());

    uniform_int_distribution<int64_t> int64_gen(0, row_count - 1);
    int64_t q_number = int64_gen(ran);
    vector<float> q_vector = insert_vectors[q_number];
    arguments.AddTargetVector(field_face_name, move(q_vector));
    cout << "Searching the No." << q_number << " entity..." << endl;

    SearchResults search_results{};
    auto status = client->Search(arguments, search_results, 100);
    CheckStatus("Failed to search:", status);
    cout << "Search success." << endl;

    for (auto& result : search_results.Results()) {
        auto& ids = result.Ids().IntIDArray();
        auto& distances = result.Scores();
        if (ids.size() != distances.size()) {
            cout << "Illegal result!" << endl;
            continue;
        }

        auto age_field = result.OutputField(field_age_name);
        Int8FieldDataPtr age_field_ptr = static_pointer_cast<Int8FieldData>(age_field);
        auto& age_data = age_field_ptr->Data();

        for (size_t i = 0; i < ids.size(); ++i) {
            cout << "ID: " << ids[i] << "\tDistance: " << distances[i] << "\tAge: " << static_cast<int32_t>(age_data[i])
                 << endl;
            // validate the age value
            if (insert_ages[ids[i]] != age_data[i]) {
                cout << "ERROR! The returned value doesn't match the inserted value" << endl;
            }
        }
    }
}

int main() {
    cout << "Example start..." << endl;

    Connect();

    // drop the collection if it exists
    DropCollection();

    CreateCollection();

    Insert();

    Search();

    return 0;
}
