//
// Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights
// Reserved. MIT License  (https://opensource.org/licenses/MIT)
//
/*
 * // 2022-2023 by zhaomingwork@qq.com
 */
// java FunasrWsClient
// usage:  FunasrWsClient [-h] [--port PORT] [--host HOST] [--audio_in AUDIO_IN] [--num_threads NUM_THREADS]
//                 [--chunk_size CHUNK_SIZE] [--chunk_interval CHUNK_INTERVAL] [--mode MODE]
package websocket;

import java.io.*;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.*;
import java.util.Map;
import net.sourceforge.argparse4j.ArgumentParsers;
import net.sourceforge.argparse4j.inf.ArgumentParser;
import net.sourceforge.argparse4j.inf.ArgumentParserException;
import net.sourceforge.argparse4j.inf.Namespace;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.drafts.Draft;
import org.java_websocket.handshake.ServerHandshake;
import org.json.simple.JSONArray;
import org.json.simple.JSONObject;
import org.json.simple.parser.JSONParser;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
/** This example demonstrates how to connect to websocket server. */
public class FunasrWsClient extends WebSocketClient {

  public class RecWavThread extends Thread {
    private FunasrWsClient funasrClient;

    public RecWavThread(FunasrWsClient funasrClient) {
      this.funasrClient = funasrClient;
    }

    public void run() {
      this.funasrClient.recWav();
    }
  }

  private static final Logger logger = LoggerFactory.getLogger(FunasrWsClient.class);

  public FunasrWsClient(URI serverUri, Draft draft) {
    super(serverUri, draft);
  }

  public FunasrWsClient(URI serverURI) {
    
	super(serverURI);
  }

  public FunasrWsClient(URI serverUri, Map<String, String> httpHeaders) {
    super(serverUri, httpHeaders);
  }

  public void getSslContext(String keyfile, String certfile) {
    // TODO
    return;
  }

