package com.yammer.telemetry.service.resources;

import com.google.common.collect.ImmutableSet;
import com.yammer.telemetry.service.models.GefxData;
import com.yammer.telemetry.service.models.GefxEdge;
import com.yammer.telemetry.service.models.GefxNode;
import com.yammer.telemetry.tracing.*;

import javax.ws.rs.GET;
import javax.ws.rs.Path;
import javax.ws.rs.Produces;
import javax.ws.rs.core.MediaType;

@Path("/data/network")
@Produces(MediaType.APPLICATION_JSON)
public class NetworkResource {
    private final SpanSource spanSource;

    public NetworkResource(SpanSource spanSource) {
        this.spanSource = spanSource;
    }

    @GET
    public GefxData getNetworkJson() {
        final ImmutableSet.Builder<GefxNode> nodeBuilder = new ImmutableSet.Builder<>();
        final ImmutableSet.Builder <GefxEdge> edgeBuilder = new ImmutableSet.Builder<>();

        for (Trace trace : spanSource.getTraces()) {
            processTrace(trace, trace.getRoot(), nodeBuilder, edgeBuilder);
        }

        return new GefxData(nodeBuilder.build(), edgeBuilder.build());
    }

    private void processTrace(Trace trace, SpanData span, ImmutableSet.Builder<GefxNode> nodeBuilder, ImmutableSet.Builder<GefxEdge> edgeBuilder) {
        for (AnnotationData annotationData : trace.getAnnotations(span)) {
            if (AnnotationNames.SERVICE_NAME.equals(annotationData.getName())) {
                nodeBuilder.add(new GefxNode(annotationData.getMessage(), annotationData.getMessage()));

                for (SpanData childSpan : trace.getChildren(span)) {
                    for (AnnotationData childAnnotationData : trace.getAnnotations(childSpan)) {
                        if (AnnotationNames.SERVICE_NAME.equals(childAnnotationData.getName())) {
                            edgeBuilder.add(new GefxEdge(annotationData.getMessage(), childAnnotationData.getMessage()));
                        }
                    }
                }
            }

            for (SpanData childSpan : trace.getChildren(span)) {
                processTrace(trace, childSpan, nodeBuilder, edgeBuilder);
            }
        }
    }
}

@@ -400,16 +400,39 @@
 diaType;
+%0Aimport java.util.List;
 %0A%0A@Path(
@@ -1262,152 +1262,339 @@
     
-for (AnnotationData annotationData : trace.getAnnotations(span)) %7B%0A            if (AnnotationNames.SERVICE_NAME.equals(annotationData.getName())
+String serviceName = determineServiceName(trace.getAnnotations(span));%0A        nodeBuilder.add(new GefxNode(serviceName, serviceName));%0A%0A        for (SpanData childSpan : trace.getChildren(span)) %7B%0A            String childServiceName = determineServiceName(trace.getAnnotations(childSpan));%0A            if (childServiceName != null
 ) %7B%0A
@@ -1601,35 +1601,35 @@
                 
-nod
+edg
 eBuilder.add(new
@@ -1637,74 +1637,46 @@
 Gefx
-Node(annotationData.getMessage(), annotationData.getMessage()
+Edge(serviceName, childServiceName
 ));%0A
-%0A
     
@@ -1687,77 +1687,209 @@
     
-    for (SpanData childSpan : trace.getChildren(span)) %7B%0A            
+%7D%0A%0A            processTrace(trace, childSpan, nodeBuilder, edgeBuilder);%0A        %7D%0A    %7D%0A%0A    private String determineServiceName(List%3CAnnotationData%3E annotations) %7B%0A        String serviceName = null;%0A
     
@@ -1912,22 +1912,17 @@
 ionData 
-childA
+a
 nnotatio
@@ -1926,62 +1926,26 @@
 tion
-Data : trace.getAnnotations(childSpan)) %7B%0A            
+ : annotations) %7B%0A
     
@@ -1992,22 +1992,17 @@
 .equals(
-childA
+a
 nnotatio
@@ -1998,28 +1998,24 @@
 s(annotation
-Data
 .getName()))
@@ -2037,84 +2037,23 @@
     
-            edgeBuilder.add(new GefxEdge(annotationData.getMessage(), childA
+serviceName = a
 nnot
@@ -2057,20 +2057,16 @@
 notation
-Data
 .getMess
@@ -2070,22 +2070,18 @@
 essage()
-));%0A  
+;%0A
         
@@ -2088,237 +2088,218 @@
     
-          %7D%0A                    %7D%0A                %7D%0A            %7D%0A%0A            for (SpanData childSpan : trace.getChildren(span)) %7B%0A                processTrace(trace, childSpan, nodeBuilder, edgeBuilder);%0A            %7D%0A        %7D
+%7D else if (serviceName == null && AnnotationNames.FALLBACK_SERVICE_NAME.equals(annotation.getName())) %7B%0A                serviceName = annotation.getMessage();%0A            %7D%0A        %7D%0A        return serviceName;
 %0A   

Fix to use the new fallback service name.