/*
 * Copyright (c) 2010-2018. Axon Framework
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.axonframework.messaging.correlation;

import org.axonframework.messaging.Message;

import java.util.HashMap;
import java.util.Map;

/**
 * CorrelationDataProvider that provides the {@link Message#getIdentifier() identifier} of a {@link Message} to
 * other messages that are created as result of processing the first message.
 *
 * @author Rene de Waele
 */
public class MessageOriginProvider implements CorrelationDataProvider {
    private static final String DEFAULT_CORRELATION_KEY = "correlationId";
    private static final String DEFAULT_TRACE_KEY = "traceId";
    private static final String DEFAULT_TRACE_PARENT_KEY = "traceparent";

    /**
     * Returns the default metadata key for the correlation id of a message.
     *
     * @return the default metadata key for the correlation id
     */
    public static String getDefaultCorrelationKey() {
        return DEFAULT_CORRELATION_KEY;
    }

    /**
     * Returns the default metadata key for the trace id of a message.
     *
     * @return the default metadata key for the trace id
     */
    public static String getDefaultTraceKey() {
        return DEFAULT_TRACE_KEY;
    }


    private final String traceParentKey;

    private final String correlationKey;
    private final String traceKey;

    /**
     * Initializes a {@link MessageOriginProvider} that uses the default correlation id key:
     * {@link #getDefaultCorrelationKey()}, trace id key: {@link #getDefaultTraceKey()}, and trace parent key:
     * {@link #getDefaultTraceParentKey()}.
     */
    public MessageOriginProvider() {
        this(DEFAULT_CORRELATION_KEY, DEFAULT_TRACE_KEY, DEFAULT_TRACE_PARENT_KEY);
    }

    /**
     * Initializes a {@link MessageOriginProvider} that uses the given {@code correlationKey} and {@code traceKey}.
     *
     * @param correlationKey the key used to store the identifier of a message in the metadata of a resulting message
     * @param traceKey       the key used to store the identifier of the original message giving rise to the current
     *                       message
     */
    public MessageOriginProvider(String correlationKey, String traceKey) {
        this(correlationKey, traceKey, DEFAULT_TRACE_PARENT_KEY);
    }

    /**
     * Initializes a {@link MessageOriginProvider} that uses the given {@code correlationKey} and {@code traceKey}.
     *
     * @param correlationKey the key used to store the identifier of a message in the metadata of a resulting message
     * @param traceKey       the key used to store the identifier of the original message giving rise to the current
     *                       message
     * @param traceParentKey the key used to store the identifier of the trace parent to the current message
     */
    public MessageOriginProvider(String correlationKey, String traceKey, String traceParentKey) {
        this.correlationKey = correlationKey;
        this.traceKey = traceKey;
        this.traceParentKey = traceParentKey;
    }

    /**
     * Returns the default metadata key for the trace parent id of a message. This is based on the W3C standard for
     * tracing. See <a href="https://www.w3.org/TR/trace-context-1/">W3C Trace Context</a>.
     *
     * @return the default metadata key for the trace parent id
     */
    public static String getDefaultTraceParentKey() {
        return DEFAULT_TRACE_PARENT_KEY;
    }

    @Override
    public Map<String, ?> correlationDataFor(Message<?> message) {
        Map<String, Object> result = new HashMap<>();
        result.put(correlationKey, message.getIdentifier());
        result.put(traceKey, message.getMetaData().getOrDefault(traceKey, message.getIdentifier()));
        if (message.getMetaData().containsKey(traceParentKey)) {
            result.put(traceParentKey, message.getMetaData().get(traceParentKey));
        }
        return result;
    }
}

@@ -1184,82 +1184,8 @@
 Id%22;
-%0A    private static final String DEFAULT_TRACE_PARENT_KEY = %22traceparent%22;
 %0A%0A  
@@ -1690,51 +1690,8 @@
  %7D%0A%0A
-%0A    private final String traceParentKey;%0A%0A
     
@@ -1859,38 +1859,38 @@
 tion id key:
-%0A     *
  %7B@link
+%0A     *
  #getDefault
@@ -1910,83 +1910,26 @@
 y()%7D
-, trace id key: %7B@link #getDefaultTraceKey()%7D, and trace parent key:%0A     *
+ and trace id key:
  %7B@l
@@ -1944,30 +1944,24 @@
 DefaultTrace
-Parent
 Key()%7D.%0A    
@@ -2060,34 +2060,8 @@
 _KEY
-, DEFAULT_TRACE_PARENT_KEY
 );%0A 
@@ -2167,37 +2167,16 @@
 tionKey%7D
- and %7B@code traceKey%7D
 .%0A     *
@@ -2547,1208 +2547,74 @@
 this
-(correlationKey, traceKey, DEFAULT_TRACE_PARENT_KEY);%0A    %7D%0A%0A    /**%0A     * Initializes a %7B@link MessageOriginProvider%7D that uses the given %7B@code correlationKey%7D and %7B@code traceKey%7D.%0A     *%0A     * @param correlationKey the key used to store the identifier of a message in the metadata of a resulting message%0A     * @param traceKey       the key used to store the identifier of the original message giving rise to the current%0A     *                       message%0A     * @param traceParentKey the key used to store the identifier of the trace parent to the current message%0A     */%0A    public MessageOriginProvider(String correlationKey, String traceKey, String traceParentKey) %7B%0A        this.correlationKey = correlationKey;%0A        this.traceKey = traceKey;%0A        this.traceParentKey = traceParentKey;%0A    %7D%0A%0A    /**%0A     * Returns the default metadata key for the trace parent id of a message. This is based on the W3C standard for%0A     * tracing. See %3Ca href=%22https://www.w3.org/TR/trace-context-1/%22%3EW3C Trace Context%3C/a%3E.%0A     *%0A     * @return the default metadata key for the trace parent id%0A     */%0A    public static String getDefaultTraceParentKey() %7B%0A        return DEFAULT_TRACE_PARENT_KEY
+.correlationKey = correlationKey;%0A        this.traceKey = traceKey
 ;%0A  
@@ -2919,166 +2919,8 @@
 ));%0A
-        if (message.getMetaData().containsKey(traceParentKey)) %7B%0A            result.put(traceParentKey, message.getMetaData().get(traceParentKey));%0A        %7D%0A
     
@@ -2944,10 +2944,11 @@
 ;%0A    %7D%0A
+%0A
 %7D%0A

Revert changes to MessageOriginProvider since they are no longer needed