package ru.yandex.webmaster3.core.tracer;

import java.util.Objects;

import lombok.extern.slf4j.Slf4j;

import ru.yandex.webmaster3.core.util.ContextTracker;
import ru.yandex.webmaster3.core.util.json.JsonMapping;

/**
 * @author aherman
 */
@Slf4j
public class YdbTracer {
    private static final InheritableThreadLocal<YdbTrace> traces = new InheritableThreadLocal<>();
    private static final YdbTrace DUMMY_TRACE = new YdbTrace(0);

    public static YdbTrace startTrace() {
        long wallTime = System.nanoTime();
        YdbTrace trace = new YdbTrace(wallTime);
        traces.set(trace);

        return trace;
    }

    public static YdbTrace getCurrentTrace() {
        YdbTrace trace = traces.get();
        return Objects.requireNonNullElse(trace, DUMMY_TRACE);
    }

    public static YdbTrace stopTrace() {
        YdbTrace trace = traces.get();
        if (trace != null) {
            trace.endRequest(System.nanoTime());
            traces.remove();
            return trace;
        }

        return null;
    }

    @FunctionalInterface
    public interface SupplierWithException<T> {
        public T get() throws Exception;
    }

    public static <T> T trace(String target, SupplierWithException<T> supplier) throws Exception {
        startTrace();
        try {
            return supplier.get();
        } finally {
            var trace = stopTrace();
            YdbMetricsService.saveSolomon(trace);
            if (trace != null) {
                log.info("Request YDB trace pretty: {} {}", target, JsonMapping.writeValueAsString(new BeautyYdbTrace(trace)));
                log.info("Request YDB trace json: {} {}", target, JsonMapping.writeValueAsString(trace));
            }
        }
    }

    public static ContextTracker CONTEXT_TRACKER = () -> {
        YdbTrace trace = traces.get();
        return new ContextTracker.Context() {
            @Override
            public void restoreContext() {
                traces.set(trace);
            }

            @Override
            public void clearContext() {
                traces.remove();
            }
        };
    };
}
