package org.datadog.jenkins.plugins.datadog.traces;

import datadog.trace.api.DDId;
import hudson.EnvVars;
import hudson.Extension;
import hudson.model.Run;
import hudson.model.TaskListener;
import org.jenkinsci.plugins.workflow.graph.FlowNode;
import org.jenkinsci.plugins.workflow.steps.StepContext;
import org.jenkinsci.plugins.workflow.steps.StepEnvironmentContributor;

import java.io.IOException;
import java.util.logging.Logger;

/**
 * Contributes the X_DATADOG_TRACE_ID and X_DATADOG_PARENT_ID environment variables to workflow steps.
 * The value of the X_DATADOG_PARENT_ID is the spanID of the span that is going to be executed, therefore
 * this method is executed as many as workflow steps the pipeline have.
 */
@Extension
public class TraceStepEnvironmentContributor extends StepEnvironmentContributor {

    private static final Logger logger = Logger.getLogger(TraceStepEnvironmentContributor.class.getName());

    public static final String TRACE_ID_ENVVAR_KEY = "X_DATADOG_TRACE_ID";
    public static final String SPAN_ID_ENVVAR_KEY = "X_DATADOG_PARENT_ID";

    @Override
    public void buildEnvironmentFor(StepContext stepContext, EnvVars envs, TaskListener listener) throws IOException, InterruptedException {
        try {
            final Run<?,?> run = stepContext.get(Run.class);
            final BuildSpanAction buildSpanAction = run.getAction(BuildSpanAction.class);
            if(buildSpanAction != null) {
                // The key 'x-datadog-trace-id' is used by the Java Tracer to inject the SpanContext
                // into the buildSpanPropagation map. As we don't have access to the tracer here,
                // we use the key directly to obtain the traceID value.
                final String traceId = buildSpanAction.getBuildSpanPropatation().get("x-datadog-trace-id");
                envs.put(TRACE_ID_ENVVAR_KEY, DDId.from(traceId).toHexString());
            }

            final FlowNode flowNode = stepContext.get(FlowNode.class);
            if(flowNode != null) {
                final GeneratedSpanIdAction idsAction = flowNode.getAction(GeneratedSpanIdAction.class);
                if(idsAction != null) {
                    envs.put(SPAN_ID_ENVVAR_KEY, idsAction.getDDSpanId().toHexString());
                }
            }
        } catch (Exception ex) {
            logger.severe("Unable to set traces IDs as environment variables before step execution. " + ex);
        }

    }
}

@@ -471,16 +471,39 @@
 TRACE_ID
+/X_DATADOG_TRACE_ID_HEX
  and X_D
@@ -518,16 +518,40 @@
 ARENT_ID
+/X_DATADOG_PARENT_ID_HEX
  environ
@@ -1079,49 +1079,215 @@
 ing 
-SPAN_ID_ENVVAR_KEY = %22X_DATADOG_PARENT_ID
+TRACE_ID_HEX_ENVVAR_KEY = %22X_DATADOG_TRACE_ID_HEX%22;%0A    public static final String SPAN_ID_ENVVAR_KEY = %22X_DATADOG_PARENT_ID%22;%0A    public static final String SPAN_ID_HEX_ENVVAR_KEY = %22X_DATADOG_PARENT_ID_HEX
 %22;%0A%0A
@@ -2047,56 +2047,173 @@
     
-envs.put(TRACE_ID_ENVVAR_KEY, DDId.from(t
+DDId ddTraceId = DDId.from(traceId);%0A                envs.put(TRACE_ID_ENVVAR_KEY, ddTraceId.toString());%0A                envs.put(TRACE_ID_HEX_ENVVAR_KEY, ddT
 raceId
-)
 .toH
@@ -2515,60 +2515,181 @@
     
-envs.put(SPAN_ID_ENVVAR_KEY, idsAction.getDD
+DDId ddSpanId = idsAction.getDDSpanId();%0A                    envs.put(SPAN_ID_ENVVAR_KEY, ddSpanId.toString());%0A                    envs.put(SPAN_ID_HEX_ENVVAR_KEY, dd
 SpanId
-()
 .toH

Add dec and hex flavour to the environment variable trace IDs. (#152)