package com.nike.wingtips.apache.httpclient;

import com.nike.internal.util.StringUtils;
import com.nike.wingtips.Span;
import com.nike.wingtips.Tracer;
import com.nike.wingtips.apache.httpclient.tag.ApacheHttpClientTagAdapter;
import com.nike.wingtips.apache.httpclient.util.WingtipsApacheHttpClientUtil;
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 org.apache.http.HttpRequest;
import org.apache.http.HttpRequestInterceptor;
import org.apache.http.HttpResponse;
import org.apache.http.HttpResponseInterceptor;
import org.apache.http.client.HttpClient;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.protocol.HttpContext;
import org.apache.http.protocol.HttpCoreContext;
import org.apache.http.protocol.HttpProcessor;
import org.jetbrains.annotations.NotNull;

import static com.nike.wingtips.apache.httpclient.util.WingtipsApacheHttpClientUtil.propagateTracingHeaders;

/**
 * (NOTE: {@link WingtipsHttpClientBuilder} is strongly recommended instead of these interceptors if you have control
 * over which {@link HttpClientBuilder} is used to create your {@link HttpClient}s. Reasons for this are described at
 * the bottom of this class javadoc.)
 *
 * <p>This class is an implementation of both {@link HttpRequestInterceptor} and {@link HttpResponseInterceptor} for
 * propagating Wingtips tracing information on a {@link HttpClient} call's request headers, with an option to surround
 * requests 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>In order for tracing to be propagated and any subspan to be completed correctly you need to add this as *both*
 * a request and response interceptor. Forgetting to add this as both a request and response interceptor could leave
 * your tracing in a broken state. Ideally this is added as the first {@link
 * HttpClientBuilder#addInterceptorFirst(HttpRequestInterceptor)} and the last {@link
 * HttpClientBuilder#addInterceptorLast(HttpResponseInterceptor)} so that any subspan surrounds as much of the request
 * as possible, including other interceptors. You can use the {@link #addTracingInterceptors(HttpClientBuilder)}
 * helper method to guarantee this interceptor gets added to both the request and response sides, although there's no
 * way to enforce the ideal request-interceptor-first and response-interceptor-last scenario with a helper method
 * when you have any other interceptors (you would have to add it as a request and response interceptor yourself at
 * the appropriate times).
 *
 * <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>As mentioned at the top of this class' javadocs, {@link WingtipsHttpClientBuilder} is recommended instead of
 * these interceptors if you have control over which {@link HttpClientBuilder} is used to create your {@link
 * HttpClient}s. Reasons for this:
 * <ul>
 *     <li>
 *         There are certain types of exceptions that can occur that prevent the response interceptor side of this
 *         class from executing, thus preventing the subspan around the request from completing. This is a consequence
 *         of how the Apache HttpClient interceptors work. The only way to guarantee subspan completion is to use
 *         {@link WingtipsHttpClientBuilder} instead of this interceptor.
 *     </li>
 *     <li>
 *         There are several ways for interceptors to be accidentally wiped out, e.g. {@link
 *         HttpClientBuilder#setHttpProcessor(HttpProcessor)}.
 *     </li>
 *     <li>
 *         {@link WingtipsHttpClientBuilder} makes sure that any subspan *fully* surrounds the request, including all
 *         other interceptors that are executed.
 *     </li>
 *     <li>
 *         You have to remember to add this interceptor as both a request interceptor ({@link HttpRequestInterceptor})
 *         *and* response interceptor ({@link HttpResponseInterceptor}) on {@link HttpClientBuilder}, or tracing will
 *         be broken.
 *     </li>
 * </ul>
 * That said, these interceptors do work perfectly well as long as they are setup correctly *and* you never experience
 * any of the exceptions that cause the response interceptor to be ignored (this is usually impossible to guarantee,
 * making it a major issue for most use cases - again, please use {@link WingtipsHttpClientBuilder} if you can).
 *
 * @author Nic Munroe
 */
@SuppressWarnings("WeakerAccess")
public class WingtipsApacheHttpClientInterceptor implements HttpRequestInterceptor, HttpResponseInterceptor {

    /**
     * Static default instance of this class. This class is thread-safe so you can reuse this default instance instead
     * of creating new objects.
     */
    public static final WingtipsApacheHttpClientInterceptor DEFAULT_IMPL = new WingtipsApacheHttpClientInterceptor();
    /**
     * This is just {@link #DEFAULT_IMPL} explicitly typed to {@link HttpRequestInterceptor} so that you can call
     * {@link HttpClientBuilder#addInterceptorFirst(HttpRequestInterceptor)} or {@link
     * HttpClientBuilder#addInterceptorLast(HttpRequestInterceptor)} without having to explicitly cast it to
     * {@link HttpRequestInterceptor}.
     */
    public static final HttpRequestInterceptor DEFAULT_REQUEST_IMPL = DEFAULT_IMPL;
    /**
     * This is just {@link #DEFAULT_IMPL} explicitly typed to {@link HttpResponseInterceptor} so that you can call
     * {@link HttpClientBuilder#addInterceptorFirst(HttpResponseInterceptor)} or {@link
     * HttpClientBuilder#addInterceptorLast(HttpResponseInterceptor)} without having to explicitly cast it to
     * {@link HttpResponseInterceptor}.
     */
    public static final HttpResponseInterceptor DEFAULT_RESPONSE_IMPL = DEFAULT_IMPL;

    protected static final String SPAN_TO_CLOSE_HTTP_CONTEXT_ATTR_KEY =
        WingtipsApacheHttpClientInterceptor.class.getSimpleName() + "-span_to_close";

    protected final boolean surroundCallsWithSubspan;
    
