package com.infoloop.tianting.tracing;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import com.infoloop.tianting.constant.ConfigConstants;
import org.slf4j.MDC;
import org.springframework.core.annotation.Order;
import org.springframework.lang.Nullable;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

/**
 * Note: tracing context is initialized here since this project is a webapp.
 * Therefore this interceptor should be placed as the first layer.
 */
@Order(1)
@SuppressWarnings("all")
public class HttpTracingInterceptor implements HandlerInterceptor {
    private final Tracing tracing;

    public HttpTracingInterceptor(final Tracing tracing) {
        this.tracing = tracing;
    }

    @Override
    public boolean preHandle(final HttpServletRequest request,
                             final HttpServletResponse response,
                             final Object handler) throws Exception {
        final var tracer = tracing.tracer();
        if (request.getAttribute(Tracer.SpanInScope.class.getName()) != null) {
            return true; // already handled (possibly due to async request)
        }
        final var span = tracer.newTrace();
        span.name(request.getMethod() + " " + request.getServletPath());
        ConfigConstants.BRAVE_PROPAGATION_DEBUG_FIELD.updateValue(span.context(), "true");
        span.start();
        MDC.put(ConfigConstants.LOGGING_TRACING_ID, span.context().traceIdString());
        MDC.put(ConfigConstants.LOGGING_UNIQUE_ID, span.context().spanIdString());
        MDC.put(ConfigConstants.LOGGING_PARENT_TRACING_ID, span.context().parentIdString());
        request.setAttribute(Tracer.SpanInScope.class.getName(), tracer.withSpanInScope(span));
        return true;
    }

    @Override
    public void afterCompletion(final HttpServletRequest request,
                                final HttpServletResponse response,
                                final Object handler,
                                @Nullable final Exception ex) throws Exception {
        final var tracer = tracing.tracer();
        final Span span = tracer.currentSpan();
//        ((Tracer.SpanInScope) request.getAttribute(Tracer.SpanInScope.class.getName())).close();
        if (ex != null) {
            span.customizer().tag("http-error", ex.toString());
        }
        span.finish();
        MDC.clear();
    }

}