package org.testng.remote.strprotocol;

import static org.testng.remote.RemoteTestNG.isVerbose;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.StringReader;
import java.io.UnsupportedEncodingException;

import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.stream.JsonReader;
import com.google.gson.stream.JsonWriter;

public class JsonMessageSender extends BaseMessageSender {

  public JsonMessageSender(String host, int port) {
    super(host, port, false);
  }

  public JsonMessageSender(String host, int port, boolean ack) {
    super(host, port, ack);
  }

  @Override
  public void sendMessage(IMessage message) throws Exception {
    if (m_outStream == null) {
      throw new IllegalStateException("Trying to send a message on a shutdown sender");
    }

    synchronized (m_outStream) {
      p("Sending message " + message);

      BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(m_outStream, "UTF-8"));
      JsonWriter writer = new JsonWriter(bw);
      writeMessage(writer, message);
      bw.newLine();
      bw.flush();

      waitForAck();
    }
  }

  void writeMessage(JsonWriter writer, IMessage message) throws IOException {
    writer.beginObject();

    MessageType messageType = getMessageType(message);
    writer.name("type").value(messageType.getValue());
    writer.name("data");
    Gson gson = new GsonBuilder().create();
    gson.toJson(message, message.getClass(), writer);

    writer.endObject();
  }

  private MessageType getMessageType(IMessage message) {
    return message.getType();
  }

  @Override
  public IMessage receiveMessage() throws Exception {
    if (m_inReader == null) {
      try {
        m_inReader = new BufferedReader(new InputStreamReader(m_inStream, "UTF-8"));
      } catch (UnsupportedEncodingException e) {
        m_inReader = new BufferedReader(new InputStreamReader(m_inStream));
      }
    }
    String msg = m_inReader.readLine();
    p("received message: " + msg);

    IMessage message = null;
    if (msg != null) {
      message = deserializeMessage(msg);
    }

    try {
      sendAck();
    } catch (Exception e) {
      if (isVerbose()) {
        System.out.println("sendAck failed with error: " + e.getMessage());
        e.printStackTrace();
      }
    }
    return message;
  }

  IMessage deserializeMessage(String jsonMsg) throws IOException {
    try (JsonReader reader = new JsonReader(new StringReader(jsonMsg))) {
      reader.beginObject();

      String name = reader.nextName();
      if (!"type".equals(name)) {
        throw new IOException("type node first");
      }

      int msgType = reader.nextInt();

      name = reader.nextName();
      if (!"data".equals(name)) {
        throw new IOException("data node then");
      }

      Gson gson = new GsonBuilder().create();
      IMessage message = null;
      switch (MessageType.fromValue(msgType)) {
      case GENERIC:
        message = gson.fromJson(reader, GenericMessage.class);
        break;
      case SUITE:
        message = gson.fromJson(reader, SuiteMessage.class);
        break;
      case TEST:
        message = gson.fromJson(reader, TestMessage.class);
        break;
      case TEST_RESULT:
        message = gson.fromJson(reader, TestResultMessage.class);
        break;
      default:
        throw new IOException("unknown message type: " + msgType + ", raw json: " + jsonMsg);
      }

      reader.endObject();

      return message;
    }
  }

  private static void p(String msg) {
    if (isVerbose()) {
      System.out.println("[JsonMessageSender] " + msg); //$NON-NLS-1$
    }
  }
}

@@ -1377,30 +1377,24 @@
 e = 
-getM
+m
 essage
+.get
 Type(
-message
 );%0A 
@@ -1604,100 +1604,8 @@
  %7D%0A%0A
-  private MessageType getMessageType(IMessage message) %7B%0A    return message.getType();%0A  %7D%0A%0A
   @O
@@ -2667,24 +2667,209 @@
 r.nextInt();
+%0A      MessageType type = MessageType.fromValue(msgType);%0A      if (type == null) %7B%0A        throw new IOException(%22unknown message type: %22 + msgType + %22, raw json: %22 + jsonMsg);%0A      %7D
 %0A%0A      name
@@ -3077,39 +3077,13 @@
 ch (
-MessageType.fromValue(msgT
+t
 ype)
-)
  %7B%0A 
@@ -3510,37 +3510,41 @@
  IOException(%22un
-known
+supported
  message type: %22

fix potential NPE in switch/case