#include "collection.h"
#include "connection.h"
#include "gridfs.h"
#include "exception.h"
#include "bson.h"

#include <library/cpp/json/json_value.h>
#include <library/cpp/json/json_reader.h>
#include <library/cpp/logger/global/global.h>

#include <contrib/libs/mongo-c-driver/libmongoc/src/mongoc/mongoc.h>

#include <util/string/vector.h>
#include <util/string/split.h>

namespace NUtil {

    TMongoCollection::TIterator::TIterator(mongoc_cursor_t* cursor, TMongoCollection* owner)
        : Cursor(cursor)
        , Owner(owner)
    {
        CHECK_WITH_LOG(!!Cursor);
        Next();
    }

    bool TMongoCollection::TIterator::IsValid() const {
        return !!Data.Get();
    }

    void TMongoCollection::TIterator::Next() {
        const bson_t* doc;
        Data.Reset(nullptr);
        if (mongoc_cursor_more(Cursor) && mongoc_cursor_next(Cursor, &doc)) {
            auto bsonStr = bson_as_json(doc, nullptr);
            if (!!bsonStr) {
                Data.Reset(new NJson::TJsonValue);
                NJson::ReadJsonFastTree(bsonStr, Data.Get(), true);
                bson_free(bsonStr);
            } else {
                ythrow TMongoException() << "Can't get json";
            }
        }
        bson_error_t error;
        if (mongoc_cursor_error(Cursor, &error)) {
            ythrow TMongoException(error);
        }
    }

    const NJson::TJsonValue& TMongoCollection::TIterator::Get() const {
        CHECK_WITH_LOG(!!IsValid());
        return *Data.Get();
    }

    TMongoCollection::TIterator::~TIterator() {
        Owner->RemoveIter();
        mongoc_cursor_destroy(Cursor);
    }

    TMongoCollection::TIterator::TPtr TMongoCollection::Find(const NJson::TJsonValue& query) {
        auto it = new TIterator(mongoc_collection_find(Collection, MONGOC_QUERY_NONE, 0, 0, 0, NPrivate::TBson(query).Get(), nullptr, nullptr), this);
        AddIter();
        return it;
    }

    TMongoCollection::TIterator::TPtr TMongoCollection::Find(const NJson::TJsonValue& query, const TString& fields) {
        TVector<TString> keys;
        StringSplitter(fields).Split(',').SkipEmpty().Collect(&keys);
        NJson::TJsonValue selector(NJson::JSON_MAP);
        for (auto&& key: keys) {
            selector[key] = 1;
        }
        auto it = new TIterator(mongoc_collection_find(Collection, MONGOC_QUERY_NONE, 0, 0, 0, NPrivate::TBson(query).Get(), NPrivate::TBson(selector).Get(), nullptr), this);
        AddIter();
        return it;
    }

    TMongoCollection::TMongoCollection(TMongoConnection* connection, const TString& db, const TString& name)
        : AliveIters(0)
    {
        CHECK_WITH_LOG(!!connection);
        CHECK_WITH_LOG(db.size() + name.size() < 128);
        Collection = mongoc_client_get_collection(connection->GetRawConnection(), db.data(), name.data());
    }

    TMongoCollection::TMongoCollection(TMongoGridFs* gridFs)
        : AliveIters(0)
    {
        CHECK_WITH_LOG(!!gridFs);
        Collection = mongoc_gridfs_get_files(gridFs->GetRawFs());
    }

    TMongoCollection::~TMongoCollection() {
        CHECK_WITH_LOG(AliveIters == 0);
        mongoc_collection_destroy(Collection);
    }

    TMongoCollection& TMongoCollection::Insert(const NJson::TJsonValue& doc) {
        bson_error_t error;
        if (!mongoc_collection_insert(Collection, MONGOC_INSERT_NONE, NPrivate::TBson(doc).Get(), nullptr, &error)) {
            ythrow TMongoException(error);
        }
        return *this;
    }

    TMongoCollection& TMongoCollection::Update(const NJson::TJsonValue& query, const NJson::TJsonValue& doc) {
        bson_error_t error;
        if (!mongoc_collection_update(Collection, MONGOC_UPDATE_UPSERT, NPrivate::TBson(query).Get(), NPrivate::TBson(doc).Get(), nullptr, &error)) {
            ythrow TMongoException(error);
        }
        return *this;
    }

    void TMongoCollection::Remove(const NJson::TJsonValue& query, bool removeAll) {
        bson_error_t error;
        if (!mongoc_collection_remove(Collection, removeAll ? MONGOC_REMOVE_NONE : MONGOC_REMOVE_SINGLE_REMOVE, NPrivate::TBson(query).Get(), nullptr, &error)) {
            ythrow TMongoException(error);
        }
    }

    void TMongoCollection::AddIter() {
        AliveIters++;
    }

    void TMongoCollection::RemoveIter() {
        CHECK_WITH_LOG(AliveIters > 0);
        AliveIters--;
    }
}