    protected final HttpTagAndSpanNamingStrategy<HttpRequest, HttpResponse> tagAndNamingStrategy;
    protected final HttpTagAndSpanNamingAdapter<HttpRequest, HttpResponse> tagAndNamingAdapter;

    /**
     * Creates a new instance with the subspan option turned on and the default {@link HttpTagAndSpanNamingStrategy}
     * and {@link HttpTagAndSpanNamingAdapter} ({@link ZipkinHttpTagStrategy} and {@link ApacheHttpClientTagAdapter}).
     */
    public WingtipsApacheHttpClientInterceptor() {
        this(true);
    }

    /**
     * Creates a new instance with the subspan option set to the value of the {@code surroundCallsWithSubspan}
     * argument, and the default {@link HttpTagAndSpanNamingStrategy} and {@link HttpTagAndSpanNamingAdapter}
     * ({@link ZipkinHttpTagStrategy} and {@link ApacheHttpClientTagAdapter}).
     *
     * @param surroundCallsWithSubspan Pass in true to have requests surrounded in a subspan, false to disable the
     * subspan option.
     */
    public WingtipsApacheHttpClientInterceptor(boolean surroundCallsWithSubspan) {
        this(
            surroundCallsWithSubspan,
            ZipkinHttpTagStrategy.<HttpRequest, HttpResponse>getDefaultInstance(),
            ApacheHttpClientTagAdapter.getDefaultInstance()
        );
    }

    /**
     * Creates a new instance with the subspan option set to the value of the {@code surroundCallsWithSubspan}
     * argument, and the given {@link HttpTagAndSpanNamingStrategy} and {@link HttpTagAndSpanNamingAdapter}.
     *
     * @param surroundCallsWithSubspan Pass in true to have requests surrounded in a subspan, false to disable the
     * subspan option.
     * @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 WingtipsApacheHttpClientInterceptor(
        boolean surroundCallsWithSubspan,
        HttpTagAndSpanNamingStrategy<HttpRequest, HttpResponse> tagAndNamingStrategy,
        HttpTagAndSpanNamingAdapter<HttpRequest, HttpResponse> 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
    public void process(HttpRequest request, HttpContext context) {
        Tracer tracer = Tracer.getInstance();

        if (surroundCallsWithSubspan) {
            // Will start a new trace if necessary, or a subspan if a trace is already in progress.
            Span spanToClose = tracer.startSpanInCurrentContext(
                getSubspanSpanName(request, tagAndNamingStrategy, tagAndNamingAdapter),
                Span.SpanPurpose.CLIENT
            );
            
            tagAndNamingStrategy.handleRequestTagging(spanToClose, request, tagAndNamingAdapter);

            // Add the subspan to the HttpContext so that the response interceptor can retrieve and close it.
            context.setAttribute(SPAN_TO_CLOSE_HTTP_CONTEXT_ATTR_KEY, spanToClose);
        }

        propagateTracingHeaders(request, tracer.getCurrentSpan());
    }

    @Override
    public void process(HttpResponse response, HttpContext context) {
        // See if there's a subspan passed to us from the request interceptor.
        Span spanToClose = (Span) context.getAttribute(SPAN_TO_CLOSE_HTTP_CONTEXT_ATTR_KEY);
        if (spanToClose != null) {
            // There was a subspan. Finalize and close it.
            try {
                // Handle response/error tagging and final span name.
                //      The request should be found in the context attributes - try to extract it from there.
                //      We have no access to any error, so we pass null for the error arg.
                HttpRequest request = null;
                Object requestRawObj = context.getAttribute(HttpCoreContext.HTTP_REQUEST);
                if (requestRawObj instanceof HttpRequest) {
                    request = (HttpRequest) requestRawObj;
                }

                tagAndNamingStrategy.handleResponseTaggingAndFinalSpanName(
                    spanToClose, request, response, null, tagAndNamingAdapter
                );
            } finally {
                // Span.close() contains the logic we want - if the spanToClose was an overall span (new trace)
                //      then tracer.completeRequestSpan() will be called, otherwise it's a subspan and
                //      tracer.completeSubSpan() will be called.
                spanToClose.close();
            }
        }
    }

    /**
     * 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 WingtipsApacheHttpClientUtil#getFallbackSubspanSpanName(HttpRequest)} 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 WingtipsApacheHttpClientUtil.getFallbackSubspanSpanName(request);
    }

    /**
     * Helper method for adding a default instance of this interceptor to the given builder's request *and* response
     * interceptors. The interceptors will have their subspan option turned on.
     *
     * @param builder The builder to add the tracing interceptors to.
     * @param <T> The type of the builder.
     * @return The same builder passed in, but with tracing interceptors added.
     */
    public static <T extends HttpClientBuilder> T addTracingInterceptors(T builder) {
        return addTracingInterceptors(builder, true);
    }

    /**
     * Helper method for adding a default instance of this interceptor to the given builder's request *and* response
     * interceptors. The interceptors will have their subspan option set to the value of the given
     * {@code surroundCallsWithSubspan} argument.
     *
     * @param builder The builder to add the tracing interceptors to.
     * @param surroundCallsWithSubspan Pass in true to have requests surrounded in a subspan, false to disable the
     * subspan option.
     * @param <T> The type of the builder.
     * @return The same builder passed in, but with tracing interceptors added.
     */
    public static <T extends HttpClientBuilder> T addTracingInterceptors(T builder, boolean surroundCallsWithSubspan) {
        WingtipsApacheHttpClientInterceptor interceptor = (surroundCallsWithSubspan)
                                                          ? DEFAULT_IMPL
                                                          : new WingtipsApacheHttpClientInterceptor(false);

        builder.addInterceptorFirst((HttpRequestInterceptor)interceptor);
        builder.addInterceptorLast((HttpResponseInterceptor)interceptor);

        return builder;
    }
}