package com.nike.wingtips.spring.interceptor;

import com.nike.internal.util.StringUtils;
import com.nike.wingtips.Span;
import com.nike.wingtips.Tracer;
import com.nike.wingtips.http.HttpRequestTracingUtils;
import com.nike.wingtips.spring.interceptor.tag.SpringHttpClientTagAdapter;
import com.nike.wingtips.spring.util.HttpRequestWrapperWithModifiableHeaders;
import com.nike.wingtips.spring.util.WingtipsSpringUtil;
import com.nike.wingtips.tags.HttpTagAndSpanNamingAdapter;
import com.nike.wingtips.tags.HttpTagAndSpanNamingStrategy;
import com.nike.wingtips.tags.NoOpHttpTagAdapter;
import com.nike.wingtips.tags.NoOpHttpTagStrategy;
import com.nike.wingtips.tags.ZipkinHttpTagStrategy;
import com.nike.wingtips.util.TracingState;

import org.jetbrains.annotations.NotNull;
import org.springframework.http.HttpMessage;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.AsyncClientHttpRequestExecution;
import org.springframework.http.client.AsyncClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.util.concurrent.ListenableFuture;
import org.springframework.util.concurrent.ListenableFutureCallback;
import org.springframework.web.client.AsyncRestTemplate;

import java.io.IOException;

import static com.nike.wingtips.spring.util.WingtipsSpringUtil.getRequestMethodAsString;
import static com.nike.wingtips.spring.util.WingtipsSpringUtil.propagateTracingHeaders;
import static com.nike.wingtips.util.AsyncWingtipsHelperJava7.runnableWithTracing;
import static com.nike.wingtips.util.AsyncWingtipsHelperJava7.unlinkTracingFromCurrentThread;

/**
 * A {@link AsyncClientHttpRequestInterceptor} which propagates Wingtips tracing information on a downstream {@link
 * AsyncRestTemplate} call's request headers, with an option to surround downstream calls in a subspan. The subspan option
 * defaults to on and is highly recommended since the subspans will provide you with timing info for your downstream
 * calls separate from any parent span that may be active at the time this interceptor executes.
 *
 * <p>If the subspan option is enabled but there's no current span on the current thread when this interceptor executes,
 * then a new root span (new trace) will be created rather than a subspan. In either case the newly created span will
 * have a {@link Span#getSpanPurpose()} of {@link Span.SpanPurpose#CLIENT} since this interceptor is for a client call.
 * The {@link Span#getSpanName()} for the newly created span will be generated by {@link
 * #getSubspanSpanName(HttpRequest, HttpTagAndSpanNamingStrategy, HttpTagAndSpanNamingAdapter)}. Instantiate this
 * class with a custom {@link HttpTagAndSpanNamingStrategy} and/or {@link HttpTagAndSpanNamingAdapter} (preferred),
 * or override that method (last resort) if you want a different span naming format.
 *
 * <p>Note that if you have the subspan option turned off then this interceptor will propagate the {@link
 * Tracer#getCurrentSpan()}'s tracing info downstream if it's available, but will do nothing if no current span exists
 * on the current thread when this interceptor executes as there's no tracing info to propagate. Turning on the
 * subspan option mitigates this as it guarantees there will be a span to propagate.
 *
 * <p>Since this interceptor works by setting request headers and we may be passed an immutable request, we wrap
 * the request in a {@link HttpRequestWrapperWithModifiableHeaders} to guarantee that the request headers are mutable.
 * Keep in mind that this will make the headers mutable for any interceptors that execute after this one.
 *
 * @author Nic Munroe
 */
@SuppressWarnings("WeakerAccess")
public class WingtipsAsyncClientHttpRequestInterceptor implements AsyncClientHttpRequestInterceptor {

    /**
     * The default implementation of this class. Since this class is thread-safe you can reuse this rather than creating
     * a new object.
     */
    public static final WingtipsAsyncClientHttpRequestInterceptor DEFAULT_IMPL =
        new WingtipsAsyncClientHttpRequestInterceptor();

