package com.centreproject.common.websocket;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import org.springframework.util.AntPathMatcher;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author zhonglei.xu
 */
public class BusinessSocketHandler extends TextWebSocketHandler {
    // Map<业务类型, Map<园区代码, Session>>
    private static final Map<String, Map<String, WebSocketSession>> sessions =
            new ConcurrentHashMap<>();
    private final AntPathMatcher pathMatcher = new AntPathMatcher();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 解析路径参数：/kpi/ws/{businessType}/{parkCode}
        Map<String, String> vars = pathMatcher.extractUriTemplateVariables(
                "/kpi/ws/{businessType}/{parkCode}",
                session.getUri().getPath()
        );

        String businessType = vars.get("businessType");
        String parkCode = vars.get("parkCode");

        // 存储会话
        sessions
                .computeIfAbsent(businessType, k -> new ConcurrentHashMap<>())
                .put(parkCode, session);

        session.sendMessage(new TextMessage(
                String.format("已连接业务通道: %s, 园区: %s", businessType, parkCode)));
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        try {
            JSONObject json = JSON.parseObject(message.getPayload());
            String targetBusiness = json.getString("businessType");
            String targetParkCode = json.getString("parkCode");
            String data = json.getString("data");

            // 定向推送
            Optional.ofNullable(sessions.get(targetBusiness))
                    .map(m -> m.get(targetParkCode))
                    .filter(WebSocketSession::isOpen)
                    .ifPresent(s -> {
                        try {
                            s.sendMessage(new TextMessage(
                                    String.format("[%s-%s] %s", targetBusiness, targetParkCode, data)));
                        } catch (IOException e) {
                            sessions.get(targetBusiness).remove(targetParkCode);
                        }
                    });

        } catch (Exception e) {
            session.sendMessage(new TextMessage("错误: " + e.getMessage()));
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        // 清理失效会话
        sessions.values().forEach(map ->
                map.entrySet().removeIf(entry -> entry.getValue().equals(session))
        );
    }

    // 定向推送方法
    public  void pushToBusiness(String businessType, String parkCode, String message) {
        Optional.ofNullable(sessions.get(businessType))
                .map(m -> m.get(parkCode))
                .filter(WebSocketSession::isOpen)
                .ifPresent(session -> {
                    try {
                        session.sendMessage(new TextMessage(message));
                    } catch (IOException e) {
                        sessions.get(businessType).remove(parkCode);
                    }
                });
    }
}
