package com.spotify.wasd.db;

import lombok.Getter;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.apache.cassandra.tools.NodeProbe;

import java.io.IOException;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

@Slf4j
@ToString(of = {"siteRecord"}, includeFieldNames = false)
public class CassandraSiteCluster {
    @Getter
    private final SiteRecord siteRecord;
    @Getter
    private final Set<Host> hostSet;
    @Getter
    private final CassandraCluster cassandraCluster;

    CassandraSiteCluster(final Hosts hosts, final SiteRecord siteRecord, CassandraCluster cassandraCluster) {
        this.siteRecord = siteRecord;
        this.cassandraCluster = cassandraCluster;
        Set<Host> readHostSet = null;

        if (siteRecord.getHostSet().size() == 0) {
            log.debug("Empty record for {}", this);
            this.hostSet = new HashSet<Host>();
            return;
        }

        for (Host node : siteRecord.getHostSet()) {
            try {
                log.debug("Contacting {} for {}", node, this);
                readHostSet = getHostSetFromNode(hosts, node);
                log.debug("{} worked for {}", node, this);
                break;
            } catch (IOException e) {
                hosts.reportFailedRingForHost(node);
                log.warn("{} failed for {}", node, this);
                log.warn("{}: {}", node, e);
            } catch (InterruptedException e) {
                hosts.reportFailedRingForHost(node);
                log.warn("{} failed for {}", node, this);
                log.warn("{}: {}", node, e);
            }
        }

        if (readHostSet == null) {
            log.error("Could not read {}", this);
            this.hostSet = new HashSet<Host>();
        } else {
            this.hostSet = readHostSet;
        }
    }

    final Set<Host> getHostSetFromNode(Hosts hosts, Host node) throws IOException, InterruptedException {
        final Set<Host> hostSet = new HashSet<Host>();

        NodeProbe probe = new NodeProbe(node.getReverseName());

        final Map<String, String> tokenToEndpointMap = probe.getTokenToEndpointMap();
        for (Map.Entry<String, String> entry : tokenToEndpointMap.entrySet()) {
            log.debug("{} has endpoint {}", this, entry.getValue());
            final Host host = hosts.getHostByName(entry.getValue());
            hostSet.add(host);
            host.addToCassandraSiteCluster(this);
        }

        probe.close();

        return hostSet;
    }
}

@@ -136,22 +136,315 @@
 ra.t
-ools.NodeProbe
+hrift.Cassandra;%0Aimport org.apache.cassandra.thrift.InvalidRequestException;%0Aimport org.apache.thrift.TException;%0Aimport org.apache.thrift.protocol.TBinaryProtocol;%0Aimport org.apache.thrift.transport.TFramedTransport;%0Aimport org.apache.thrift.transport.TSocket;%0Aimport org.apache.thrift.transport.TTransport
 ;%0A%0Ai
@@ -2310,78 +2310,286 @@
     
-NodeProbe probe = new NodeProbe(node.getReverseName());%0A%0A        final
+final TTransport transport = new TFramedTransport(new TSocket(node.getAddress().getHostAddress(), 9160));%0A        final TBinaryProtocol protocol = new TBinaryProtocol(transport);%0A%0A        final Cassandra.Client client = new Cassandra.Client(protocol);%0A%0A        try %7B%0A           
  Map
@@ -2630,47 +2630,46 @@
 p = 
-probe.getTokenToEndpointM
+client.describe_token_m
 ap();%0A
+%0A
         
 for 
@@ -2660,24 +2660,28 @@
 );%0A%0A        
+    
 for (Map.Ent
@@ -2732,32 +2732,36 @@
 p.entrySet()) %7B%0A
+    
             log.
@@ -2821,24 +2821,28 @@
             
+    
 final Host h
@@ -2878,32 +2878,36 @@
 ry.getValue());%0A
+    
             host
@@ -2929,24 +2929,28 @@
             
+    
 host.addToCa
@@ -2987,33 +2987,229 @@
     
-%7D%0A%0A        probe.close();
+    %7D%0A        %7D catch (InvalidRequestException e) %7B%0A            throw new IOException(e);%0A        %7D catch (TException e) %7B%0A            throw new IOException(e);%0A        %7D finally %7B%0A            transport.close();%0A        %7D
 %0A%0A  

use thrift instead of NodeProbe