// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include <vespa/searchcore/proton/bucketdb/bucketdbhandler.h>
#include <vespa/searchcore/proton/bucketdb/bucket_create_notifier.h>
#include <vespa/searchcore/proton/test/bucketfactory.h>
#include <vespa/searchcore/proton/feedoperation/moveoperation.h>
#include <vespa/searchcore/proton/server/i_move_operation_limiter.h>
#include <vespa/searchcore/proton/server/idocumentmovehandler.h>
#include <vespa/searchcore/proton/server/imaintenancejobrunner.h>
#include <vespa/searchcore/proton/server/maintenancedocumentsubdb.h>
#include <vespa/searchcore/proton/server/ibucketmodifiedhandler.h>
#include <vespa/searchcore/proton/server/i_maintenance_job.h>
#include <vespa/searchcore/proton/test/buckethandler.h>
#include <vespa/searchcore/proton/test/clusterstatehandler.h>
#include <vespa/searchcore/proton/test/disk_mem_usage_notifier.h>
#include <vespa/searchcore/proton/test/test.h>
#include <vespa/document/test/make_bucket_space.h>

namespace proton {
    class DocumentMetaStore;
}
namespace proton::move::test {

struct MyMoveOperationLimiter : public IMoveOperationLimiter {
    uint32_t beginOpCount;

    MyMoveOperationLimiter() : beginOpCount(0) {}

    vespalib::IDestructorCallback::SP beginOperation() override {
        ++beginOpCount;
        return {};
    }
    size_t numPending() const override { return beginOpCount; }
};

struct MyMoveHandler : public IDocumentMoveHandler {
    using MoveOperationVector = std::vector<MoveOperation>;
    bucketdb::BucketDBOwner &_bucketDb;
    MoveOperationVector _moves;
    size_t _numCachedBuckets;
    bool _storeMoveDoneContexts;
    std::vector<vespalib::IDestructorCallback::SP> _moveDoneContexts;

    MyMoveHandler(bucketdb::BucketDBOwner &bucketDb, bool storeMoveDoneContext = false);
    ~MyMoveHandler() override;
    void handleMove(MoveOperation &op, vespalib::IDestructorCallback::SP moveDoneCtx) override;

    void reset() {
        _moves.clear();
        _numCachedBuckets = 0;
    }

    void clearMoveDoneContexts() {
        _moveDoneContexts.clear();
    }
};

struct MyDocumentRetriever : public DocumentRetrieverBaseForTest {
    using DocumentTypeRepo = document::DocumentTypeRepo;
    using DocumentMetaData = search::DocumentMetaData;
    using Document = document::Document;
    using DocumentId = document::DocumentId;
    using DocumentIdT = search::DocumentIdT;
    using DocumentVector = std::vector<Document::SP>;
    std::shared_ptr<const DocumentTypeRepo> _repo;
    DocumentVector _docs;
    uint32_t _lid2Fail;

    MyDocumentRetriever(std::shared_ptr<const DocumentTypeRepo> repo)
        : _repo(std::move(repo)),
          _docs(),
          _lid2Fail(0)
    {
        _docs.push_back(Document::UP()); // lid 0 invalid
    }

    const DocumentTypeRepo &getDocumentTypeRepo() const override { return *_repo; }

    void getBucketMetaData(const storage::spi::Bucket &, DocumentMetaData::Vector &) const override {}

    DocumentMetaData getDocumentMetaData(const DocumentId &) const override { return DocumentMetaData(); }

    Document::UP getFullDocument(DocumentIdT lid) const override {
        return (lid != _lid2Fail) ? Document::UP(_docs[lid]->clone()) : Document::UP();
    }

    void failRetrieveForLid(uint32_t lid) { _lid2Fail = lid; }

    CachedSelect::SP parseSelect(const vespalib::string &) const override {
        return {};
    }
};

struct MyBucketModifiedHandler : public IBucketModifiedHandler {
    using BucketId = document::BucketId;
    std::vector<BucketId> _modified;

    ~MyBucketModifiedHandler() override;
    void notifyBucketModified(const BucketId &bucket) override;

    void reset() { _modified.clear(); }
};

struct MySubDb {
    using BucketId = document::BucketId;
    using Document = document::Document;
    using DocumentTypeRepo = document::DocumentTypeRepo;
    using DocumentVector = proton::test::DocumentVector;
    using UserDocuments = proton::test::UserDocuments;
    std::shared_ptr<DocumentMetaStore>    _metaStoreSP;
    DocumentMetaStore                    &_metaStore;
    std::shared_ptr<MyDocumentRetriever>  _realRetriever;
    std::shared_ptr<IDocumentRetriever>   _retriever;
    MaintenanceDocumentSubDB              _subDb;
    UserDocuments                         _docs;
    bucketdb::BucketDBHandler             _bucketDBHandler;

    MySubDb(const std::shared_ptr<const DocumentTypeRepo> &repo, std::shared_ptr<bucketdb::BucketDBOwner> bucketDB,
            uint32_t subDbId, SubDbType subDbType);

    ~MySubDb();

    void insertDocs(const UserDocuments &docs_);

    void failRetrieveForLid(uint32_t lid) {
        _realRetriever->failRetrieveForLid(lid);
    }

    BucketId bucket(uint32_t userId) const {
        return _docs.getBucket(userId);
    }

    DocumentVector docs(uint32_t userId) {
        return _docs.getGidOrderDocs(userId);
    }

    void setBucketState(const BucketId &bucketId, bool active);
};

struct MyCountJobRunner : public IMaintenanceJobRunner {
    uint32_t runCount;
    explicit MyCountJobRunner(IMaintenanceJob &job) : runCount(0) {
        job.registerRunner(this);
    }
    void run() override { ++runCount; }
};

bool
assertEqual(const document::BucketId &bucket, const proton::test::Document &doc,
            uint32_t sourceSubDbId, uint32_t targetSubDbId, const MoveOperation &op);

}