/*
 * 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.
 */

package org.apache.bookkeeper.clients.impl.kv;

import static org.apache.bookkeeper.clients.impl.kv.KvUtils.toProtoCompare;
import static org.apache.bookkeeper.clients.impl.kv.KvUtils.toProtoRequest;

import com.google.common.collect.Lists;
import com.google.protobuf.UnsafeByteOperations;
import io.netty.buffer.ByteBuf;
import io.netty.util.ReferenceCountUtil;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.api.kv.PTable;
import org.apache.bookkeeper.api.kv.Txn;
import org.apache.bookkeeper.api.kv.impl.result.KeyValueFactory;
import org.apache.bookkeeper.api.kv.impl.result.ResultFactory;
import org.apache.bookkeeper.api.kv.op.CompareOp;
import org.apache.bookkeeper.api.kv.op.Op;
import org.apache.bookkeeper.api.kv.op.OpFactory;
import org.apache.bookkeeper.api.kv.options.DeleteOption;
import org.apache.bookkeeper.api.kv.options.IncrementOption;
import org.apache.bookkeeper.api.kv.options.PutOption;
import org.apache.bookkeeper.api.kv.options.RangeOption;
import org.apache.bookkeeper.api.kv.result.DeleteResult;
import org.apache.bookkeeper.api.kv.result.IncrementResult;
import org.apache.bookkeeper.api.kv.result.PutResult;
import org.apache.bookkeeper.api.kv.result.RangeResult;
import org.apache.bookkeeper.api.kv.result.TxnResult;
import org.apache.bookkeeper.clients.impl.container.StorageContainerChannel;
import org.apache.bookkeeper.common.util.Backoff;
import org.apache.bookkeeper.stream.proto.RangeProperties;
import org.apache.bookkeeper.stream.proto.kv.rpc.RoutingHeader;
import org.apache.bookkeeper.stream.proto.kv.rpc.TxnRequest;

/**
 * A range of a table.
 */
@Slf4j
class PByteBufTableRangeImpl implements PTable<ByteBuf, ByteBuf> {

    private final long streamId;
    private final RangeProperties rangeProps;
    private final StorageContainerChannel scChannel;
    private final ScheduledExecutorService executor;
    private final OpFactory<ByteBuf, ByteBuf> opFactory;
    private final ResultFactory<ByteBuf, ByteBuf> resultFactory;
    private final KeyValueFactory<ByteBuf, ByteBuf> kvFactory;
    private final Backoff.Policy backoffPolicy;

    PByteBufTableRangeImpl(long streamId,
                           RangeProperties rangeProps,
                           StorageContainerChannel scChannel,
                           ScheduledExecutorService executor,
                           OpFactory<ByteBuf, ByteBuf> opFactory,
                           ResultFactory<ByteBuf, ByteBuf> resultFactory,
                           KeyValueFactory<ByteBuf, ByteBuf> kvFactory,
                           Backoff.Policy backoffPolicy) {
        this.streamId = streamId;
        this.rangeProps = rangeProps;
        this.scChannel = scChannel;
        this.executor = executor;
        this.opFactory = opFactory;
        this.resultFactory = resultFactory;
        this.kvFactory = kvFactory;
        this.backoffPolicy = backoffPolicy;
    }

    private RoutingHeader.Builder newRoutingHeader(ByteBuf pKey) {
        return RoutingHeader.newBuilder()
            .setStreamId(streamId)
            .setRangeId(rangeProps.getRangeId())
            .setRKey(UnsafeByteOperations.unsafeWrap(pKey.nioBuffer()));
    }

    @Override
    public CompletableFuture<RangeResult<ByteBuf, ByteBuf>> get(
        ByteBuf pKey, ByteBuf lKey, RangeOption<ByteBuf> option) {
        pKey.retain();
        lKey.retain();
        if (null != option.endKey()) {
            option.endKey().retain();
        }
        return RangeRequestProcessor.of(
            KvUtils.newRangeRequest(lKey, option)
                .setHeader(newRoutingHeader(pKey))
                .build(),
            response -> KvUtils.newRangeResult(response, resultFactory, kvFactory),
            scChannel,
            executor,
            backoffPolicy
        ).process().whenComplete((value, cause) -> {
            ReferenceCountUtil.release(pKey);
            ReferenceCountUtil.release(lKey);
            if (null != option.endKey()) {
                ReferenceCountUtil.release(option.endKey());
            }
        });
    }

