package com.cisco.trex.stateless;

import com.cisco.trex.stateless.model.RPCResponse;
import com.cisco.trex.stateless.util.IDataCompressor;
import com.cisco.trex.stateless.util.TRexDataCompressor;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zeromq.SocketType;
import org.zeromq.ZMQ;
import org.zeromq.ZMQException;
import zmq.ZError;

/** TRex Transport class to create zmq socket for connection to trex server */
public class TRexTransport {

  private static final Logger LOGGER = LoggerFactory.getLogger(TRexTransport.class);
  private static final int DEFAULT_TIMEOUT = 3000;
  private static final String PROTOCOL = "tcp";
  private final String connectionString;
  private final IDataCompressor dataCompressor;
  private final ZMQ.Context zmqCtx = ZMQ.context(1);
  private final ZMQ.Socket zmqSocket;
  private final String host;
  private final String port;

  /**
   * @param host Server address
   * @param port Server port
   * @param timeout How long to wait for server response
   * @param dataCompressor
   */
  public TRexTransport(String host, String port, int timeout, IDataCompressor dataCompressor) {
    this.host = host;
    this.port = port;
    this.zmqSocket = zmqCtx.socket(SocketType.REQ);
    int actualTimeout = timeout <= 0 ? DEFAULT_TIMEOUT : timeout;
    zmqSocket.setReceiveTimeOut(actualTimeout);
    zmqSocket.setSendTimeOut(actualTimeout);
    this.connectionString = PROTOCOL + "://" + this.host + ":" + this.port;
    zmqSocket.connect(connectionString);
    this.dataCompressor = dataCompressor;
  }

  /**
   * @param host Server address
   * @param port Server port
   * @param timeout How long to wait for server response
   */
  public TRexTransport(String host, String port, int timeout) {
    this(host, port, timeout, new TRexDataCompressor());
  }

  /**
   * Send TRex command to the server
   *
   * @param command
   * @return RPCResponse
   * @throws IOException
   */
  public RPCResponse sendCommand(TRexCommand command) throws IOException {
    String json = new ObjectMapper().writeValueAsString(command.getParameters());
    String response = sendJson(json);
    ObjectMapper objectMapper = new ObjectMapper();
    if (objectMapper.readTree(response).isArray()) {
      // for versions of TRex before v2.61, single entry response also wrapped with json array
      RPCResponse[] rpcResult = objectMapper.readValue(response, RPCResponse[].class);
      return rpcResult[0];
    }
    return objectMapper.readValue(response, RPCResponse.class);
  }

  /**
   * Send TRex command list to the server
   *
   * @param commands
   * @return RPCResponse
   * @throws IOException
   */
  public RPCResponse[] sendCommands(List<TRexCommand> commands) throws IOException {
    if (commands.size() == 1) {
      return new RPCResponse[] {sendCommand(commands.get(0))};
    }

    List<Map<String, Object>> commandList =
        commands.stream().map(TRexCommand::getParameters).collect(Collectors.toList());

    if (commandList.isEmpty()) {
      return new RPCResponse[0];
    }

    String json = new ObjectMapper().writeValueAsString(commandList);
    String response = sendJson(json);
    return new ObjectMapper().readValue(response, RPCResponse[].class);
  }

  /** @return server host address */
  public String getHost() {
    return host;
  }

  /** @return server port */
  public String getPort() {
    return port;
  }

  /** close zmq connection to server */
  public synchronized void close() {
    zmqSocket.disconnect(connectionString);
    zmqSocket.close();
    zmqCtx.close();
  }

  /** @return zmq socket */
  public ZMQ.Socket getSocket() {
    return zmqSocket;
  }

  /**
   * Send json string to server
   *
   * @param json
   * @return json string
   */
  public synchronized String sendJson(String json) {
    LOGGER.debug("JSON Req: {}", json);

    byte[] compressed = this.dataCompressor.compressStringToBytes(json);

    try {
      zmqSocket.send(compressed);
    } catch (ZMQException e) {
      throwTimeoutException(e);
    }
    byte[] msg = zmqSocket.recv();
    // add this if-block to deal with the connect lost error  which is caused
    // after a long time the client-server don't interactive
    if (msg == null && zmqSocket.errno() == ZError.ENOTCONN) {
      zmqSocket.connect(this.connectionString);
      try {
        zmqSocket.send(compressed);
      } catch (ZMQException e) {
        throwTimeoutException(e);
      }
      msg = zmqSocket.recv();
    }

    String response = this.dataCompressor.decompressBytesToString(msg);
    LOGGER.debug("JSON Resp: {}", response);
    return response;
  }

  private void throwTimeoutException(Exception e) {
    throw new IllegalStateException(
        "Did not get any response from server "
            + getHost()
            + " within timeout "
            + zmqSocket.getReceiveTimeOut(),
        e);
  }
}

@@ -496,27 +496,8 @@
 ion;
-%0Aimport zmq.ZError;
 %0A%0A/*
@@ -4131,39 +4131,49 @@
 e) %7B%0A      throw
-Timeout
+ new IllegalState
 Exception(e);%0A  
@@ -4170,322 +4170,135 @@
 ion(
-e);
 %0A    
-%7D%0A
     
-byte%5B%5D msg = zmqSocket.recv();%0A    // add this if-block to deal with the connect lost error  which is caused%0A    // after a long time the client-server don't interactive%0A    if (msg == null && zmqSocket.errno() == ZError.ENOTCONN) %7B%0A      zmqSocket.connect(this.connectionString);
+  %22Did not get any response from server %22%0A              + getHost()%0A              + %22 within timeout %22
 %0A      
-try %7B%0A
+ 
        
++
  zmq
@@ -4308,97 +4308,47 @@
 ket.
-send(compressed);%0A      %7D catch (ZMQException e) %7B%0A        throwTimeoutException(
+getReceiveTimeOut(),%0A          
 e);%0A    
   %7D%0A
@@ -4347,17 +4347,20 @@
     
-  
 %7D%0A    
- 
+byte%5B%5D
  msg
@@ -4379,22 +4379,16 @@
 .recv();
-%0A    %7D
 %0A%0A    St
@@ -4526,263 +4526,7 @@
 %0A  %7D
-%0A%0A  private void throwTimeoutException(Exception e) %7B%0A    throw new IllegalStateException(%0A        %22Did not get any response from server %22%0A            + getHost()%0A            + %22 within timeout %22%0A            + zmqSocket.getReceiveTimeOut(),%0A        e);%0A  %7D
 %0A%7D%0A

Revert "auto-reconnect mechinism"