#include <unistd.h>
#include <uuid/uuid.h>

#include <iostream>
#include <nlohmann/json.hpp>
#include <random>
#include <string>

#include "MilvusClientImpl.h"

using namespace std;
using namespace VectorDB;

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

std::shared_ptr<MilvusClientImpl> client = std::make_shared<MilvusClientImpl>();
const string collection_name = "TEST";
const string field_id_name = "identity";
const string field_age_name = "age";
const string field_face_name = "face";
const string field_metadata_name = "metadata";
const uint32_t dimension = 4;
const int64_t row_count = 5;
vector<std::string> insert_ids;
vector<std::string> insert_ages;
vector<std::vector<float>> insert_vectors;
vector<nlohmann::json> insert_jsons;
default_random_engine ran(time(nullptr));

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 HasCollection() {
    bool exist;
    auto status = client->HasCollection(collection_name, exist);
    CheckStatus("Failed to has TEST collection:", status);
    cout << "Have TEST collection or not: " << exist << endl;
}

void CreateCollection() {
    CollectionSchema collection_schema(collection_name, "++++----", true);
    collection_schema.AddField(FieldSchema(field_id_name, DataType::VARCHAR, "userid", true, false).WithMaxLength(50));
    collection_schema.AddField(FieldSchema(field_age_name, DataType::VARCHAR, "userage").WithMaxLength(1024));
    collection_schema.AddField(
        FieldSchema(field_face_name, DataType::FLOAT_VECTOR, "facesignature").WithDimension(dimension));

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

void ShowCollection() {
    CollectionDesc desc;
    auto status = client->DescribeCollection(collection_name, desc);
    CheckStatus("Failed to describe TEST collection:", status);
    cout << "--------------------" << endl;
    cout << "Have TEST collection desc: " << desc.Schema().ShardsNum() << endl
         << desc.Schema().Name() << endl
         << desc.Schema().Description() << endl;

    for (const auto& item : desc.Schema().AnnsFieldNames()) {
        cout << "AnnsFieldNames: " << item << endl;
    }
    cout << "--------------------" << endl;
    for (const auto& item : desc.Schema().Fields()) {
        cout << "Name: " << item.Name() << endl
             << "Description: " << item.Description() << endl
             << "FieldDataType: " << static_cast<int>(item.FieldDataType()) << endl
             << "AutoID: " << item.AutoID() << endl
             << "Dimension: " << item.Dimension() << endl
             << "IsPrimaryKey: " << item.IsPrimaryKey() << endl;
        cout << "++++++++++++++++++++++++" << endl;
        for (const auto& tmp : item.TypeParams()) {
            cout << "key: " << tmp.first << ", val: " << tmp.second << endl;
        }
    }
}

void CreateIndex() {
    IndexDesc index_desc(field_face_name, "", 0);
    auto status = client->CreateIndex(collection_name, index_desc);
    CheckStatus("Failed to create index:", status);
    cout << "Create index success." << endl;
}

void LoadCollection() {
    auto status = client->LoadCollection(collection_name);
    CheckStatus("Failed to load collection:", status);
    cout << "Load collection success." << endl;
}

void ReleaseCollection() {
    auto status = client->ReleaseCollection(collection_name);
    CheckStatus("Failed to release collection:", status);
    cout << "Release collection success." << endl;
}

std::string generateUUID() {
    uuid_t uuid;
    char uuidStr[37];

    uuid_generate_random(uuid);
    uuid_unparse(uuid, uuidStr);

    return uuidStr;
}

void Insert() {
    uniform_real_distribution<float> float_gen(0.0, 1.0);
    string seed = "abcdefghijklmn";
    for (auto i = 0; i < row_count; ++i) {
        insert_ids.push_back(generateUUID());
        insert_ages.push_back(std::to_string(i));
        vector<float> vector(dimension);

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

        nlohmann::json value = nlohmann::json::object();
        string key = std::to_string(seed[i]);
        value[key] = 100 * (i + 1);
        insert_jsons.push_back(value);
    }

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

void Search() {
    SearchArguments arguments(collection_name, 10, MetricType::COSINE);
    arguments.AddOutputField(field_age_name);
    arguments.AddOutputField(field_metadata_name);
    arguments.SetExpression(field_age_name + " > '1'");
    // 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. find results: " << search_results.Results().size() << endl;

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

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

        auto metadata_field = result.OutputField(field_metadata_name);
        JsonFieldDataPtr metadata_field_ptr = static_pointer_cast<JsonFieldData>(metadata_field);
        auto& metadata = metadata_field_ptr->Data();

        for (size_t i = 0; i < ids.size(); ++i) {
            cout << "ID: " << ids[i] << "\tDistance: " << distances[i] << "\tAge: " << age_data[i]
                 << "\tMetadata: " << metadata[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;
            // }
        }
    }
}

void Query() {
    QueryArguments queryArguments;
    queryArguments.SetCollectionName(collection_name);
    queryArguments.AddOutputField(field_age_name);
    queryArguments.SetExpression(field_age_name + " == '2'");

    QueryResults queryResults;
    auto status = client->Query(queryArguments, queryResults, 100);
    CheckStatus("Failed to query:", status);
    cout << "Query success." << endl;

    FieldDataPtr idField = queryResults.GetFieldByName(field_id_name);
    shared_ptr<VarCharFieldData> idFieldPtr = static_pointer_cast<VarCharFieldData>(idField);
    for (const auto& item : idFieldPtr->Data()) {
        cout << idFieldPtr->Name() << ": " << item << endl;
    }

    FieldDataPtr ageField = queryResults.GetFieldByName(field_age_name);
    shared_ptr<VarCharFieldData> ageFieldPtr = static_pointer_cast<VarCharFieldData>(ageField);
    for (const auto& item : ageFieldPtr->Data()) {
        cout << ageFieldPtr->Name() << ": " << item << endl;
    }
}

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

    Connect();

    // ShowCollection();

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

    // create a collection
    CreateCollection();
    // CreateCollectionTmp();
    ShowCollection();

    // create index
    CreateIndex();

    // insert some rows
    Insert();

    // tell server prepare to load collection
    // LoadCollection();

    // query by id
    Query();

    // do search
    Search();
    // SearchTmp();

    // release collection after search
    ReleaseCollection();

    // drop collection
    HasCollection();
    DropCollection();
    HasCollection();  // Failed to has TEST collection: : collection not found

    return 0;
}
