// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 "MilvusClientImpl.h"

#include <nlohmann/json.hpp>

#include "TypeUtils.h"
#include "common.pb.h"
#include "milvus.pb.h"
#include "schema.pb.h"
#include "types/Utils.h"

namespace VectorDB {
using namespace milvus;

MilvusClientImpl::~MilvusClientImpl() { Disconnect(); }

Status MilvusClientImpl::Connect(const ConnectParam& param) {
    if (connection_ != nullptr) {
        connection_->Disconnect();
    }

    connection_ = std::make_shared<MilvusConnection>();
    int retryTime = 0;
    Status status = connection_->Connect(param);
    while (status.Code() == StatusCode::NOT_CONNECTED && retryTime <= RETRY_TIME) {
        if (retryTime == RETRY_TIME) {
            std::cerr << "ERROR connect milvus retry failed, code: " << static_cast<int>(status.Code())
                      << ", msg: " << status.Message() << std::endl;
            return status;
        }
        std::cout << "WARN connect milvus failed, code: " << static_cast<int>(status.Code())
                  << ", msg: " << status.Message() << ", retry: " << retryTime << std::endl;

        sleep(1);
        retryTime++;
        status = connection_->Connect(param);
    }
    return status;
}

Status MilvusClientImpl::Disconnect() {
    if (connection_ != nullptr) {
        return connection_->Disconnect();
    }

    return Status::OK();
}

Status MilvusClientImpl::CreateCollection(const CollectionSchema& schema, const IndexDesc& index_desc) {
    // create collection
    auto status = CreateCollection(schema);
    if (!status.IsOk()) {
        return status;
    }

    // create index
    status = CreateIndex(schema.Name(), index_desc);
    if (!status.IsOk()) {
        return status;
    }

    // load collection，目前先不加载，看查询性能，后续再修改
    // LoadCollection(schema.Name());

    return Status();
}

Status MilvusClientImpl::CreateCollection(const std::string& collection_name, int dim, bool auto_id,
                                          bool enable_dynamic_field) {
    // create collection schema
    CollectionSchema schema(collection_name, "default schema", enable_dynamic_field);
    schema.AddField({DEFAULT_ID_FIELD_NAME, DataType::INT64, "user id", true, auto_id});
    schema.AddField(FieldSchema(DEFAULT_VECTOR_FIELD_NAME, DataType::FLOAT_VECTOR, "vector").WithDimension(dim));

    auto status = CreateCollection(schema);
    if (!status.IsOk()) {
        return status;
    }

    // create index
    IndexDesc desc(DEFAULT_VECTOR_FIELD_NAME, "vector_index", 0);
    status = CreateIndex(schema.Name(), desc);
    if (!status.IsOk()) {
        return status;
    }

    // load collection，目前先不加载，看查询性能，后续再修改
    // LoadCollection(schema.Name());

    return Status();
}

Status MilvusClientImpl::CreateCollection(const CollectionSchema& schema) {
    auto pre = [&schema]() {
        proto::milvus::CreateCollectionRequest rpc_request;
        rpc_request.set_collection_name(schema.Name());
        rpc_request.set_shards_num(schema.ShardsNum());
        rpc_request.set_consistency_level(proto::common::ConsistencyLevel::Bounded);

        proto::schema::CollectionSchema rpc_collection;
        rpc_collection.set_name(schema.Name());
        rpc_collection.set_description(schema.Description());
        rpc_collection.set_enable_dynamic_field(schema.EnableDynamicField());

        for (auto& field : schema.Fields()) {
            proto::schema::FieldSchema* rpc_field = rpc_collection.add_fields();
            rpc_field->set_name(field.Name());
            rpc_field->set_description(field.Description());
            rpc_field->set_data_type(static_cast<proto::schema::DataType>(field.FieldDataType()));
            rpc_field->set_is_primary_key(field.IsPrimaryKey());
            rpc_field->set_autoid(field.AutoID());
            rpc_field->set_is_dynamic(field.IsDynamic());
            rpc_field->set_is_clustering_key(false);
            rpc_field->set_is_partition_key(false);

            if (!field.TypeParams().empty()) {
                proto::common::KeyValuePair* kv = rpc_field->add_type_params();
                for (auto& pair : field.TypeParams()) {
                    kv->set_key(pair.first);
                    kv->set_value(pair.second);
                }
            }
        }
        // std::cout << "rpc collection: " << rpc_collection.DebugString() << std::endl;
        std::string binary;
        rpc_collection.SerializeToString(&binary);
        rpc_request.set_schema(binary);
        return rpc_request;
    };

    return apiHandler<proto::milvus::CreateCollectionRequest, proto::common::Status>(
        pre, &MilvusConnection::CreateCollection);
}

Status MilvusClientImpl::HasCollection(const std::string& collection_name, bool& has) {
    auto pre = [&collection_name]() {
        proto::milvus::HasCollectionRequest rpc_request;
        rpc_request.set_collection_name(collection_name);
        rpc_request.set_time_stamp(0);
        return rpc_request;
    };

    auto post = [&has](const proto::milvus::BoolResponse& response) { has = response.value(); };

    return apiHandler<proto::milvus::HasCollectionRequest, proto::milvus::BoolResponse>(
        pre, &MilvusConnection::HasCollection, post);
}

Status MilvusClientImpl::DropCollection(const std::string& collection_name) {
    auto pre = [&collection_name]() {
        proto::milvus::DropCollectionRequest rpc_request;
        rpc_request.set_collection_name(collection_name);
        return rpc_request;
    };

    return apiHandler<proto::milvus::DropCollectionRequest, proto::common::Status>(pre,
                                                                                   &MilvusConnection::DropCollection);
}

Status MilvusClientImpl::LoadCollection(const std::string& collection_name) {
    auto pre = [&collection_name]() {
        proto::milvus::LoadCollectionRequest rpc_request;
        rpc_request.set_collection_name(collection_name);
        // rpc_request.set_replica_number(replica_number);
        return rpc_request;
    };

    // todo 要不要查询加载进度
    return apiHandler<proto::milvus::LoadCollectionRequest, proto::common::Status>(pre,
                                                                                   &MilvusConnection::LoadCollection);
}

Status MilvusClientImpl::ReleaseCollection(const std::string& collection_name) {
    auto pre = [&collection_name]() {
        proto::milvus::ReleaseCollectionRequest rpc_request;
        rpc_request.set_collection_name(collection_name);
        return rpc_request;
    };

    return apiHandler<proto::milvus::ReleaseCollectionRequest, proto::common::Status>(
        pre, &MilvusConnection::ReleaseCollection);
}

Status MilvusClientImpl::DescribeCollection(const std::string& collection_name, CollectionDesc& collection_desc) {
    auto pre = [&collection_name]() {
        proto::milvus::DescribeCollectionRequest rpc_request;
        rpc_request.set_collection_name(collection_name);
        return rpc_request;
    };

    auto post = [&collection_desc](const proto::milvus::DescribeCollectionResponse& response) {
        CollectionSchema schema;
        ConvertCollectionSchema(response.schema(), schema);
        schema.SetShardsNum(response.shards_num());
        collection_desc.SetSchema(std::move(schema));
        collection_desc.SetID(response.collectionid());

        std::vector<std::string> aliases;
        aliases.reserve(response.aliases_size());
        aliases.insert(aliases.end(), response.aliases().begin(), response.aliases().end());

        collection_desc.SetAlias(std::move(aliases));
        collection_desc.SetCreatedTime(response.created_timestamp());
    };

    return apiHandler<proto::milvus::DescribeCollectionRequest, proto::milvus::DescribeCollectionResponse>(
        pre, &MilvusConnection::DescribeCollection, post);
}

Status MilvusClientImpl::CreateIndex(const std::string& collection_name, const IndexDesc& index_desc) {
    auto validate = [&index_desc] { return index_desc.Validate(); };

    auto pre = [&collection_name, &index_desc]() {
        proto::milvus::CreateIndexRequest rpc_request;
        rpc_request.set_collection_name(collection_name);
        rpc_request.set_field_name(index_desc.FieldName());

        auto kv_pair = rpc_request.add_extra_params();
        kv_pair->set_key(KeyIndexType());
        kv_pair->set_value(std::to_string(index_desc.IndexType()));

        kv_pair = rpc_request.add_extra_params();
        kv_pair->set_key(KeyMetricType());
        kv_pair->set_value(std::to_string(index_desc.MetricType()));

        kv_pair = rpc_request.add_extra_params();
        kv_pair->set_key(KeyParams());
        kv_pair->set_value(index_desc.ExtraParams());

        return rpc_request;
    };

    // todo 要不要查询加载进度
    return apiHandler<proto::milvus::CreateIndexRequest, proto::common::Status>(validate, pre,
                                                                                &MilvusConnection::CreateIndex);
}

Status MilvusClientImpl::DropIndex(const std::string& collection_name, const std::string& field_name) {
    auto pre = [&collection_name, &field_name]() {
        proto::milvus::DropIndexRequest rpc_request;
        rpc_request.set_collection_name(collection_name);
        rpc_request.set_field_name(field_name);
        return rpc_request;
    };
    return apiHandler<proto::milvus::DropIndexRequest, proto::common::Status>(pre, &MilvusConnection::DropIndex);
}

Status MilvusClientImpl::Insert(const std::string& collection_name, const std::vector<FieldDataPtr>& fields,
                                DmlResults& results) {
    // TODO(matrixji): add common validations check for fields
    // TODO(matrixji): add scheme based validations check for fields
    // auto validate = nullptr;

    auto pre = [&collection_name, &fields] {
        proto::milvus::InsertRequest rpc_request;

        rpc_request.set_collection_name(collection_name);
        rpc_request.set_num_rows((*fields.front()).Count());
        for (const auto& field : fields) {
            proto::schema::FieldData data = CreateProtoFieldData(*field);
            rpc_request.add_fields_data()->CopyFrom(data);
        }
        // std::cout << rpc_request.DebugString() << std::endl;
        return rpc_request;
    };

    auto post = [&results](const proto::milvus::MutationResult& response) {
        auto id_array = CreateIDArray(response.ids());
        results.SetIdArray(std::move(id_array));
        results.SetTimestamp(response.timestamp());
    };

    return apiHandler<proto::milvus::InsertRequest, proto::milvus::MutationResult>(pre, &MilvusConnection::Insert,
                                                                                   post);
}

Status MilvusClientImpl::Delete(const std::string& collection_name, const std::string& expression,
                                DmlResults& results) {
    auto pre = [&collection_name, &expression]() {
        proto::milvus::DeleteRequest rpc_request;
        rpc_request.set_collection_name(collection_name);
        rpc_request.set_expr(expression);
        return rpc_request;
    };

    auto post = [&results](const proto::milvus::MutationResult& response) {
        auto id_array = CreateIDArray(response.ids());
        results.SetIdArray(std::move(id_array));
        results.SetTimestamp(response.timestamp());
    };

    return apiHandler<proto::milvus::DeleteRequest, proto::milvus::MutationResult>(pre, &MilvusConnection::Delete,
                                                                                   post);
}

Status MilvusClientImpl::Upsert(const std::string& collection_name, const std::vector<FieldDataPtr>& fields,
                                DmlResults& results) {
    auto pre = [&collection_name, &fields]() {
        proto::milvus::UpsertRequest rpc_request;
        rpc_request.set_collection_name(collection_name);
        rpc_request.set_num_rows((*fields.front()).Count());
        for (const auto& field : fields) {
            proto::schema::FieldData data = CreateProtoFieldData(*field);
            rpc_request.add_fields_data()->CopyFrom(data);
        }
        return rpc_request;
    };

    auto post = [&results](const proto::milvus::MutationResult& response) {
        auto id_array = CreateIDArray(response.ids());
        results.SetIdArray(std::move(id_array));
        results.SetTimestamp(response.timestamp());
    };

    return apiHandler<proto::milvus::UpsertRequest, proto::milvus::MutationResult>(pre, &MilvusConnection::Upsert,
                                                                                   post);
}

Status MilvusClientImpl::Search(const SearchArguments& arguments, SearchResults& results, int timeout) {
    std::string anns_field;
    auto validate = [this, &arguments, &anns_field]() {
        CollectionDesc collection_desc;
        auto status = DescribeCollection(arguments.CollectionName(), collection_desc);
        if (status.IsOk()) {
            // check anns fields
            auto& field_name = arguments.TargetVectors()->Name();
            auto anns_fileds = collection_desc.Schema().AnnsFieldNames();
            if (anns_fileds.find(field_name) != anns_fileds.end()) {
                anns_field = field_name;
            } else {
                return Status{StatusCode::INVALID_AGUMENT, std::string(field_name + " is not a valid anns field")};
            }
            // basic check for extra params
            status = arguments.Validate();
        }
        return status;
    };

    auto pre = [&arguments, &anns_field]() {
        proto::milvus::SearchRequest rpc_request;
        rpc_request.set_collection_name(arguments.CollectionName());
        rpc_request.set_dsl_type(proto::common::DslType::BoolExprV1);
        if (!arguments.Expression().empty()) {
            rpc_request.set_dsl(arguments.Expression());
        }
        for (const auto& partition_name : arguments.PartitionNames()) {
            rpc_request.add_partition_names(partition_name);
        }
        for (const auto& output_field : arguments.OutputFields()) {
            rpc_request.add_output_fields(output_field);
        }

        // placeholders
        proto::common::PlaceholderGroup placeholder_group;
        auto& placeholder_value = *placeholder_group.add_placeholders();
        placeholder_value.set_tag("$0");
        auto target = arguments.TargetVectors();
        if (target->Type() == DataType::BINARY_VECTOR) {
            // bins
            placeholder_value.set_type(proto::common::PlaceholderType::BinaryVector);
            auto& bins_vec = dynamic_cast<BinaryVecFieldData&>(*target);
            for (const auto& bins : bins_vec.Data()) {
                std::string placeholder_data(reinterpret_cast<const char*>(bins.data()), bins.size());
                placeholder_value.add_values(std::move(placeholder_data));
            }
            rpc_request.set_nq(bins_vec.Data().size());
        } else {
            // floats
            placeholder_value.set_type(proto::common::PlaceholderType::FloatVector);
            auto& floats_vec = dynamic_cast<FloatVecFieldData&>(*target);
            for (const auto& floats : floats_vec.Data()) {
                std::string placeholder_data(reinterpret_cast<const char*>(floats.data()),
                                             floats.size() * sizeof(float));
                placeholder_value.add_values(std::move(placeholder_data));
            }
            rpc_request.set_nq(floats_vec.Data().size());
        }
        rpc_request.set_placeholder_group(std::move(placeholder_group.SerializeAsString()));

        auto kv_pair = rpc_request.add_search_params();
        kv_pair->set_key("anns_field");
        kv_pair->set_value(anns_field);

        kv_pair = rpc_request.add_search_params();
        kv_pair->set_key("topk");
        kv_pair->set_value(std::to_string(arguments.TopK()));

        kv_pair = rpc_request.add_search_params();
        kv_pair->set_key(KeyMetricType());
        kv_pair->set_value(std::to_string(arguments.MetricType()));

        kv_pair = rpc_request.add_search_params();
        kv_pair->set_key("round_decimal");
        kv_pair->set_value(std::to_string(arguments.RoundDecimal()));

        kv_pair = rpc_request.add_search_params();
        kv_pair->set_key(KeyParams());
        // merge extra params with range search
        auto json = nlohmann::json::parse(arguments.ExtraParams());
        if (arguments.RangeSearch()) {
            json["range_filter"] = arguments.RangeFilter();
            json["radius"] = arguments.Radius();
        }
        kv_pair->set_value(json.dump());

        rpc_request.set_travel_timestamp(arguments.TravelTimestamp());
        rpc_request.set_guarantee_timestamp(arguments.GuaranteeTimestamp());
        return rpc_request;
    };

    auto post = [&results](const proto::milvus::SearchResults& response) {
        auto& result_data = response.results();
        const auto& ids = result_data.ids();
        const auto& scores = result_data.scores();
        const auto& fields_data = result_data.fields_data();
        auto num_of_queries = result_data.num_queries();
        std::vector<int> topks{};
        topks.reserve(result_data.topks_size());
        for (int i = 0; i < result_data.topks_size(); ++i) {
            topks.emplace_back(result_data.topks(i));
        }
        std::vector<SingleResult> single_results;
        single_results.reserve(num_of_queries);
        int offset{0};
        for (int i = 0; i < num_of_queries; ++i) {
            std::vector<float> item_scores;
            std::vector<FieldDataPtr> item_field_data;
            auto item_topk = topks[i];
            item_scores.reserve(item_topk);
            for (int j = 0; j < item_topk; ++j) {
                item_scores.emplace_back(scores.Get(offset + j));
            }
            item_field_data.reserve(fields_data.size());
            for (const auto& field_data : fields_data) {
                item_field_data.emplace_back(std::move(CreateMilvusFieldData(field_data, offset, item_topk)));
            }
            single_results.emplace_back(std::move(CreateIDArray(ids, offset, item_topk)), std::move(item_scores),
                                        std::move(item_field_data));
            offset += item_topk;
        }

        results = std::move(SearchResults(std::move(single_results)));
    };

    return apiHandler<proto::milvus::SearchRequest, proto::milvus::SearchResults>(
        validate, pre, &MilvusConnection::Search, nullptr, post, GrpcOpts{timeout});
}

Status MilvusClientImpl::Query(const QueryArguments& arguments, QueryResults& results, int timeout) {
    auto pre = [&arguments]() {
        proto::milvus::QueryRequest rpc_request;
        rpc_request.set_collection_name(arguments.CollectionName());
        for (const auto& partition_name : arguments.PartitionNames()) {
            rpc_request.add_partition_names(partition_name);
        }

        rpc_request.set_expr(arguments.Expression());
        for (const auto& field : arguments.OutputFields()) {
            rpc_request.add_output_fields(field);
        }

        rpc_request.set_travel_timestamp(arguments.TravelTimestamp());
        rpc_request.set_guarantee_timestamp(arguments.GuaranteeTimestamp());
        return rpc_request;
    };

    auto post = [&results](const proto::milvus::QueryResults& response) {
        std::vector<FieldDataPtr> return_fields{};
        return_fields.reserve(response.fields_data_size());
        for (const auto& field_data : response.fields_data()) {
            return_fields.emplace_back(std::move(CreateMilvusFieldData(field_data)));
        }

        results = std::move(QueryResults(std::move(return_fields)));
    };
    return apiHandler<proto::milvus::QueryRequest, proto::milvus::QueryResults>(pre, &MilvusConnection::Query, post,
                                                                                GrpcOpts{timeout});
}

}  // namespace VectorDB
