// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed 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.

#pragma once

#include "db/snapshot/Operations.h"

namespace milvus {
namespace engine {
namespace snapshot {

class CollectionCommitOperation : public CommitOperation<CollectionCommit> {
 public:
    using BaseT = CommitOperation<CollectionCommit>;
    CollectionCommitOperation(OperationContext context, ScopedSnapshotT prev_ss) : BaseT(context, prev_ss) {
    }
    CollectionCommitOperation(OperationContext context, ID_TYPE collection_id, ID_TYPE commit_id = 0)
        : BaseT(context, collection_id, commit_id) {
    }

    CollectionCommitPtr
    GetPrevResource() const override {
        return prev_ss_->GetCollectionCommit();
    }

    Status
    DoExecute(Store&) override;
};

class PartitionCommitOperation : public CommitOperation<PartitionCommit> {
 public:
    using BaseT = CommitOperation<PartitionCommit>;
    PartitionCommitOperation(const OperationContext& context, ScopedSnapshotT prev_ss);
    PartitionCommitOperation(const OperationContext& context, ID_TYPE collection_id, ID_TYPE commit_id = 0);

    PartitionCommitPtr
    GetPrevResource() const override;

    Status
    DoExecute(Store&) override;

    Status
    PreCheck() override;
};

class PartitionOperation : public CommitOperation<Partition> {
 public:
    using BaseT = CommitOperation<Partition>;
    PartitionOperation(const PartitionContext& context, ScopedSnapshotT prev_ss);
    PartitionOperation(const PartitionContext& context, ID_TYPE collection_id, ID_TYPE commit_id = 0);

    Status
    DoExecute(Store& store) override;

    Status
    PreCheck() override;

 protected:
    PartitionContext context_;
};

class SegmentCommitOperation : public CommitOperation<SegmentCommit> {
 public:
    using BaseT = CommitOperation<SegmentCommit>;
    SegmentCommitOperation(const OperationContext& context, ScopedSnapshotT prev_ss);
    SegmentCommitOperation(const OperationContext& context, ID_TYPE collection_id, ID_TYPE commit_id = 0);

    SegmentCommit::Ptr
    GetPrevResource() const override;

    Status
    DoExecute(Store&) override;

    Status
    PreCheck() override;
};

class SegmentOperation : public CommitOperation<Segment> {
 public:
    using BaseT = CommitOperation<Segment>;
    SegmentOperation(const OperationContext& context, ScopedSnapshotT prev_ss);
    SegmentOperation(const OperationContext& context, ID_TYPE collection_id, ID_TYPE commit_id = 0);

    Status
    DoExecute(Store& store) override;

    Status
    PreCheck() override;
};

class SegmentFileOperation : public CommitOperation<SegmentFile> {
 public:
    using BaseT = CommitOperation<SegmentFile>;
    SegmentFileOperation(const SegmentFileContext& sc, ScopedSnapshotT prev_ss);
    SegmentFileOperation(const SegmentFileContext& sc, ID_TYPE collection_id, ID_TYPE commit_id = 0);

    Status
    DoExecute(Store& store) override;

 protected:
    SegmentFileContext context_;
};

template <>
class LoadOperation<Collection> : public Operations {
 public:
    explicit LoadOperation(const LoadOperationContext& context)
        : Operations(OperationContext(), ScopedSnapshotT(), OperationsType::O_Leaf), context_(context) {
    }

    Status
    ApplyToStore(Store& store) override {
        if (done_) {
            Done();
            return status_;
        }
        Status status;
        if (context_.id == 0 && context_.name != "") {
            status = store.GetCollection(context_.name, resource_);
        } else {
            status = store.GetResource<Collection>(context_.id, resource_);
        }
        SetStatus(status);
        Done();
        return status_;
    }

    Status
    GetResource(CollectionPtr& res, bool wait = true) {
        if (wait) {
            WaitToFinish();
        }
        auto status = DoneRequired();
        if (!status.ok())
            return status;
        if (!resource_) {
            return Status(SS_NOT_FOUND_ERROR, "No specified resource");
        }
        res = resource_;
        return status;
    }

 protected:
    LoadOperationContext context_;
    CollectionPtr resource_;
};

}  // namespace snapshot
}  // namespace engine
}  // namespace milvus
