package io.fluxcapacitor.javaclient.tracking.metrics;

import io.fluxcapacitor.common.handling.Handler;
import io.fluxcapacitor.javaclient.FluxCapacitor;
import io.fluxcapacitor.javaclient.common.serialization.DeserializingMessage;
import io.fluxcapacitor.javaclient.tracking.handling.HandlerInterceptor;
import lombok.extern.slf4j.Slf4j;

import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.function.Function;

@Slf4j
public class HandlerMonitor implements HandlerInterceptor {
    @Override
    public Function<DeserializingMessage, Object> interceptHandling(Function<DeserializingMessage, Object> function,
                                                                    Handler<DeserializingMessage> handler,
                                                                    String consumer) {
        return message -> {
            Instant start = Instant.now();
            try {
                Object result = function.apply(message);
                publishMetrics(handler, consumer, message, false, start, result);
                return result;
            } catch (Throwable e) {
                publishMetrics(handler, consumer, message, true, start, e);
                throw e;
            }
        };
    }

    protected void publishMetrics(Handler<DeserializingMessage> handler, String consumer, DeserializingMessage message,
                                  boolean exceptionalResult, Instant start, Object result) {
        try {
            boolean completed = !(result instanceof CompletableFuture<?>) || ((CompletableFuture<?>) result).isDone();
            FluxCapacitor.publishMetrics(new HandleMessageEvent(
                    FluxCapacitor.get().client().name(), FluxCapacitor.get().client().id(), consumer,
                    handler.getTarget().getClass().getSimpleName(), message.getSerializedObject().getIndex(),
                    message.getPayloadClass().getSimpleName(), exceptionalResult,
                    start.until(Instant.now(), ChronoUnit.NANOS), completed));
            if (!completed) {
                ((CompletionStage<?>) result).whenComplete((r, e) -> FluxCapacitor.publishMetrics(
                        new CompleteMessageEvent(
                                FluxCapacitor.get().client().name(), FluxCapacitor.get().client().id(), consumer,
                                handler.getTarget().getClass().getSimpleName(),
                                message.getSerializedObject().getIndex(),
                                message.getPayloadClass().getSimpleName(), e != null,
                                start.until(Instant.now(), ChronoUnit.NANOS))));
            }
        } catch (Exception e) {
            log.error("Failed to publish handler metrics", e);
        }
    }
}

@@ -2228,16 +2228,220 @@
 r, e) -%3E
+ %7B%0A                    DeserializingMessage current = DeserializingMessage.getCurrent();%0A                    try %7B%0A                        DeserializingMessage.setCurrent(message);%0A                       
  FluxCap
@@ -2487,16 +2487,24 @@
         
+        
 new Comp
@@ -2513,32 +2513,40 @@
 teMessageEvent(%0A
+        
                 
@@ -2667,32 +2667,40 @@
                 
+ 
+       
 handler.getTarge
@@ -2755,32 +2755,40 @@
                 
+     
+   
 message.getSeria
@@ -2805,32 +2805,40 @@
 t().getIndex(),%0A
+        
                 
@@ -2931,32 +2931,40 @@
                 
+ 
+       
 start.until(Inst
@@ -2993,16 +2993,155 @@
 NANOS)))
+;%0A                    %7D finally %7B%0A                        DeserializingMessage.setCurrent(current);%0A                    %7D%0A                %7D
 );%0A     

Add correlation data when publishing CompleteMessageEvent metrics