Set zone levels

1 files changed, 12 insertions, 0 deletions

diff --git a/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/Stanik.java b/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/Stanik.java
index 999c193..f572efe 100644
--- a/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/Stanik.java
+++ b/src/main/java/pl/edu/mimuw/cloudatlas/agent/modules/Stanik.java
@@ -61,6 +61,7 @@ public class Stanik extends Module {
 
     public void handleGetState(GetStateMessage message) throws InterruptedException {
         pruneHierarchy();
+        addLevels();
         StateMessage response = new StateMessage(
             "",
             message.getRequestingModule(),
@@ -78,6 +79,17 @@ public class Stanik extends Module {
         pruneZMI(hierarchy, now);
     }
 
+    private void addLevels() {
+        addLevelsRecursive(hierarchy, 0);
+    }
+
+    private void addLevelsRecursive(ZMI zmi, long level) {
+        zmi.getAttributes().addOrChange("level", new ValueInt(level));
+        for (ZMI son : zmi.getSons()) {
+            addLevelsRecursive(son, level + 1);
+        }
+    }
+
     private boolean pruneZMI(ZMI zmi, ValueTime time) {
         Value timestamp = zmi.getAttributes().get("timestamp");