package xyz.wbsite.frame.sse;

import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import xyz.wbsite.frame.utils.MapperUtil;

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

public class Sser {

    private static ConcurrentHashMap<String, SseEmitter> sseMap = new ConcurrentHashMap();

    /**
     * 注册推送服务
     */
    public static SseEmitter register(String key, SseEmitter sseEmitter) {
        if (sseMap.get(key) != null) {
            sseMap.remove(key);
        }
        sseMap.put(key, sseEmitter);
        return sseEmitter;
    }

    /**
     * 消息推送
     *
     * @param key  key
     * @param data 推送数据
     */
    public static void push(String key, Object data) {
        SseEmitter sseEmitter = sseMap.get(key);
        if (sseEmitter == null) {
            return;
        }
        try {
            sseEmitter.send(MapperUtil.toJson(data), MediaType.APPLICATION_JSON);
        } catch (IOException e) {
            sseMap.remove(key);
        }
    }

    /**
     * 消息批量推送
     *
     * @param data 推送数据
     */
    public static void pushAll(Object data) {
        for (String s : sseMap.keySet()) {
            try {
                sseMap.get(s).send(MapperUtil.toJson(data), MediaType.APPLICATION_JSON);
            } catch (IOException e) {
                sseMap.remove(s);
            }
        }
    }

}
