package io.quarkus.jaeger.runtime;

import org.jboss.logging.MDC;

import io.jaegertracing.internal.JaegerSpanContext;
import io.opentracing.Scope;
import io.opentracing.Span;

/**
 * Scope that sets span context into MDC.
 */
public class MDCScope implements Scope {

    /**
     * MDC keys
     */
    private static final String TRACE_ID = "traceId";
    private static final String SPAN_ID = "spanId";
    private static final String SAMPLED = "sampled";

    private final Scope wrapped;
    private final Object originalTraceId;
    private final Object originalSpanId;
    private final Object originalSampled;

    public MDCScope(Scope scope) {
        this.wrapped = scope;
        this.originalTraceId = MDC.get(TRACE_ID);
        this.originalSpanId = MDC.get(SPAN_ID);
        this.originalSampled = MDC.get(SAMPLED);
        if (scope.span().context() instanceof JaegerSpanContext) {
            putContext((JaegerSpanContext) scope.span().context());
        }
    }

    @Override
    public void close() {
        wrapped.close();
        MDC.remove(TRACE_ID);
        MDC.remove(SPAN_ID);
        MDC.remove(SAMPLED);

        if (originalTraceId != null) {
            MDC.put(TRACE_ID, originalTraceId);
        }
        if (originalSpanId != null) {
            MDC.put(SPAN_ID, originalSpanId);
        }
        if (originalSampled != null) {
            MDC.put(SAMPLED, originalSampled);
        }
    }

    @Override
    public Span span() {
        return wrapped.span();
    }

    protected void putContext(JaegerSpanContext spanContext) {
        MDC.put(TRACE_ID, spanContext.getTraceId());
        MDC.put(SPAN_ID, Long.toHexString(spanContext.getSpanId()));
        MDC.put(SAMPLED, Boolean.toString(spanContext.isSampled()));
    }
}
