package org.springblade.modules.sse.server;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

@Slf4j
public class SSEServer {

	/**
	 * 当前连接数
	 */
	private static AtomicInteger count = new AtomicInteger(0);

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

	public static SseEmitter connect(String userId){
		//设置超时时间，0表示不过期，默认是30秒，超过时间未完成会抛出异常
		SseEmitter sseEmitter = new SseEmitter(0L);
		//注册回调
		sseEmitter.onCompletion(completionCallBack(userId));
		sseEmitter.onError(errorCallBack(userId));
		sseEmitter.onTimeout(timeOutCallBack(userId));
		sseEmitterMap.put(userId,sseEmitter);
		//数量+1
		count.getAndIncrement();
		log.info("create new sse connect ,current user:{}",userId);
		log.info("count",count.getAndIncrement());
		return sseEmitter;
	}
	/**
	 * 给指定用户发消息
	 */
	public static void sendMessage(String userId, String message){
		if(sseEmitterMap.containsKey(userId)){
			try{
				sseEmitterMap.get(userId).send(message);
			}catch (IOException e){
				log.error("user id:{}, send message error:{}",userId,e.getMessage());
				e.printStackTrace();
			}
		}
	}

	/**
	 * 想多人发送消息，组播
	 */
	public static void groupSendMessage(String groupId, String message){
		if(sseEmitterMap!=null&&!sseEmitterMap.isEmpty()){
			sseEmitterMap.forEach((k,v) -> {
				try{
					if(k.startsWith(groupId)){
						v.send(message, MediaType.APPLICATION_JSON);
					}
				}catch (IOException e){
					log.error("user id:{}, send message error:{}",groupId,message);
					removeUser(k);
				}
			});
		}
	}

	/**
	 * 批量发送消息
	 * @param message
	 */
	public static void batchSendMessage(String message) {
		sseEmitterMap.forEach((k,v)->{
			try{
				v.send(message,MediaType.APPLICATION_JSON);
			}catch (IOException e){
				log.error("user id:{}, send message error:{}",k,e.getMessage());
				removeUser(k);
			}
		});
	}

	/**
	 * 群发消息
	 */
	public static void batchSendMessage(String message, Set<String> userIds){
		userIds.forEach(userId->sendMessage(userId,message));
	}

	/**
	 * 用户离线删除用户
	 * @param userId
	 */
	public static void removeUser(String userId){
		sseEmitterMap.remove(userId);
		//数量-1
		count.getAndDecrement();
		log.info("remove user id:{}",userId);
	}


	public static List<String> getIds(){
		return new ArrayList<>(sseEmitterMap.keySet());
	}

	public static int getUserCount(){
		return count.intValue();
	}

	/**
	 * 结束回调
	 * @param userId
	 * @return
	 */
	private static Runnable completionCallBack(String userId) {
		return () -> {
			log.info("结束连接,{}",userId);
			removeUser(userId);
		};
	}

	/**
	 * 超时回调
	 * @param userId
	 * @return
	 */
	private static Runnable timeOutCallBack(String userId){
		return ()->{
			log.info("连接超时,{}",userId);
			removeUser(userId);
		};
	}

	/**
	 * 错误回调
	 * @param userId
	 * @return
	 */
	private static Consumer<Throwable> errorCallBack(String userId){
		return throwable -> {
			log.error("连接异常，{}",userId);
			removeUser(userId);
		};
	}
}