package net.geedge.api.config;

import cn.hutool.log.Log;
import net.geedge.api.entity.EnvApiYml;
import net.geedge.common.T;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.Socket;

public class VncProxyHandler extends TextWebSocketHandler {

    private static final Log log = Log.get();

    private EnvApiYml.Vnc vnc;

    public VncProxyHandler(EnvApiYml.Vnc vnc) {
        this.vnc = vnc;
    }

    @Override
    public synchronized void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("[afterConnectionEstablished] [WebSocket connection established] [websocket uri: {}]", session.getUri());
        super.afterConnectionEstablished(session);

        // connect to VNC Server
        Socket vncSocket = new Socket(vnc.getHost(), vnc.getPort());
        session.getAttributes().put("vncSocket", vncSocket);
        log.info("[afterConnectionEstablished] [vnc server: {}] [isConnected: {}]", T.JSONUtil.toJsonStr(vnc), vncSocket.isConnected());

        // vnc server -> web
        T.ThreadUtil.execute(() -> {
            this.forwardFromVncToWeb(session, vncSocket);
        });
    }

    /**
     * vnc server -> web
     *
     * @param session
     * @param vncSocket
     */
    private void forwardFromVncToWeb(WebSocketSession session, Socket vncSocket) {
        try (InputStream inputStream = vncSocket.getInputStream()) {
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                session.sendMessage(new BinaryMessage(buffer, 0, bytesRead, true));
            }
        } catch (IOException e) {
            log.error(e, "[forwardFromVncToWeb] [error]");
        }
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) {
        try {
            // web -> vnc server
            Socket vncSocket = (Socket) session.getAttributes().get("vncSocket");
            if (vncSocket != null && !vncSocket.isClosed()) {
                OutputStream outputStream = vncSocket.getOutputStream();
                outputStream.write(message.getPayload().array());
                outputStream.flush();
            }
        } catch (IOException e) {
            log.error(e, "[handleBinaryMessage] [error]");
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("[afterConnectionClosed] [WebSocket connection closed] [websocket uri: {}]", session.getUri());
        Socket vncSocket = (Socket) session.getAttributes().get("vncSocket");
        if (vncSocket != null && !vncSocket.isClosed()) {
            vncSocket.close();
        }
        super.afterConnectionClosed(session, status);
    }
}