package com.nike.wingtips.spring.interceptor;

import com.nike.internal.util.StringUtils;
import com.nike.wingtips.Span;
import com.nike.wingtips.Span.SpanPurpose;
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 org.jetbrains.annotations.NotNull;
import org.springframework.http.HttpMessage;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.web.client.RestTemplate;

import java.io.IOException;

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

/**
 * A {@link ClientHttpRequestInterceptor} which propagates Wingtips tracing information on a downstream {@link
 * RestTemplate} 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 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.
 */
@SuppressWarnings("WeakerAccess")
public class WingtipsClientHttpRequestInterceptor implements ClientHttpRequestInterceptor {

    /**
     * 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 WingtipsClientHttpRequestInterceptor DEFAULT_IMPL = new WingtipsClientHttpRequestInterceptor();

    /**
     * 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 WingtipsClientHttpRequestInterceptor() {
        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 WingtipsClientHttpRequestInterceptor(boolean surroundCallsWithSubspan) {
        this(
            surroundCallsWithSubspan,
            ZipkinHttpTagStrategy.<HttpRequest, ClientHttpResponse>getDefaultInstance(),
            SpringHttpClientTagAdapter.getDefaultInstance()
        );
    }

    /**
     * Constructor that lets you choose whether downstream calls will be surrounded with a subspan and supply the relevant tag strategy
     * for the subspan.
     * 
     * @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).
     * @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 WingtipsClientHttpRequestInterceptor(
        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
    public ClientHttpResponse intercept(
        HttpRequest request, byte[] body, ClientHttpRequestExecution 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 createNewSpanAndExecuteRequest(wrapperRequest, body, execution);
        }

        return propagateTracingHeadersAndExecuteRequest(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 ClientHttpRequestExecution#execute(HttpRequest, byte[])} to execute the request.
     *
     * @return The result of calling {@link ClientHttpRequestExecution#execute(HttpRequest, byte[])}.
     */
    protected ClientHttpResponse propagateTracingHeadersAndExecuteRequest(
        HttpRequestWrapperWithModifiableHeaders wrapperRequest, byte[] body, ClientHttpRequestExecution execution
    ) throws IOException {
        propagateTracingHeaders(wrapperRequest, Tracer.getInstance().getCurrentSpan());
        
        // Execute the request/interceptor chain.
        return execution.execute(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 #propagateTracingHeadersAndExecuteRequest(HttpRequestWrapperWithModifiableHeaders,
     * byte[], ClientHttpRequestExecution)} to actually execute the request. Span naming and tagging is done here using
     * {@link #tagAndNamingStrategy} and {@link #tagAndNamingAdapter}.
     *
     * @return The result of calling {@link
     * #propagateTracingHeadersAndExecuteRequest(HttpRequestWrapperWithModifiableHeaders, byte[],
     * ClientHttpRequestExecution)} after surrounding the request with a subspan (or new trace if no current span
     * exists).
     */
    protected ClientHttpResponse createNewSpanAndExecuteRequest(
        HttpRequestWrapperWithModifiableHeaders wrapperRequest, byte[] body, ClientHttpRequestExecution execution
    ) throws IOException {
        // Will start a new trace if necessary, or a subspan if a trace is already in progress.
        Span spanAroundCall = Tracer.getInstance().startSpanInCurrentContext(
            getSubspanSpanName(wrapperRequest, tagAndNamingStrategy, tagAndNamingAdapter),
            SpanPurpose.CLIENT
        );

        Throwable errorForTagging = null;
        ClientHttpResponse response = null;
        try {
            tagAndNamingStrategy.handleRequestTagging(spanAroundCall, wrapperRequest, tagAndNamingAdapter);
            response = propagateTracingHeadersAndExecuteRequest(wrapperRequest, body, execution);

            return response;
        } catch(Throwable exception) {
            errorForTagging = exception;
            throw exception;
        }
        finally {
            try {
                // Handle response/error tagging and final span name.
                tagAndNamingStrategy.handleResponseTaggingAndFinalSpanName(
                    spanAroundCall, wrapperRequest, response, errorForTagging, tagAndNamingAdapter
                );
            }
            finally {
                // Span.close() contains the span-finishing 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.
                spanAroundCall.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 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(
            "resttemplate_downstream_call", getRequestMethodAsString(request.getMethod())
        );
    }
}