    @Override
    public CompletableFuture<PutResult<ByteBuf, ByteBuf>> put(ByteBuf pKey,
                                                              ByteBuf lKey,
                                                              ByteBuf value,
                                                              PutOption<ByteBuf> option) {
        pKey.retain();
        lKey.retain();
        value.retain();
        return PutRequestProcessor.of(
            KvUtils.newPutRequest(lKey, value, option)
                .setHeader(newRoutingHeader(pKey))
                .build(),
            response -> KvUtils.newPutResult(response, resultFactory, kvFactory),
            scChannel,
            executor,
            backoffPolicy
        ).process().whenComplete((ignored, cause) -> {
            ReferenceCountUtil.release(pKey);
            ReferenceCountUtil.release(lKey);
            ReferenceCountUtil.release(value);
        });
    }

    @Override
    public CompletableFuture<DeleteResult<ByteBuf, ByteBuf>> delete(ByteBuf pKey,
                                                                    ByteBuf lKey,
                                                                    DeleteOption<ByteBuf> option) {
        pKey.retain();
        lKey.retain();
        if (null != option.endKey()) {
            option.endKey().retain();
        }
        return DeleteRequestProcessor.of(
            KvUtils.newDeleteRequest(lKey, option)
                .setHeader(newRoutingHeader(pKey))
                .build(),
            response -> KvUtils.newDeleteResult(response, resultFactory, kvFactory),
            scChannel,
            executor,
            backoffPolicy
        ).process().whenComplete((ignored, cause) -> {
            ReferenceCountUtil.release(pKey);
            ReferenceCountUtil.release(lKey);
            if (null != option.endKey()) {
                ReferenceCountUtil.release(option.endKey());
            }
        });
    }

    @Override
    public CompletableFuture<IncrementResult<ByteBuf, ByteBuf>> increment(ByteBuf pKey,
                                                                          ByteBuf lKey,
                                                                          long amount,
                                                                          IncrementOption<ByteBuf> option) {
        pKey.retain();
        lKey.retain();
        return IncrementRequestProcessor.of(
            KvUtils.newIncrementRequest(lKey, amount, option)
                .setHeader(newRoutingHeader(pKey))
                .build(),
            response -> KvUtils.newIncrementResult(response, resultFactory, kvFactory),
            scChannel,
            executor,
            backoffPolicy
        ).process().whenComplete((ignored, cause) -> {
            ReferenceCountUtil.release(pKey);
            ReferenceCountUtil.release(lKey);
        });
    }

    @Override
    public Txn<ByteBuf, ByteBuf> txn(ByteBuf pKey) {
        return new TxnImpl(pKey);
    }

    @Override
    public void close() {
        // no-op
    }

    @Override
    public OpFactory<ByteBuf, ByteBuf> opFactory() {
        return opFactory;
    }

    //
    // Txn Implementation
    //

    class TxnImpl implements Txn<ByteBuf, ByteBuf> {

        private final ByteBuf pKey;
        private final TxnRequest.Builder txnBuilder;
        private final List<AutoCloseable> resourcesToRelease;

        TxnImpl(ByteBuf pKey) {
            this.pKey = pKey.retain();
            this.txnBuilder = TxnRequest.newBuilder();
            this.resourcesToRelease = Lists.newArrayList();
        }

        @SuppressWarnings("unchecked")
        @Override
        public Txn<ByteBuf, ByteBuf> If(CompareOp... cmps) {
            for (CompareOp<ByteBuf, ByteBuf> cmp : cmps) {
                txnBuilder.addCompare(toProtoCompare(cmp));
                resourcesToRelease.add(cmp);
            }
            return this;
        }

        @SuppressWarnings("unchecked")
        @Override
        public Txn<ByteBuf, ByteBuf> Then(Op... ops) {
            for (Op<ByteBuf, ByteBuf> op : ops) {
                txnBuilder.addSuccess(toProtoRequest(op));
                resourcesToRelease.add(op);
            }
            return this;
        }

        @SuppressWarnings("unchecked")
        @Override
        public Txn<ByteBuf, ByteBuf> Else(Op... ops) {
            for (Op<ByteBuf, ByteBuf> op : ops) {
                txnBuilder.addFailure(toProtoRequest(op));
                resourcesToRelease.add(op);
            }
            return this;
        }

        @Override
        public CompletableFuture<TxnResult<ByteBuf, ByteBuf>> commit() {
            return TxnRequestProcessor.of(
                txnBuilder.setHeader(newRoutingHeader(pKey)).build(),
                response -> KvUtils.newKvTxnResult(response, resultFactory, kvFactory),
                scChannel,
                executor,
                backoffPolicy
            ).process().whenComplete((ignored, cause) -> {
                ReferenceCountUtil.release(pKey);
                for (AutoCloseable resource : resourcesToRelease) {
                    closeResource(resource);
                }
            });
        }

        private void closeResource(AutoCloseable resource) {
            try {
                resource.close();
            } catch (Exception e) {
                log.warn("Fail to close resource {}", resource, e);
            }
        }
    }
}