    /**
     * If this is true then all downstream calls that this interceptor intercepts will be surrounded by a
     * subspan which will be started immediately before the call and completed as soon as the call completes.
     */
    protected final boolean surroundCallsWithSubspan;
    /**
     * Controls span naming and tagging when {@link #surroundCallsWithSubspan} is true.
     */
    protected final HttpTagAndSpanNamingStrategy<HttpRequest, ClientHttpResponse> tagAndNamingStrategy;
    /**
     * Used by {@link #tagAndNamingStrategy} for span naming and tagging when {@link #surroundCallsWithSubspan} is true.
     */
    protected final HttpTagAndSpanNamingAdapter<HttpRequest, ClientHttpResponse> tagAndNamingAdapter;
    
    /**
     * Default constructor - sets {@link #surroundCallsWithSubspan} to true, and uses the default
     * {@link HttpTagAndSpanNamingStrategy} and {@link HttpTagAndSpanNamingAdapter} ({@link ZipkinHttpTagStrategy} and
     * {@link SpringHttpClientTagAdapter}).
     */
    public WingtipsAsyncClientHttpRequestInterceptor() {
        this(true);
    }

    /**
     * Constructor that lets you choose whether downstream calls will be surrounded with a subspan. The default
     * {@link HttpTagAndSpanNamingStrategy} and {@link HttpTagAndSpanNamingAdapter} will be used
     * ({@link ZipkinHttpTagStrategy} and {@link SpringHttpClientTagAdapter}).
     *
     * @param surroundCallsWithSubspan pass in true to have downstream calls surrounded with a new span, false to only
     * propagate the current span's info downstream (no subspan).
     */
    public WingtipsAsyncClientHttpRequestInterceptor(boolean surroundCallsWithSubspan) {
        this(
            surroundCallsWithSubspan,
            ZipkinHttpTagStrategy.<HttpRequest, ClientHttpResponse>getDefaultInstance(),
            SpringHttpClientTagAdapter.getDefaultInstance()
        );
    }
    
    /**
     * Constuctor that lets you define whether downstream calls will be surrounded with a subspan and provide
     * a different span tag strategy. 
     * @param surroundCallsWithSubspan pass in true to have downstream calls surrounded with a new span, false to only
     * @param tagAndNamingStrategy The span tag and naming strategy to use - cannot be null. If you really want no
     * tag and naming strategy, then pass in {@link NoOpHttpTagStrategy#getDefaultInstance()}.
     * @param tagAndNamingAdapter The tag and naming adapter to use - cannot be null. If you really want no tag and
     * naming adapter, then pass in {@link NoOpHttpTagAdapter#getDefaultInstance()}.
     */
    public WingtipsAsyncClientHttpRequestInterceptor(
        boolean surroundCallsWithSubspan,
        HttpTagAndSpanNamingStrategy<HttpRequest, ClientHttpResponse> tagAndNamingStrategy,
        HttpTagAndSpanNamingAdapter<HttpRequest, ClientHttpResponse> tagAndNamingAdapter
    ) {
        if (tagAndNamingStrategy == null) {
            throw new IllegalArgumentException(
                "tagAndNamingStrategy cannot be null - if you really want no strategy, use NoOpHttpTagStrategy"
            );
        }

        if (tagAndNamingAdapter == null) {
            throw new IllegalArgumentException(
                "tagAndNamingAdapter cannot be null - if you really want no adapter, use NoOpHttpTagAdapter"
            );
        }

        this.surroundCallsWithSubspan = surroundCallsWithSubspan;
        this.tagAndNamingStrategy = tagAndNamingStrategy;
        this.tagAndNamingAdapter = tagAndNamingAdapter;
    }

    @Override
    @SuppressWarnings("deprecation")
    public ListenableFuture<ClientHttpResponse> intercept(
        HttpRequest request, byte[] body, AsyncClientHttpRequestExecution execution
    ) throws IOException {
        // We need to wrap the request with HttpRequestWrapperWithModifiableHeaders so that tracing info can be
        //      propagated on the headers.
        HttpRequestWrapperWithModifiableHeaders wrapperRequest = new HttpRequestWrapperWithModifiableHeaders(request);

        if (surroundCallsWithSubspan) {
            return createAsyncSubSpanAndExecute(wrapperRequest, body, execution);
        }
        
        return propagateTracingHeadersAndExecute(wrapperRequest, body, execution);
    }

