package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.AbstractStatusException;
import com.google.common.base.Throwables;
import io.grpc.stub.StreamObserver;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.annotation.Order;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

/**
 * Advice for decorating {@link io.grpc.BindableService gRPC services}.
 */
@Aspect
@Order(-1) //Wired as the innermost advice for gRPC services
public class GrpcServiceAdvice {

    private static final Logger log = LoggerFactory.getLogger(GrpcServiceAdvice.class);

    /**
     * Inspects the provided {@link Throwable} and propagates the exception to the provided {@link StreamObserver},
     * or rethrows the exception is {@code observer} is {@code null}.
     *
     * @param observer the observer
     * @param t the throwable
     * @throws Throwable the provided {@code t} if {@code observer} is {@code null}
     */
    public static void forwardThrowable(@Nullable StreamObserver<?> observer, @Nonnull Throwable t) throws Throwable {
        if (t instanceof AbstractStatusException) {
            //If the exception is an AbstractStatusException, translate it to a StatusRuntimeException to allow
            //clients to handle it better
            t = ((AbstractStatusException) t).toStatusException();
        } else if (isClosed(t)) {
            //Logging the exception here will show the _second_ place the RPC was closed from. The first place was
            //most likely a ServerInterceptor or AOP interceptor; we may be able to infer it from the second closer
            log.debug("The RPC was closed twice", t);
            return;
        }

        if (observer == null) {
            //If we couldn't find a StreamObserver, propagate the exception (which may have been translated, if it
            //was an AbstractStatusException)
            throw t;
        }

        //Otherwise, suppress the exception and forward it to the observer instead. If the exception isn't a
        //StatusException or StatusRuntimeException, the observer will convert it into a generic one automatically
        try {
            observer.onError(t);
        } catch (IllegalStateException e) {
            if (isClosed(e)) {
                //This means _we_ were the second caller to try and close the RPC. Since we can't forward the
                //exception, log it locally instead
                log.debug("The exception could not be forwarded; the RPC is already closed", t);
            } else {
                throw e;
            }
        }
    }

    /**
     * Catches exceptions thrown by gRPC service methods and attempts to forward them to a {@link StreamObserver}.
     * <p>
     * gRPC expects services to dispatch errors using {@link StreamObserver#onError} rather than throwing exceptions.
     * Actually doing that in practice can be awkward, though, requiring the {@link StreamObserver} to be passed into
     * all sorts of types. To avoid that, this advisor catches thrown exceptions and forwards them to the observer.
     * It also {@link AbstractStatusException#toStatusException translates} {@link AbstractStatusException}s to
     * {@code StatusRuntimeException}s, to simplify providing additional context that may be useful for clients.
     *
     * @param joinPoint describes the invocation being intercepted
     * @return the returned value from the intercepted method
     * @throws Throwable if the intercepted method throws and doesn't have a {@link StreamObserver} the advisor can
     *                   forward exceptions to
     */
    @Around("target(io.grpc.BindableService)") //Apply to all types which implement BindableService
    public Object forwardExceptions(ProceedingJoinPoint joinPoint) throws Throwable {
        //Try to find a StreamObserver so we can forward the exception to it instead of
        //just propagating the exception.
        StreamObserver<?> observer = null;
        for (Object arg : joinPoint.getArgs()) {
            if (arg instanceof StreamObserver) {
                observer = (StreamObserver<?>) arg;
                break;
            }
        }

        try {
            return maybeWrap(joinPoint.proceed(), observer);
        } catch (Throwable t) {
            forwardThrowable(observer, t);
            return null;
        }
    }

    private static boolean isClosed(Throwable t) {
        if (t instanceof IllegalStateException) {
            //gRPC doesn't offer any API for detecting whether a call is already closed. The only
            //way we can detect it is string matching
            String message = t.getMessage();

            return "call already closed".equals(message) || "call is closed".equals(message);
        }

        return false;
    }

    @SuppressWarnings({"rawtypes", "unchecked"})
    private static Object maybeWrap(Object returnValue, StreamObserver<?> observer) {
        if (returnValue instanceof StreamObserver) {
            return new ErrorTranslatingStreamObserver((StreamObserver) returnValue, observer);
        }
        return returnValue;
    }

    private static class ErrorTranslatingStreamObserver<ReqT, RespT> implements StreamObserver<RespT> {

        private final StreamObserver<RespT> delegate;
        private final StreamObserver<ReqT> requestObserver;

        ErrorTranslatingStreamObserver(StreamObserver<RespT> delegate, StreamObserver<ReqT> requestObserver) {
            this.delegate = delegate;
            this.requestObserver = requestObserver;
        }

        @Override
        public void onNext(RespT value) {
            try {
                delegate.onNext(value);
            } catch (RuntimeException e) {
                handleException(e);
            }
        }

        @Override
        public void onError(Throwable t) {
            delegate.onError(t);
        }

        @Override
        public void onCompleted() {
            try {
                delegate.onCompleted();
            } catch (RuntimeException e) {
                handleException(e);
            }
        }

        void handleException(RuntimeException e) {
            try {
                forwardThrowable(requestObserver, e);
            } catch (Throwable t) {
                Throwables.throwIfUnchecked(t);
                throw new RuntimeException(e); //should never happen
            }
        }
    }
}