package com.atlassian.bitbucket.mesh.execution;

import com.atlassian.bitbucket.mesh.grpc.GrpcServiceAdvice;
import com.google.common.base.Throwables;
import io.grpc.stub.StreamObserver;

import javax.annotation.Nonnull;
import java.time.Duration;
import java.util.concurrent.*;

import static java.util.Objects.requireNonNull;

public class DefaultGrpcAsyncExecutor implements GrpcAsyncExecutor {

    private final ExecutionManager executionManager;
    private final ThreadPoolExecutor executor;

    public DefaultGrpcAsyncExecutor(ExecutionManager executionManager, int coreThreads, int maxThreads,
                                    Duration keepAliveTime, ThreadFactory threadFactory) {
        this.executionManager = executionManager;

        executor = new ThreadPoolExecutor(coreThreads, maxThreads, keepAliveTime.getSeconds(), TimeUnit.SECONDS,
                new SynchronousQueue<>(), threadFactory);
    }

    @Override
    public <T> void execute(@Nonnull StreamObserver<T> streamObserver, @Nonnull Runnable task) {
        requireNonNull(streamObserver, "streamObserver");
        requireNonNull(task, "task");

        executor.execute(new GrpcTask(streamObserver, executionManager.maybeWrap(task)));
    }

    @Nonnull
    @Override
    public <T> Future<Void> submit(@Nonnull StreamObserver<T> streamObserver, @Nonnull Runnable task) {
        requireNonNull(streamObserver, "streamObserver");
        requireNonNull(task, "task");

        return executor.submit(new GrpcTask(streamObserver, executionManager.maybeWrap(task)), null);
    }

    @Override
    public void shutdown() {
        executor.shutdown();
    }

    private static class GrpcTask implements Runnable {

        private final Runnable task;
        private final StreamObserver<?> observer;

        GrpcTask(StreamObserver<?> observer, Runnable task) {
            this.task = task;
            this.observer = observer;
        }

        @Override
        public void run() {
            try {
                task.run();
            } catch (Throwable t) {
                try {
                    GrpcServiceAdvice.forwardThrowable(observer, t);
                } catch (Throwable throwable) {
                    //This should never throw because GrpcServiceAdvice.forwardThrowable only throws if observer is
                    //null, which is never the case here. But we'll handle it anyway
                    Throwables.throwIfUnchecked(t);
                    throw new RuntimeException(t);
                }
            }
        }
    }
}
