package com.atlassian.bitbucket.mesh.git;

import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.git.builder.GitProcessBuilderFactory;
import com.atlassian.bitbucket.mesh.git.exception.*;
import com.atlassian.bitbucket.mesh.repository.Repository;
import com.atlassian.bitbucket.mesh.repository.RepositoryManager;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.google.protobuf.ByteString;
import com.google.protobuf.Timestamp;
import io.grpc.Status;
import io.grpc.stub.CallStreamObserver;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.WritableByteChannel;
import java.time.Instant;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

public class GrpcGitDataService extends GitDataServiceGrpc.GitDataServiceImplBase {

    private final GitAgent agent;
    private final GitProcessBuilderFactory builderFactory;
    private final EventPublisher eventPublisher;
    private final ExecutionManager executionManager;
    private final RepositoryManager repositoryManager;
    private final UploadManager uploadManager;

    public GrpcGitDataService(GitAgent agent, GitProcessBuilderFactory builderFactory,
                              EventPublisher eventPublisher, ExecutionManager executionManager,
                              RepositoryManager repositoryManager, UploadManager uploadManager) {
        this.agent = agent;
        this.builderFactory = builderFactory;
        this.eventPublisher = eventPublisher;
        this.executionManager = executionManager;
        this.repositoryManager = repositoryManager;
        this.uploadManager = uploadManager;
    }

    @Nonnull
    @Override
    public StreamObserver<RpcUploadFileFragment> uploadFile(StreamObserver<RpcUploadFileResponse> responseObserver) {
        return new UploadFileRequestObserver(repositoryManager, uploadManager, responseObserver);
    }

    @Nonnull
    @Override
    public StreamObserver<RpcUploadFilesFragment> uploadFiles(StreamObserver<RpcUploadFilesResponse> responseObserver) {
        return new UploadFilesRequestObserver(repositoryManager, uploadManager, responseObserver);
    }

    @Nonnull
    @Override
    public StreamObserver<RpcUploadPackFileFragment> uploadPackFile(
            StreamObserver<RpcUploadPackFileResponse> responseObserver) {
        return new UploadPackFileRequestObserver(agent, builderFactory, eventPublisher, executionManager,
                repositoryManager, responseObserver);
    }

    private abstract static class AbstractUploadRequestObserver<ReqT, RespT> implements StreamObserver<ReqT> {

        protected final AtomicBoolean completed;
        protected final Logger log;
        protected final RepositoryManager repositoryManager;
        protected final UploadManager uploadManager;

        protected StreamObserver<RespT> responseObserver;

        private volatile WritableByteChannel outputChannel;
        private volatile UploadFileParameters uploading;

        AbstractUploadRequestObserver(RepositoryManager repositoryManager, UploadManager uploadManager,
                                      StreamObserver<RespT> responseObserver) {
            this.repositoryManager = repositoryManager;
            this.responseObserver = responseObserver;
            this.uploadManager = uploadManager;

            //Request 2 to configure flow control to allow 2 incoming messages "on the wire" before blocking waiting
            //for space.
            ((CallStreamObserver<RespT>) responseObserver).request(2);

            completed = new AtomicBoolean(false);
            log = LoggerFactory.getLogger(getClass());
        }

        @Override
        public void onCompleted() {
            if (uploading != null) {
                throw new FailedPreconditionException("Early EOF detected for upload of " + uploading.getPath() +
                        " to " + uploading.getRepository() + ".");
            }

            maybeComplete();
        }

        @Override
        public void onError(Throwable t) {
            Status status = Status.fromThrowable(t);
            if (status.getCode() == Status.CANCELLED.getCode()) {
                log.debug("[{}] Client canceled {} upload to {} ", uploading.getRepository(),
                        uploading.getType(), uploading.getPath());
            } else {
                //Client encountered an error
                log.warn("[{}] Client encountered an error while processing {} upload to {}",
                        uploading.getRepository(), uploading.getType(), uploading.getPath(), t);
            }
            maybeComplete();
        }

        protected static UploadFileParameters buildParameters(Repository repository, UploadType uploadType,
                                                              RpcOsFile file) {
            return new UploadFileParameters.Builder(repository, uploadType, file.getPath().toStringUtf8())
                    .lastModifiedTime(toInstant(file.getLastModified()))
                    .size(file.getSize())
                    .build();
        }

        protected void maybeComplete() {
            if (completed.compareAndSet(false, true)) {
                responseObserver.onCompleted();
            }
        }

        protected void onChunk(RpcStreamChunk chunk) {
            if (outputChannel == null) {
                throw new FailedPreconditionException("Send an upload request before sending the data!");
            }

            try {
                ByteString data = chunk.getData();
                for (ByteBuffer buffer : data.asReadOnlyByteBufferList()) {
                    outputChannel.write(buffer);
                }
                if (log.isTraceEnabled()) {
                    log.trace("[{}] Wrote {} bytes to {} upload to {}",
                            uploading.getRepository(), chunk.getData().size(), uploading.getType(), uploading.getPath());
                }
                if (chunk.getClosed()) {
                    if (log.isTraceEnabled()) {
                        log.trace("[{}] Completed {} upload to {}",
                                uploading.getRepository(), uploading.getType(), uploading.getPath());
                    }
                    outputChannel.close();
                    outputChannel = null;
                    uploading = null;
                }
            } catch (IOException e) {
                log.warn("Error while processing chunk", e);
                responseObserver.onError(Status.INTERNAL
                        .withDescription("Failed to write " + uploading.getType() +
                                " chunk to disk (path = " + uploading.getPath() + ")")
                        .asRuntimeException());
            }
        }

