package org.yraft.network;

import com.google.common.collect.Maps;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.yraft.Messages;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.net.SocketException;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.Map;

public class UDPCommunicator extends AbstractCommunicator {

  private static final Logger LOG = LoggerFactory.getLogger(UDPCommunicator.class);

  private Map<Integer, InetSocketAddress> serverTable = Collections.EMPTY_MAP;
  private final DatagramSocket socket;
  private final ByteBuffer buf;
  private static final int MAX_BUF_SIZE = 8192;

  public UDPCommunicator(SocketAddress addr) throws SocketException {
    this.socket = new DatagramSocket(addr);
    this.buf = ByteBuffer.allocate(MAX_BUF_SIZE);
  }

  @Override
  public void setMembers(Map<Integer, InetSocketAddress> servers) {
    serverTable = Maps.newHashMap(servers);
  }

  @Override
  public void broadcast(Messages msg, byte[] data) {

    ByteBuffer packet = ByteBuffer.allocate(4 + data.length);
    packet.putInt(msg.getId());
    packet.put(data);

    for (InetSocketAddress addr : this.serverTable.values()) {
      sendTo(addr, packet.array());
    }
  }

  @Override
  public void sendTo(int candidateId, Messages msg, byte[] data) {

    // LOG.debug("Send message from " + this.server.getCandidateId() + " to " + candidateId);
    InetSocketAddress addr = this.serverTable.get(candidateId);
    if (addr != null) {
      sendTo(addr, msg, data);
    }
  }

  @Override
  public void run() {
    while (true) {
      DatagramPacket request = new DatagramPacket(this.buf.array(), this.buf.capacity());
      try {
        socket.receive(request);
        ByteBuffer buf = ByteBuffer.wrap(request.getData(), 0, request.getLength());
        Messages msg = Messages.fromId(buf.getInt());
        byte[] data = new byte[buf.remaining()];
        buf.get(data);
        server.receive(msg, data);

      } catch (IOException e) {
        LOG.error("", e);
      } catch (Exception e) {
        LOG.error("", e);
      }
    }
  }

  public void sendTo(SocketAddress addr, Messages msg, byte[] data) {

    ByteBuffer packet = ByteBuffer.allocate(4 + data.length);
      packet.putInt(msg.getId());
      packet.put(data);
    sendTo(addr, packet.array());
  }

  public void sendTo(SocketAddress addr, byte[] data) {

    try {
      this.socket.send(new DatagramPacket(data, data.length, addr));
    } catch (IOException e) {
      LOG.error(e.toString());
    }
  }
}

@@ -652,22 +652,16 @@
 private 
-final 
 Datagram
@@ -675,16 +675,52 @@
 socket;%0A
+  private final SocketAddress addr;%0A
   privat
@@ -875,41 +875,19 @@
 his.
-socket = new DatagramSocket(
+addr = 
 addr
-)
 ;%0A  
@@ -1695,24 +1695,163 @@
 oid run() %7B%0A
+    try %7B%0A      this.socket = new DatagramSocket(this.addr);%0A    %7D catch (SocketException e) %7B%0A      throw new RuntimeException(e);%0A    %7D%0A%0A
     while (t
@@ -2246,34 +2246,32 @@
 %0A      %7D catch (
-IO
 Exception e) %7B%0A 
@@ -2305,57 +2305,107 @@
     
-%7D catch (Exception e) %7B%0A        LOG.error(%22%22, e);
+  if (this.socket != null) %7B%0A          this.socket.close();%0A          this.socket = null;%0A        %7D
 %0A   

move UDP socket initilization to run method