package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.network.MeshNode;
import io.grpc.*;
import org.springframework.beans.factory.config.BeanPostProcessor;

import javax.annotation.Nonnull;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArraySet;

import static com.atlassian.bitbucket.mesh.network.MeshChannel.OPT_TARGET_NODE;
import static com.atlassian.bitbucket.mesh.rpc.MetadataKeys.KEY_TARGET_ID;

/**
 * A {@link ClientInterceptor} which applies {@link MetadataProvider}s to populate the headers that are dispatched
 * when gRPC requests are started.
 */
public class MetadataClientInterceptor implements ClientInterceptor, BeanPostProcessor {

    private final Set<MetadataProvider> providers;

    public MetadataClientInterceptor() {
        this.providers = new CopyOnWriteArraySet<>();
    }

    @Override
    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method,
                                                               CallOptions callOptions, Channel next) {
        //Run the chain to create the delegate. If this throws, no reason to populate metadata
        ClientCall<ReqT, RespT> delegate = next.newCall(method, callOptions);
        if (providers.isEmpty()) {
            return delegate;
        }

        //Populate the metadata to send. This is done here, rather than inside start(Listener, Metadata)
        //below, because this is called on the request thread where that is called on a gRPC thread. If
        //we call it in start some request state may not be available
        Metadata metadata = new Metadata();
        for (MetadataProvider provider : providers) {
            provider.apply(metadata);
        }
        //Include the intended node's ID to detect cases to detect configuration issues / incorrect name
        //resolution due to stale DNS data.
        MeshNode targetNode = callOptions.getOption(OPT_TARGET_NODE);
        if (targetNode != null) {
            metadata.put(KEY_TARGET_ID, targetNode.getId());
        }

        return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(delegate) {

            @Override
            public void start(Listener<RespT> responseListener, Metadata headers) {
                headers.merge(metadata);

                super.start(responseListener, headers);
            }
        };
    }

    /**
     * We can't use constructor or setter injection for several of the {@link MetadataProvider} implementations
     * because they end up triggering circular dependencies. Because we'd inject the providers using an array,
     * the circular dependencies don't even trigger the {@code ApplicationContext} to fail to start; instead,
     * any {@link MetadataProvider} instances that haven't yet been created are <i>silently omitted</i> from the
     * array and don't get invoked.
     * <p>
     * Rather than fighting dependencies to allow injecting providers, we use {@link BeanPostProcessor} to scan
     * beans as they're initialized and register any that implement {@link MetadataProvider}.
     *
     * @param bean     the bean instance to analyze
     * @param beanName the bean's name in the {@code ApplicationContext}
     * @return the provided {@code bean}, unchanged
     */
    @Override
    public Object postProcessAfterInitialization(@Nonnull Object bean, @Nonnull String beanName) {
        if (bean instanceof MetadataProvider) {
            providers.add((MetadataProvider) bean);
        }

        return bean;
    }
}
