package com.jivesoftware.os.miru.service.partition;

import com.google.common.base.Optional;
import com.google.common.collect.Maps;
import com.jivesoftware.os.miru.api.MiruHost;
import com.jivesoftware.os.miru.api.MiruPartitionCoord;
import com.jivesoftware.os.miru.api.MiruPartitionCoordInfo;
import com.jivesoftware.os.miru.api.activity.MiruPartitionId;
import com.jivesoftware.os.miru.api.topology.MiruClusterClient;
import com.jivesoftware.os.miru.api.topology.MiruHeartbeatRequest;
import com.jivesoftware.os.miru.api.topology.MiruHeartbeatResponse;
import com.jivesoftware.os.miru.api.topology.MiruPartitionActive;
import com.jivesoftware.os.miru.api.topology.MiruPartitionActiveUpdate;
import com.jivesoftware.os.miru.api.topology.NamedCursor;
import com.jivesoftware.os.miru.api.topology.PartitionInfo;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;

/**
 *
 */
public class MiruPartitionHeartbeatHandler {

    private final MiruClusterClient clusterClient;

    private final AtomicReference<Map<MiruPartitionCoord, PartitionInfo>> heartbeats = new AtomicReference<>(
        (Map<MiruPartitionCoord, PartitionInfo>) new ConcurrentHashMap<MiruPartitionCoord, PartitionInfo>());
    private final AtomicReference<Map<MiruPartitionCoord, MiruPartitionActive>> active = new AtomicReference<>();

    private final Object cursorLock = new Object();
    private final Map<String, NamedCursor> partitionActiveUpdatesSinceCursors = Maps.newHashMap();
    private final Map<String, NamedCursor> topologyUpdatesSinceCursors = Maps.newHashMap();

    public MiruPartitionHeartbeatHandler(MiruClusterClient clusterClient) {
        this.clusterClient = clusterClient;
    }

    public void heartbeat(MiruPartitionCoord coord,
        Optional<MiruPartitionCoordInfo> info,
        Optional<Long> refreshTimestamp)
        throws Exception {

        Map<MiruPartitionCoord, PartitionInfo> beats;
        do {
            beats = heartbeats.get();
            PartitionInfo got = beats.get(coord);
            if (got == null) {
                got = new PartitionInfo(coord.tenantId,
                    coord.partitionId.getId(),
                    refreshTimestamp.or(-1L),
                    info.orNull());
                beats.put(coord, got);
            } else {
                got = new PartitionInfo(coord.tenantId,
                    coord.partitionId.getId(),
                    refreshTimestamp.or(got.activeTimestamp),
                    info.isPresent() ? info.get() : got.info);
                beats.put(coord, got);
            }
        }
        while (beats != heartbeats.get());
    }

    public MiruHeartbeatResponse thumpthump(MiruHost host) throws Exception {
        synchronized (cursorLock) {
            MiruHeartbeatResponse thumpthump = clusterClient.thumpthump(host,
                new MiruHeartbeatRequest(heartbeats(), partitionActiveUpdatesSinceCursors.values(), topologyUpdatesSinceCursors.values()));

            setActive(host, thumpthump.activeHasChanged.result);
            handleCursors(partitionActiveUpdatesSinceCursors, thumpthump.activeHasChanged.cursors);
            handleCursors(topologyUpdatesSinceCursors, thumpthump.topologyHasChanged.cursors);

            return thumpthump;
        }
    }

    public MiruPartitionActive getPartitionActive(MiruPartitionCoord coord) throws Exception {
        Map<MiruPartitionCoord, MiruPartitionActive> got = active.get();
        if (got != null) {
            MiruPartitionActive partitionActive = got.get(coord);
            if (partitionActive != null) {
                return partitionActive;
            }
        }
        return new MiruPartitionActive(false, false);
    }

    private Collection<PartitionInfo> heartbeats() {
        Map<MiruPartitionCoord, PartitionInfo> beats = heartbeats.get();
        if (beats != null) {
            heartbeats.compareAndSet(beats, new ConcurrentHashMap<MiruPartitionCoord, PartitionInfo>());
            return beats.values();
        } else {
            return Collections.emptyList();
        }
    }

    private void setActive(MiruHost host, Collection<MiruPartitionActiveUpdate> updates) {
        Map<MiruPartitionCoord, MiruPartitionActive> coordActive = new HashMap<>(updates.size());
        for (MiruPartitionActiveUpdate update : updates) {
            MiruPartitionCoord coord = new MiruPartitionCoord(update.tenantId, MiruPartitionId.of(update.partitionId), host);
            MiruPartitionActive partitionActive = new MiruPartitionActive(update.active, update.idle);
            coordActive.put(coord, partitionActive);
        }
        active.set(coordActive);
    }

    /**
     * Must be holding cursor lock.
     */
    private void handleCursors(Map<String, NamedCursor> cursors, Collection<NamedCursor> updates) {
        for (NamedCursor update : updates) {
            NamedCursor existing = cursors.get(update.name);
            if (existing == null || existing.id < update.id) {
                cursors.put(update.name, update);
            }
        }
    }

}

@@ -867,34 +867,8 @@
 ns;%0A
-import java.util.HashMap;%0A
 impo
@@ -928,24 +928,67 @@
 entHashMap;%0A
+import java.util.concurrent.ConcurrentMap;%0A
 import java.
@@ -1375,32 +1375,26 @@
 e final 
-AtomicReference%3C
+Concurrent
 Map%3CMiru
@@ -1429,17 +1429,16 @@
 nActive%3E
-%3E
  active 
@@ -1439,37 +1439,37 @@
 ctive = 
-new AtomicReference%3C%3E
+Maps.newConcurrentMap
 ();%0A%0A   
@@ -3526,152 +3526,51 @@
    M
-ap%3CMiruPartitionCoord, MiruPartitionActive%3E got = active.get();%0A        if (got != null) %7B%0A            MiruPartitionActive partitionActive = got
+iruPartitionActive partitionActive = active
 .get
@@ -3570,36 +3570,32 @@
 ive.get(coord);%0A
-    
         if (part
@@ -3625,28 +3625,24 @@
             
-  
-  
 return parti
@@ -3653,30 +3653,16 @@
 Active;%0A
-            %7D%0A
         
@@ -4192,106 +4192,8 @@
 ) %7B%0A
-        Map%3CMiruPartitionCoord, MiruPartitionActive%3E coordActive = new HashMap%3C%3E(updates.size());%0A
     
@@ -4488,22 +4488,17 @@
         
-coordA
+a
 ctive.pu
@@ -4538,41 +4538,8 @@
   %7D%0A
-        active.set(coordActive);%0A
     

fix active checks