package io.fluxcapacitor.javaclient.publishing.dataprotection;

import io.fluxcapacitor.common.api.SerializedMessage;
import io.fluxcapacitor.common.handling.Handler;
import io.fluxcapacitor.common.reflection.ReflectionUtils;
import io.fluxcapacitor.javaclient.common.Message;
import io.fluxcapacitor.javaclient.common.serialization.DeserializingMessage;
import io.fluxcapacitor.javaclient.keyvalue.KeyValueStore;
import io.fluxcapacitor.javaclient.publishing.DispatchInterceptor;
import io.fluxcapacitor.javaclient.tracking.handling.HandlerInterceptor;
import lombok.AllArgsConstructor;

import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

import static io.fluxcapacitor.common.reflection.ReflectionUtils.getAnnotatedFields;
import static io.fluxcapacitor.common.reflection.ReflectionUtils.getProperty;
import static io.fluxcapacitor.common.reflection.ReflectionUtils.setField;
import static java.util.UUID.randomUUID;

@AllArgsConstructor
public class DataProtectionInterceptor implements DispatchInterceptor, HandlerInterceptor {

    public static String METADATA_KEY = "$protectedData";

    private final KeyValueStore keyValueStore;

    @Override
    public Function<Message, SerializedMessage> interceptDispatch(Function<Message, SerializedMessage> function) {
        return m -> {
            if (!m.getMetadata().containsKey(METADATA_KEY)) {
                Map<String, String> protectedFields = new HashMap<>();
                getAnnotatedFields(m.getPayload(), ProtectData.class).forEach(field -> {
                    Object value = getProperty(field, m.getPayload());
                    String key = randomUUID().toString();
                    keyValueStore.store(key, value);
                    protectedFields.put(field.getName(), key);
                    setField(field, m.getPayload(), null);
                });
                if (!protectedFields.isEmpty()) {
                    m.getMetadata().put(METADATA_KEY, protectedFields);
                }
            }
            return function.apply(m);
        };
    }

    @Override
    @SuppressWarnings("unchecked")
    public Function<DeserializingMessage, Object> interceptHandling(Function<DeserializingMessage, Object> function,
                                                                    Handler<DeserializingMessage> handler,
                                                                    String consumer) {
        return m -> {
            if (m.getMetadata().containsKey(METADATA_KEY)) {
                Object payload = m.getPayload();
                Map<String, String> protectedFields = m.getMetadata().get(METADATA_KEY, Map.class);
                boolean dropProtectedData = handler.getMethod(m).isAnnotationPresent(DropProtectedData.class);
                protectedFields.forEach((fieldName, key) -> {
                    ReflectionUtils.setField(fieldName, payload, keyValueStore.get(key));
                    if (dropProtectedData) {
                        keyValueStore.delete(key);
                    }
                });
            }
            return function.apply(m);
        };
    }
}

@@ -1166,17 +1166,21 @@
 eStore;%0A
+    
 %0A
-
     @Ove
@@ -1181,24 +1181,59 @@
   @Override%0A
+    @SuppressWarnings(%22unchecked%22)%0A
     public F
@@ -1361,74 +1361,8 @@
 %3E %7B%0A
-            if (!m.getMetadata().containsKey(METADATA_KEY)) %7B%0A    
     
@@ -1420,24 +1420,192 @@
 ashMap%3C%3E();%0A
+            if (m.getMetadata().containsKey(METADATA_KEY)) %7B%0A                protectedFields.putAll(m.getMetadata().get(METADATA_KEY, Map.class));%0A            %7D else %7B%0A
             
@@ -1930,67 +1930,8 @@
 y);%0A
-                    setField(field, m.getPayload(), null);%0A
     
@@ -2078,32 +2078,32 @@
               %7D%0A
-
             %7D%0A  
@@ -2092,32 +2092,122 @@
 %7D%0A            %7D%0A
+            protectedFields.forEach((name, key) -%3E setField(name, m.getPayload(), null));%0A
             retu

Fix in DataProtectionInterceptor