        protected void startUpload(UploadFileParameters parameters) {
            if (uploading != null) {
                throw new FailedPreconditionException("Upload of " + uploading.getPath() + " to " +
                        uploading.getRepository() + " in progress while receiving request to upload " +
                        parameters.getPath() + ".");
            }

            try {
                outputChannel = uploadManager.uploadFile(parameters);
                uploading = parameters;
            } catch (IOException e) {
                responseObserver.onError(Status.INTERNAL
                        .withDescription("Failed to open file for writing (type = " + parameters.getType() +
                                ", path = " + parameters.getPath() + ")")
                        .asRuntimeException());
            }
        }

        private static Instant toInstant(Timestamp timestamp) {
            return Instant.ofEpochSecond(timestamp.getSeconds(), timestamp.getNanos());
        }
    }

    private static class UploadFileRequestObserver
            extends AbstractUploadRequestObserver<RpcUploadFileFragment, RpcUploadFileResponse> {

        UploadFileRequestObserver(RepositoryManager repositoryManager, UploadManager uploadManager,
                                  StreamObserver<RpcUploadFileResponse> responseObserver) {
            super(repositoryManager, uploadManager, responseObserver);
        }

        @Override
        public void onNext(RpcUploadFileFragment value) {
            switch (value.getRequestOneofCase()) {
                case CHUNK:
                    onChunk(value.getChunk());
                    break;
                case REQUEST:
                    onRequest(value.getRequest());
                    break;
                default:
                    completed.set(true);
                    throw new InvalidArgumentException("Received [" + value.getRequestOneofCase() +
                            "] where [case CHUNK or REQUEST] was expected");
            }
        }

        private void onRequest(RpcUploadFileRequest request) {
            Repository repository = repositoryManager.getById(request.getRepository());
            UploadFileParameters parameters = buildParameters(
                    repository, UploadType.fromRpc(request.getType()), request.getFile());

            boolean want = uploadManager.want(parameters);
            if (log.isTraceEnabled()) {
                log.trace("[{}] Server {} {} upload of {}",
                        repository, want ? "wants" : "does not want", parameters.getType(), parameters.getPath());
            }
            //Let the client know whether we're interested or not
            responseObserver.onNext(RpcUploadFileResponse.newBuilder()
                    .setWant(want)
                    .build());

            //And set up everything for the file in question
            if (want) {
                startUpload(parameters);
            }
        }
    }

    private static class UploadFilesRequestObserver
            extends AbstractUploadRequestObserver<RpcUploadFilesFragment, RpcUploadFilesResponse> {

        private final Map<String, UploadFileParameters> requestedFilesByPath;

        UploadFilesRequestObserver(RepositoryManager repositoryManager, UploadManager uploadManager,
                                   StreamObserver<RpcUploadFilesResponse> responseObserver) {
            super(repositoryManager, uploadManager, responseObserver);

            requestedFilesByPath = new HashMap<>();
        }

        @Override
        public void onNext(RpcUploadFilesFragment value) {
            switch (value.getRequestOneofCase()) {
                case CHUNK:
                    onChunk(value.getChunk());
                    break;
                case FILE_START:
                    onFileStart(value.getFileStart());
                    break;
                case REQUEST:
                    onRequest(value.getRequest());
                    break;
                default:
                    completed.set(true);
                    throw new InvalidArgumentException("Received [" + value.getRequestOneofCase() +
                            "] where [case CHUNK, FILE_START or REQUEST] was expected");
            }
        }

        private void onFileStart(RpcUploadFileStart fileStart) {
            String path = fileStart.getPath().toStringUtf8();
            UploadFileParameters parameters = requestedFilesByPath.remove(path);
            if (parameters == null) {
                throw new InvalidArgumentException("File " + path + " is offered for upload but was not requested");
            }

            startUpload(parameters);
        }

        private void onRequest(RpcUploadFilesRequest request) {
            Repository repository = repositoryManager.getById(request.getRepository());
            UploadType uploadType = UploadType.fromRpc(request.getType());

            requestedFilesByPath.clear();
            RpcUploadFilesResponse.Builder builder = RpcUploadFilesResponse.newBuilder();
            for (RpcOsFile file : request.getFilesList()) {
                UploadFileParameters parameters = buildParameters(repository, uploadType, file);
                boolean want = uploadManager.want(parameters);
                if (log.isTraceEnabled()) {
                    log.trace("[{}] Server {} {} upload of {}",
                            repository, want ? "wants" : "does not want", parameters.getType(), parameters.getPath());
                }

                if (want) {
                    //Inform the client that we're interested
                    builder.addWantPaths(file.getPath());
                    //Record the file attributes for later when we receive the file
                    requestedFilesByPath.put(parameters.getPath(), parameters);
                }
            }

            responseObserver.onNext(builder.build());
        }
    }
}
