package cn.reghao.tnb.search.esapp.ws.handler;

import cn.reghao.jutil.jdk.converter.DateTimeConverter;
import cn.reghao.jutil.jdk.web.log.AppLog;
import cn.reghao.jutil.jdk.serializer.JdkSerializer;
import cn.reghao.jutil.jdk.serializer.JsonConverter;
import cn.reghao.jutil.jdk.web.log.GatewayLog;
import cn.reghao.jutil.jdk.web.log.NginxLog;
import cn.reghao.tnb.search.esapp.log.mq.KafkaPub;
import cn.reghao.tnb.search.esapp.ws.config.WebSocketPath;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author reghao
 * @date 2022-03-11 16:45:52
 */
@Slf4j
@Component
public class LogHandler implements WebSocketHandler {
    private final Map<String, String> appMap = new HashMap<>();
    private final Map<String, WebSocketSession> pullRuntimeMap = new ConcurrentHashMap<>();
    private final Map<String, WebSocketSession> pullAccessMap = new ConcurrentHashMap<>();
    private final Map<String, WebSocketSession> pushSessionMap = new ConcurrentHashMap<>();
    private final Map<String, WebSocketSession> pullSessionMap = new ConcurrentHashMap<>();
    private final KafkaPub kafkaPub;

    public LogHandler(KafkaPub kafkaPub) {
        this.kafkaPub = kafkaPub;
    }

    private WebSocketSession getPullSession(String app, String host) {
        String suffix = String.format("%s@%s", app, host);
        Set<String> keys = pullSessionMap.keySet();
        for (String key : keys) {
            if (key.endsWith(suffix)) {
                return pullSessionMap.get(key);
            }
        }

        return null;
    }

    private WebSocketSession getPullRuntimeSession(String app, String host) {
        String suffix = String.format("%s@%s", app, host);
        Set<String> keys = pullRuntimeMap.keySet();
        for (String key : keys) {
            if (key.endsWith(suffix)) {
                return pullRuntimeMap.get(key);
            }
        }

        return null;
    }

    private WebSocketSession getPullAccessSession(String app, String host) {
        String suffix = String.format("%s@%s", app, host);
        Set<String> keys = pullAccessMap.keySet();
        for (String key : keys) {
            if (key.endsWith(suffix)) {
                return pullAccessMap.get(key);
            }
        }

        return null;
    }

    private Map<String, String> parseParams(String query) {
        String[] params = query.split("&");
        Map<String, String> map = new HashMap<>();
        for (String param : params) {
            String[] arr = param.split("=");
            map.put(arr[0], arr[1]);
        }
        return map;
    }

    public void pushGatewayLog(GatewayLog gatewayLog) {
        WebSocketSession pullSession = getPullAccessSession("tnb", "127.0.0.1");
        if (pullSession != null) {
            String jsonData = JsonConverter.objectToJson(gatewayLog);
            WebSocketMessage<String> message1 = new TextMessage(jsonData);
            try {
                pullSession.sendMessage(message1);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    public void pushAppLog(AppLog appLog) {
        String app = appLog.getApp();
        String host = appLog.getHost();
        WebSocketSession pullSession = getPullRuntimeSession(app, host);
        if (pullSession != null) {
            String jsonData = JsonConverter.objectToJson(appLog);
            WebSocketMessage<String> message1 = new TextMessage(jsonData);
            try {
                pullSession.sendMessage(message1);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws IOException {
        String sessionId = webSocketSession.getId();
        String query = webSocketSession.getUri().getQuery();
        Map<String, String> map = parseParams(query);
        String app = map.get("app");
        String host = map.get("host");
        if (app == null || host == null) {
            if (webSocketSession.isOpen()) {
                webSocketSession.close();
            }
        }

        String appKey = String.format("%s@%s", app, host);
        String path = webSocketSession.getUri().getPath();
        if (path.equals(WebSocketPath.wsAgent)) {
            appMap.put(sessionId, appKey);
            pushSessionMap.put(appKey, webSocketSession);
        } else if (path.equals(WebSocketPath.wsPullLog)) {
            pullSessionMap.put(appKey, webSocketSession);
        }

        log.info("WebSocket 建立连接");
    }

    @Override
    public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> webSocketMessage)
            throws IOException {
        try {
            String appKey = appMap.get(webSocketSession.getId());
            if (webSocketMessage instanceof TextMessage) {
                String payload = (String) webSocketMessage.getPayload();
            } else if (webSocketMessage instanceof BinaryMessage) {
                //log.info("接收到 WebSocket 二进制消息");
                BinaryMessage binaryMessage = (BinaryMessage) webSocketMessage;
                Object object = JdkSerializer.deserialize(binaryMessage.getPayload().array());
                if (object instanceof String) {
                    //log.info("{} -> {}", appKey, object);
                    log.info("WebSocket heartbeat");
                    webSocketSession.sendMessage(new TextMessage(""));
                } else if (object instanceof AppLog) {
                    AppLog appLog = (AppLog) object;
                    String dateTimeStr = DateTimeConverter.format(appLog.getTimestamp());
                } else if (object instanceof NginxLog) {
                    NginxLog nginxLog = (NginxLog) object;
                    //logProducer.put(nginxLog);
                    kafkaPub.produce(nginxLog);
                }
            } else if (webSocketMessage instanceof PingMessage) {
                log.info("接收到 WebSocket PingMessage");
                webSocketSession.sendMessage(new PongMessage());
            } else if (webSocketMessage instanceof PongMessage) {
                log.info("接收到 WebSocket PongMessage");
            } else {
                log.error("接收到未知类型的 WebSocket 消息");
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Override
    public void handleTransportError(WebSocketSession webSocketSession, Throwable throwable) {
        log.error("WebSocket 数据传输错误: {}", throwable.getMessage());
        removeSession(webSocketSession);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus closeStatus) {
        log.info("WebSocket 断开连接: {} -> {}", closeStatus.getCode(), closeStatus.getReason());
        removeSession(webSocketSession);
    }

    private void removeSession(WebSocketSession webSocketSession) {
        String sessionId = webSocketSession.getId();
        String appKey = appMap.get(sessionId);
        if (appKey == null) {
            return;
        }

        String path = webSocketSession.getUri().getPath();
        if (path.equals("/ws/log/push")) {
            pushSessionMap.remove(appKey);
        } else if (path.equals("/ws/log/pull")) {
            pullSessionMap.remove(appKey);
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}
