package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.git.exception.LfsStorageException;
import com.atlassian.bitbucket.mesh.io.IoSupplier;
import com.atlassian.bitbucket.mesh.rpc.util.ByteStringUtils;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamChunk;
import com.google.common.io.ByteStreams;
import io.grpc.stub.ServerCallStreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.InputStream;

import static java.util.Objects.requireNonNull;

public class LfsReadObjectHandler implements Runnable {

    private static final int IO_BUFFER_SIZE = 32768;
    private static final Logger log = LoggerFactory.getLogger(LfsReadObjectHandler.class);

    private final byte[] buffer;
    private final IoSupplier<InputStream> inputSupplier;
    private final ServerCallStreamObserver<RpcStreamChunk> responseObserver;

    private volatile boolean done;
    private volatile InputStream inputStream;

    public LfsReadObjectHandler(@Nonnull IoSupplier<InputStream> inputSupplier,
                                @Nonnull ServerCallStreamObserver<RpcStreamChunk> responseObserver) {
        this.responseObserver = requireNonNull(responseObserver, "responseObserver");
        this.inputSupplier = requireNonNull(inputSupplier, "inputSupplier");
        buffer = new byte[IO_BUFFER_SIZE];
    }

    @Override
    public void run() {
        try {
            if (done) {
                return;
            }
            if (responseObserver.isCancelled()) {
                markDone();
                return;
            }
            if (inputStream == null) {
                inputStream = inputSupplier.get();
            }

            while (!done && responseObserver.isReady() && !responseObserver.isCancelled()) {
                int bytesRead = ByteStreams.read(inputStream, buffer, 0, IO_BUFFER_SIZE);

                RpcStreamChunk.Builder chunkBuilder = RpcStreamChunk.newBuilder()
                        .setClosed(bytesRead < IO_BUFFER_SIZE)
                        .setData(ByteStringUtils.unsafeWrap(buffer, 0, bytesRead));
                responseObserver.onNext(chunkBuilder.build());

                if (chunkBuilder.getClosed()) {
                    markDone();
                    responseObserver.onCompleted();
                }
            }
        } catch (IOException e) {
            log.error("LFS read object handler failure", e);
            markDone();
            responseObserver.onError(new LfsStorageException(e).toStatusException());
        }
    }

    private void markDone() {
        done = true;
        try {
            if (inputStream != null) {
                inputStream.close();
            }
        } catch (IOException e) {
            log.debug("Error during close of InputStream", e);
        }
    }
}