    /**
     * Calls {@link WingtipsSpringUtil#propagateTracingHeaders(HttpMessage, Span)} to propagate the current span's
     * tracing state on the given request's headers, then returns
     * {@link AsyncClientHttpRequestExecution#executeAsync(HttpRequest, byte[])} to execute the request.
     *
     * @return The result of calling {@link AsyncClientHttpRequestExecution#executeAsync(HttpRequest, byte[])}.
     */
    protected ListenableFuture<ClientHttpResponse> propagateTracingHeadersAndExecute(
        HttpRequestWrapperWithModifiableHeaders wrapperRequest, byte[] body, AsyncClientHttpRequestExecution execution
    ) throws IOException {
        propagateTracingHeaders(wrapperRequest, Tracer.getInstance().getCurrentSpan());

        // Execute the request/interceptor chain.
        return execution.executeAsync(wrapperRequest, body);
    }

    /**
     * Creates a subspan (or new trace if no current span exists) to surround the HTTP request, then returns the
     * result of calling {@link #propagateTracingHeadersAndExecute(HttpRequestWrapperWithModifiableHeaders, byte[],
     * AsyncClientHttpRequestExecution)} to actually execute the request. A {@link SpanAroundAsyncCallFinisher} will
     * be registered as a callback to finish the subspan when the request finishes. Request tagging (and initial span
     * naming) is done here, and response tagging (and final span naming) is done in the {@link
     * SpanAroundAsyncCallFinisher}.
     *
     * @return The result of calling {@link #propagateTracingHeadersAndExecute(HttpRequestWrapperWithModifiableHeaders,
     * byte[], AsyncClientHttpRequestExecution)} after surrounding the request with a subspan (or new trace if no
     * current span exists).
     */
    protected ListenableFuture<ClientHttpResponse> createAsyncSubSpanAndExecute(
        HttpRequestWrapperWithModifiableHeaders wrapperRequest, byte[] body, AsyncClientHttpRequestExecution execution
    ) throws IOException {
        // Handle subspan stuff. Start by getting the current thread's tracing state (so we can restore it before
        //      this method returns).
        TracingState originalThreadInfo = TracingState.getCurrentThreadTracingState();

        SpanAroundAsyncCallFinisher subspanFinisher = null;

        try {
            // This will start a new trace if necessary, or a subspan if a trace is already in progress.
            Span subspan = Tracer.getInstance().startSpanInCurrentContext(
                getSubspanSpanName(wrapperRequest, tagAndNamingStrategy, tagAndNamingAdapter),
                Span.SpanPurpose.CLIENT
            );

            // Add request tags to the subspan.
            tagAndNamingStrategy.handleRequestTagging(subspan, wrapperRequest, tagAndNamingAdapter);

            // Create the callback that will complete the subspan when the request finishes.
            subspanFinisher = new SpanAroundAsyncCallFinisher(
                TracingState.getCurrentThreadTracingState(), wrapperRequest, tagAndNamingStrategy, tagAndNamingAdapter
            );

            // Execute the request/interceptor chain, and add the callback to finish the subspan (if one exists).
            ListenableFuture<ClientHttpResponse> result = propagateTracingHeadersAndExecute(
                wrapperRequest, body, execution
            );
            result.addCallback(subspanFinisher);

            return result;
        }
        catch(Throwable t) {
            // Something went wrong, probably in the execution.executeAsync(...) call. Complete the subspan now
            //      (if one exists).
            if (subspanFinisher != null) {
                subspanFinisher.finishCallSpan(null, t);
            }

            throw t;
        }
        finally {
            // Reset back to the original tracing state that was on this thread when this method began.
            //noinspection deprecation
            unlinkTracingFromCurrentThread(originalThreadInfo);
        }
    }

