package org.restheart.handlers.metrics;

import org.restheart.Bootstrapper;
import org.slf4j.MDC;

import java.util.Optional;

import io.undertow.server.HttpServerExchange;
import io.undertow.util.HttpString;
import org.restheart.handlers.PipedHttpHandler;
import org.restheart.handlers.RequestContext;

/**
 * Handler to write tracing headers to the logging MDC. Pick it up via the
 * default way with "%X{name}", e.g. "%X{x-b3-traceid}".
 */
public class TracingInstrumentationHandler extends PipedHttpHandler {
    public TracingInstrumentationHandler(PipedHttpHandler next) {
        super(next);
    }

    @Override
    public void handleRequest(HttpServerExchange exchange, RequestContext context) throws Exception {
        for (String traceIdHeader : Bootstrapper.getConfiguration().getTraceHeaders()) {
            Optional.ofNullable(exchange.getRequestHeaders().get(traceIdHeader))
                    .flatMap(x -> Optional.ofNullable(x.peekFirst()))
                    .ifPresent(value -> {
                        MDC.put(traceIdHeader, value);
                        exchange.getResponseHeaders().put(HttpString.tryFromString(traceIdHeader), value);
                    });
        }

        if (!exchange.isResponseComplete() && getNext() != null) {
            next(exchange, context);
        }

        for (String traceIdHeader : Bootstrapper.getConfiguration().getTraceHeaders()) {
            MDC.remove(traceIdHeader);
        }
    }
}

@@ -673,16 +673,28 @@
 xchange,
+%0A           
  Request
@@ -740,36 +740,8 @@
     
- for (String traceIdHeader :
  Boo
@@ -778,36 +778,87 @@
 etTraceHeaders()
-) %7B%0A
+%0A                .forEach((traceIdHeader) -%3E %7B%0A        
             Opti
@@ -901,16 +901,45 @@
 eaders()
+%0A                            
 .get(tra
@@ -971,16 +971,24 @@
         
+     
+   
 .flatMap
@@ -1049,16 +1049,24 @@
         
+        
 .ifPrese
@@ -1099,24 +1099,32 @@
             
+  
+      
 MDC.put(trac
@@ -1166,16 +1166,24 @@
         
+      
+  
 exchange
@@ -1207,53 +1207,191 @@
 rs()
-.put(HttpString.tryFromString(traceIdHeader),
+%0A                                        .put(HttpString%0A                                                .tryFromString(traceIdHeader),%0A                                               
  val
@@ -1387,32 +1387,40 @@
         value);%0A
+        
                 
@@ -1427,33 +1427,43 @@
     %7D);%0A        
-%7D
+        %7D);
 %0A%0A        if (!e
@@ -1574,36 +1574,8 @@
     
- for (String traceIdHeader :
  Boo
@@ -1594,32 +1594,49 @@
 tConfiguration()
+%0A                
 .getTraceHeaders
@@ -1637,20 +1637,54 @@
 eaders()
-) %7B%0A
+.forEach((traceIdHeader) -%3E %7B%0A        
         
@@ -1722,17 +1722,27 @@
         
-%7D
+        %7D);
 %0A    %7D%0A%7D

Refactor TracingInstrumentationHandler to use lambda expressions