package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.config.GrpcConfig;
import com.atlassian.bitbucket.mesh.rpc.ZeroCopyMarshaller;
import com.atlassian.bitbucket.mesh.rpc.ZeroCopyServerCallListener;
import com.atlassian.bitbucket.mesh.rpc.ZeroCopyUtils;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcMarshallerType;
import io.grpc.*;
import io.grpc.MethodDescriptor.PrototypeMarshaller;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.Field;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * Interceptor that inspects the RPC method and service options, and for RPC methods / services configured to use
 * the {@link RpcMarshallerType#MARSHALLER_ZERO_COPY zero-copy marshaller} overrides the standard request marshaller
 * with the zero-copy marshaller.
 */
public class ZeroCopyServerInterceptor implements ServerInterceptor {

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

    private final boolean available;
    private final GrpcConfig config;
    private final Field marshallerField;
    private final ConcurrentMap<String, Boolean> zeroCopyDescriptors;

    public ZeroCopyServerInterceptor(GrpcConfig config) {
        this.config = config;

        Field field = null;
        try {
            field = MethodDescriptor.class.getDeclaredField("requestMarshaller");
            ReflectionUtils.makeAccessible(field);
        } catch (NoSuchFieldException | SecurityException e) {
            log.warn("Cannot access MethodDescriptor.responseMarshaller field. Zero-copy marshalling is disabled.");
        }
        available = field != null;
        marshallerField = field;
        zeroCopyDescriptors = new ConcurrentHashMap<>();
    }

    @Override
    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers,
                                                                 ServerCallHandler<ReqT, RespT> next) {
        if (available && config.isZeroCopy()) {
            MethodDescriptor<ReqT, RespT> method = call.getMethodDescriptor();
            if (zeroCopyDescriptors.computeIfAbsent(method.getFullMethodName(), fqName -> {
                if (ZeroCopyUtils.getMarshallerType(method) == RpcMarshallerType.MARSHALLER_ZERO_COPY) {
                    //Try to override the response marshaller in the MethodDescriptor with the zero-copy marshaller
                    try {
                        return injectZeroCopyMarshaller(method);
                    } catch (IllegalAccessException e) {
                        log.warn("Failed to update {} responseMarshaller with the zero-copy marshaller.", fqName, e);
                        return Boolean.FALSE;
                    }
                }
                //Method should not use the zero-copy marshaller
                return Boolean.FALSE;
            })) {
                //Zero copy marshaller will be used; make sure the detached streams are closed by
                //wrapping the ServerCall.Listener in a ZeroCopyServerCallListener
                return new ZeroCopyServerCallListener<>(next.startCall(call, headers));
            }
        }
        return next.startCall(call, headers);
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    private boolean injectZeroCopyMarshaller(MethodDescriptor descriptor) throws IllegalAccessException {
        MethodDescriptor.Marshaller<?> marshaller = descriptor.getRequestMarshaller();
        if (marshaller instanceof PrototypeMarshaller) {
            marshallerField.set(
                    descriptor, new ZeroCopyMarshaller(config::isZeroCopy, (PrototypeMarshaller) marshaller));
            return true;
        }
        if (marshaller instanceof ZeroCopyMarshaller) {
            log.info("{} is already configured with a zero-copy marshaller", descriptor.getFullMethodName());
            return true;
        }

        log.warn("Cannot configure a zero-copy marshaller for {}", descriptor.getFullMethodName());
        return false;
    }
}