    /**
     * Returns the name that should be used for the subspan surrounding the call. Defaults to whatever {@link
     * HttpTagAndSpanNamingStrategy#getInitialSpanName(Object, HttpTagAndSpanNamingAdapter)} returns, with a fallback
     * of {@link HttpRequestTracingUtils#getFallbackSpanNameForHttpRequest(String, String)} if the naming strategy
     * returned null or blank string. You can override this method to return something else if you want different
     * behavior and you don't want to adjust the naming strategy or adapter.
     *
     * @param request The request that is about to be executed.
     * @param namingStrategy The {@link HttpTagAndSpanNamingStrategy} being used.
     * @param adapter The {@link HttpTagAndSpanNamingAdapter} being used.
     * @return The name that should be used for the subspan surrounding the call.
     */
    protected @NotNull String getSubspanSpanName(
        @NotNull HttpRequest request,
        @NotNull HttpTagAndSpanNamingStrategy<HttpRequest, ?> namingStrategy,
        @NotNull HttpTagAndSpanNamingAdapter<HttpRequest, ?> adapter
    ) {
        // Try the naming strategy first.
        String subspanNameFromStrategy = namingStrategy.getInitialSpanName(request, adapter);

        if (StringUtils.isNotBlank(subspanNameFromStrategy)) {
            return subspanNameFromStrategy;
        }

        // The naming strategy didn't have anything for us. Fall back to something reasonable.
        return HttpRequestTracingUtils.getFallbackSpanNameForHttpRequest(
            "asyncresttemplate_downstream_call", getRequestMethodAsString(request.getMethod())
        );
    }

    /**
     * A {@link ListenableFutureCallback} that will complete the given {@link TracingState} (e.g. tracing state
     * representing a subspan) when executed. This should be attached as a callback to the result of {@link
     * #intercept(HttpRequest, byte[], AsyncClientHttpRequestExecution)}.
     */
    @SuppressWarnings("WeakerAccess")
    protected static class SpanAroundAsyncCallFinisher implements ListenableFutureCallback<ClientHttpResponse> {

        protected final TracingState spanAroundCallTracingState;
        protected final HttpRequest request;
        protected final HttpTagAndSpanNamingStrategy<HttpRequest, ClientHttpResponse> tagAndNamingStrategy;
        protected final HttpTagAndSpanNamingAdapter<HttpRequest, ClientHttpResponse> tagAndNamingAdapter;

        protected SpanAroundAsyncCallFinisher(
            TracingState spanAroundCallTracingState,
            HttpRequest request,
            HttpTagAndSpanNamingStrategy<HttpRequest, ClientHttpResponse> tagAndNamingStrategy,
            HttpTagAndSpanNamingAdapter<HttpRequest, ClientHttpResponse> tagAndNamingAdapter
        ) {
            this.spanAroundCallTracingState = spanAroundCallTracingState;
            this.request = request;
            this.tagAndNamingStrategy = tagAndNamingStrategy;
            this.tagAndNamingAdapter = tagAndNamingAdapter;
        }

        @Override
        public void onFailure(Throwable ex) {
            finishCallSpan(null, ex);
        }

        @Override
        public void onSuccess(ClientHttpResponse result) {
            finishCallSpan(result, null);
        }

        @SuppressWarnings("deprecation")
        protected void finishCallSpan(final ClientHttpResponse response, final Throwable error) {
            if (spanAroundCallTracingState != null) {
                runnableWithTracing(
                    new Runnable() {
                        @Override
                        public void run() {
                            Span span = Tracer.getInstance().getCurrentSpan();
                            //noinspection TryFinallyCanBeTryWithResources
                            try {
                                // Add the tags from the response.
                                tagAndNamingStrategy.handleResponseTaggingAndFinalSpanName(
                                    span, request, response, error, tagAndNamingAdapter
                                );
                            }
                            finally {
                                // Span.close() contains the logic we want - if the spanAroundCall was an overall span
                                //      (new trace) then tracer.completeRequestSpan() will be called, otherwise it's
                                //      a subspan and tracer.completeSubSpan() will be called.
                                span.close();
                            }
                        }
                    },
                    spanAroundCallTracingState
                ).run();
            }
        }
    }
}