  // send json at first time
  public void sendJson(
      String mode, String strChunkSize, int chunkInterval, String wavName, boolean isSpeaking,String suffix) {
    try {

      JSONObject obj = new JSONObject();
      obj.put("mode", mode);
      JSONArray array = new JSONArray();
      String[] chunkList = strChunkSize.split(",");
      for (int i = 0; i < chunkList.length; i++) {
        array.add(Integer.valueOf(chunkList[i].trim()));
      }

      obj.put("chunk_size", array);
      obj.put("chunk_interval", new Integer(chunkInterval));
      obj.put("wav_name", wavName);

	  if(FunasrWsClient.hotwords.trim().length()>0)
	  {
		  String regex = "\\d+";
		  JSONObject jsonitems = new JSONObject();
		  String[] items=FunasrWsClient.hotwords.trim().split(" ");
          Pattern pattern = Pattern.compile(regex);
          String tmpWords="";
		  for(int i=0;i<items.length;i++)
		  {

			  Matcher matcher = pattern.matcher(items[i]);

			  if (matcher.matches()) {

				jsonitems.put(tmpWords.trim(), items[i].trim());
				tmpWords="";
				continue;
			  }
			  tmpWords=tmpWords+items[i]+" ";

		  }



		  obj.put("hotwords", jsonitems.toString());
	  }

	  if(suffix.equals("wav")){
	      suffix="pcm";
	  }
	  obj.put("wav_format", suffix);
      if (isSpeaking) {
        obj.put("is_speaking", new Boolean(true));
      } else {
        obj.put("is_speaking", new Boolean(false));
      }
      logger.info("sendJson: " + obj);
      // return;

      send(obj.toString());

      return;
    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  // send json at end of wav
  public void sendEof() {
    try {
      JSONObject obj = new JSONObject();

      obj.put("is_speaking", new Boolean(false));

      logger.info("sendEof: " + obj);
      // return;

      send(obj.toString());
      iseof = true;
      return;
    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  // function for rec wav file
  public void recWav() {
	String fileName=FunasrWsClient.wavPath;
	String suffix=fileName.split("\\.")[fileName.split("\\.").length-1];
    sendJson(mode, strChunkSize, chunkInterval, wavName, true,suffix);
    File file = new File(FunasrWsClient.wavPath);

    int chunkSize = sendChunkSize;
    byte[] bytes = new byte[chunkSize];

    int readSize = 0;
    try (FileInputStream fis = new FileInputStream(file)) {
      if (FunasrWsClient.wavPath.endsWith(".wav")) {
        fis.read(bytes, 0, 44); //skip first 44 wav header
      }
      readSize = fis.read(bytes, 0, chunkSize);
      while (readSize > 0) {
        // send when it is chunk size
        if (readSize == chunkSize) {
          send(bytes); // send buf to server

        } else {
          // send when at last or not is chunk size
          byte[] tmpBytes = new byte[readSize];
          for (int i = 0; i < readSize; i++) {
            tmpBytes[i] = bytes[i];
          }
          send(tmpBytes);
        }
        // if not in offline mode, we simulate online stream by sleep
        if (!mode.equals("offline")) {
          Thread.sleep(Integer.valueOf(chunkSize / 32));
        }

        readSize = fis.read(bytes, 0, chunkSize);
      }

      if (!mode.equals("offline")) {
        // if not offline, we send eof and wait for 3 seconds to close
        Thread.sleep(2000);
        sendEof();
        Thread.sleep(3000);
        close();
      } else {
        // if offline, just send eof
        sendEof();
      }

    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  @Override
  public void onOpen(ServerHandshake handshakedata) {

    RecWavThread thread = new RecWavThread(this);
    thread.start();
  }

  @Override
  public void onMessage(String message) {
    JSONObject jsonObject = new JSONObject();
    JSONParser jsonParser = new JSONParser();
    logger.info("received: " + message);
    try {
      jsonObject = (JSONObject) jsonParser.parse(message);
      logger.info("text: " + jsonObject.get("text"));
	  if(jsonObject.containsKey("timestamp"))
	  {
		  logger.info("timestamp: " + jsonObject.get("timestamp"));
	  }
    } catch (org.json.simple.parser.ParseException e) {
      e.printStackTrace();
    }
    if (iseof && mode.equals("offline") && !jsonObject.containsKey("is_final")) {
      close();
    }
	 
    if (iseof && mode.equals("offline") && jsonObject.containsKey("is_final") && jsonObject.get("is_final").equals("false")) {
      close();
    }
  }

  @Override
  public void onClose(int code, String reason, boolean remote) {

    logger.info(
        "Connection closed by "
            + (remote ? "remote peer" : "us")
            + " Code: "
            + code
            + " Reason: "
            + reason);
  }

  @Override
  public void onError(Exception ex) {
    logger.info("ex: " + ex);
    ex.printStackTrace();
    // if the error is fatal then onClose will be called additionally
  }

  private boolean iseof = false;
  public static String wavPath;
  static String mode = "online";
  static String strChunkSize = "5,10,5";
  static int chunkInterval = 10;
  static int sendChunkSize = 1920;
  static String hotwords="";
  static String fsthotwords="";

  String wavName = "javatest";

  public static void main(String[] args) throws URISyntaxException {
    ArgumentParser parser = ArgumentParsers.newArgumentParser("ws client").defaultHelp(true);
    parser
        .addArgument("--port")
        .help("Port on which to listen.")
        .setDefault("8889")
        .type(String.class)
        .required(false);
    parser
        .addArgument("--host")
        .help("the IP address of server.")
        .setDefault("127.0.0.1")
        .type(String.class)
        .required(false);
    parser
        .addArgument("--audio_in")
        .help("wav path for decoding.")
        .setDefault("asr_example.wav")
        .type(String.class)
        .required(false);
    parser
        .addArgument("--num_threads")
        .help("num of threads for test.")
        .setDefault(1)
        .type(Integer.class)
        .required(false);
    parser
        .addArgument("--chunk_size")
        .help("chunk size for asr.")
        .setDefault("5, 10, 5")
        .type(String.class)
        .required(false);
    parser
        .addArgument("--chunk_interval")
        .help("chunk for asr.")
        .setDefault(10)
        .type(Integer.class)
        .required(false);

    parser
        .addArgument("--mode")
        .help("mode for asr.")
        .setDefault("offline")
        .type(String.class)
        .required(false);
    parser
        .addArgument("--hotwords")
        .help("hotwords, splited by space, hello 30 nihao 40")
        .setDefault("")
        .type(String.class)
        .required(false);
    String srvIp = "";
    String srvPort = "";
    String wavPath = "";
    int numThreads = 1;
    String chunk_size = "";
    int chunk_interval = 10;
    String strmode = "offline";
    String hot="";
    try {
      Namespace ns = parser.parseArgs(args);
      srvIp = ns.get("host");
      srvPort = ns.get("port");
      wavPath = ns.get("audio_in");
      numThreads = ns.get("num_threads");
      chunk_size = ns.get("chunk_size");
      chunk_interval = ns.get("chunk_interval");
      strmode = ns.get("mode");
	  hot=ns.get("hotwords");
      System.out.println(srvPort);

    } catch (ArgumentParserException ex) {
      ex.getParser().handleError(ex);
      return;
    }

    FunasrWsClient.strChunkSize = chunk_size;
    FunasrWsClient.chunkInterval = chunk_interval;
    FunasrWsClient.wavPath = wavPath;
    FunasrWsClient.mode = strmode;
	FunasrWsClient.hotwords=hot;
    System.out.println(
        "serIp="
            + srvIp
            + ",srvPort="
            + srvPort
            + ",wavPath="
            + wavPath
            + ",strChunkSize"
            + strChunkSize);

    class ClientThread implements Runnable {

      String srvIp;
      String srvPort;

      ClientThread(String srvIp, String srvPort, String wavPath) {
        this.srvIp = srvIp;
        this.srvPort = srvPort;
      }

      public void run() {
        try {

          int RATE = 16000;
          String[] chunkList = strChunkSize.split(",");
          int int_chunk_size = 60 * Integer.valueOf(chunkList[1].trim()) / chunkInterval;
          int CHUNK = Integer.valueOf(RATE / 1000 * int_chunk_size);
          int stride =
              Integer.valueOf(
                  60 * Integer.valueOf(chunkList[1].trim()) / chunkInterval / 1000 * 16000 * 2);
          System.out.println("chunk_size:" + String.valueOf(int_chunk_size));
          System.out.println("CHUNK:" + CHUNK);
          System.out.println("stride:" + String.valueOf(stride));
          FunasrWsClient.sendChunkSize = CHUNK * 2;

          String wsAddress = "ws://" + srvIp + ":" + srvPort;

          FunasrWsClient c = new FunasrWsClient(new URI(wsAddress));

          c.connect();

          System.out.println("wsAddress:" + wsAddress);
        } catch (Exception e) {
          e.printStackTrace();
          System.out.println("e:" + e);
        }
      }
    }
    for (int i = 0; i < numThreads; i++) {
      System.out.println("Thread1 is running...");
      Thread t = new Thread(new ClientThread(srvIp, srvPort, wavPath));
      t.start();
    }
  }
}