package cn.reghao.bnt.common.ws;

import cn.reghao.jutil.jdk.serializer.JdkSerializer;
import cn.reghao.jutil.jdk.thread.ThreadPoolWrapper;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import okio.ByteString;

import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * @author reghao
 * @date 2023-02-23 09:26:50
 */
public class WsClient {
    private final ScheduledExecutorService scheduler = ThreadPoolWrapper.scheduledThreadPool("ws-heartbeat", 1);
    private Future<?> heartbeatFuture;

    private final String url;
    private String token;
    private WebSocket webSocket;
    private boolean connected;
    private final WebSocketListener webSocketListener;
    private boolean retry;
    private int retryCount;

    public WsClient(String url) {
        this.url = url;
        //this.token = token;
        this.webSocketListener = new WebSocketListenerImpl(this);
        this.retry = true;
        this.retryCount = 0;
    }

    public void setRetry(boolean retry) {
        this.retry = retry;
    }

    public boolean isRetry() {
        return retry;
    }

    public void retryCountIncr() {
        this.retryCount += 1;
    }

    public void resetRetryCount() {
        this.retryCount = 0;
    }

    public int getRetryCount() {
        return retryCount;
    }

    public void connect() {
        Request request = new Request.Builder()
                .url(url)
                //.header("Authorization", "Bearer " + token)
                .build();

        OkHttpClient okHttpClient = new OkHttpClient.Builder()
                .connectTimeout(30, TimeUnit.SECONDS)
                .readTimeout(30, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();
        this.webSocket = okHttpClient.newWebSocket(request, webSocketListener);
    }

    public void setConnected(boolean status) {
        this.connected = status;
        if (status) {
            heartbeatFuture = scheduler.scheduleAtFixedRate(new HeartbeatTask(), 0, 10, TimeUnit.SECONDS);
        } else {
            if (heartbeatFuture != null && !heartbeatFuture.isCancelled()) {
                heartbeatFuture.cancel(true);
            }
        }
    }

    public boolean isConnected() {
        return connected;
    }

    public void send(Object message) {
        if (isConnected()) {
            byte[] bytes = JdkSerializer.serialize(message);
            webSocket.send(ByteString.of(bytes));
        }
    }

    public void close() {
        setRetry(false);
        webSocket.close(1000, "Client Close Connection");
    }

    class HeartbeatTask implements Runnable {
        @Override
        public void run() {
            send("");
        }
    }
}
