package org.dromara.common.push.core;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.map.MapUtil;
import lombok.extern.slf4j.Slf4j;
import org.dromara.common.core.utils.ThreadUtils;
import org.dromara.common.json.utils.JsonUtils;
import org.dromara.common.push.constant.MessageConstants;
import org.dromara.common.push.dto.PushDTO;
import org.dromara.common.push.properties.MessageProperties;
import org.dromara.common.redis.utils.RedisUtils;
import org.dromara.system.api.domain.PushPayloadDTO;
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.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

/**
 * 管理 Server-Sent Events (SSE) 连接
 *
 * @author Lion Li
 */
@Slf4j
public class SseEmitterSessionManager implements PushSessionManager {

    private final static Map<Long, Map<String, SseEmitter>> USER_TOKEN_EMITTERS = new ConcurrentHashMap<>();

    private final MessageProperties messageProperties;

    /**
     * 构造 SSE 会话管理器并启动心跳检测。
     *
     * @param scheduledExecutorService 定时任务线程池
     * @param messageProperties 消息推送配置
     */
    public SseEmitterSessionManager(ScheduledExecutorService scheduledExecutorService, MessageProperties messageProperties) {
        this.messageProperties = messageProperties;
        // 定时执行 SSE 心跳检测
        scheduledExecutorService.scheduleWithFixedDelay(
            this::sseMonitor,
            messageProperties.getHeartbeatInterval(),
            messageProperties.getHeartbeatInterval(),
            TimeUnit.SECONDS
        );
    }

    /**
     * 建立与指定用户的 SSE 连接
     *
     * @param userId 用户的唯一标识符，用于区分不同用户的连接
     * @param token  用户的唯一令牌，用于识别具体的连接
     * @return 返回一个 SseEmitter 实例，客户端可以通过该实例接收 SSE 事件
     */
    public SseEmitter connect(Long userId, String token) {
        // 从 USER_TOKEN_EMITTERS 中获取或创建当前用户的 SseEmitter 映射表（ConcurrentHashMap）
        // 每个用户可以有多个 SSE 连接，通过 token 进行区分
        Map<String, SseEmitter> emitters = USER_TOKEN_EMITTERS.computeIfAbsent(userId, k -> new ConcurrentHashMap<>());

        // 关闭已存在的SseEmitter，防止超过最大连接数
        SseEmitter oldEmitter = emitters.remove(token);
        if (oldEmitter != null) {
            oldEmitter.complete();
        }

        // 创建一个新的 SseEmitter 实例，避免连接之后直接关闭浏览器导致连接停滞
        SseEmitter emitter = new SseEmitter(messageProperties.getSseTimeout());

        emitters.put(token, emitter);

        // 当 emitter 完成、超时或发生错误时，从映射表中移除对应的 token
        emitter.onCompletion(() -> {
            SseEmitter remove = emitters.remove(token);
            if (remove != null) {
                remove.complete();
            }
        });
        emitter.onTimeout(() -> {
            SseEmitter remove = emitters.remove(token);
            if (remove != null) {
                remove.complete();
            }
        });
        emitter.onError((e) -> {
            SseEmitter remove = emitters.remove(token);
            if (remove != null) {
                remove.complete();
            }
        });

        try {
            // 向客户端发送一条连接成功的事件
            emitter.send(SseEmitter.event().comment("connected"));
        } catch (IOException e) {
            // 如果发送消息失败，则从映射表中移除 emitter
            emitters.remove(token);
        }
        return emitter;
    }

    /**
     * 断开指定用户的 SSE 连接
     *
     * @param userId 用户的唯一标识符，用于区分不同用户的连接
     * @param token  用户的唯一令牌，用于识别具体的连接
     */
    public void disconnect(Long userId, String token) {
        if (userId == null || token == null) {
            return;
        }
        Map<String, SseEmitter> emitters = USER_TOKEN_EMITTERS.get(userId);
        if (MapUtil.isNotEmpty(emitters)) {
            try {
                SseEmitter sseEmitter = emitters.get(token);
                sseEmitter.send(SseEmitter.event().comment("disconnected"));
                sseEmitter.complete();
            } catch (Exception ignore) {
            }
            emitters.remove(token);
        } else {
            USER_TOKEN_EMITTERS.remove(userId);
        }
    }

    /**
     * 执行 SSE 心跳检测并清理失效连接。
     */
    public void sseMonitor() {
        final SseEmitter.SseEventBuilder heartbeat = SseEmitter.event().comment("heartbeat");
        // 记录需要移除的用户ID
        List<Long> toRemoveUsers = new ArrayList<>();

        USER_TOKEN_EMITTERS.forEach((userId, emitterMap) -> {
            if (CollUtil.isEmpty(emitterMap)) {
                toRemoveUsers.add(userId);
                return;
            }

            emitterMap.entrySet().removeIf(entry -> {
                try {
                    entry.getValue().send(heartbeat);
                    return false;
                } catch (Exception ex) {
                    try {
                        entry.getValue().complete();
                    } catch (Exception ignore) {
                        // 忽略重复关闭异常
                    }
                    return true; // 发送失败 → 移除该连接
                }
            });

            // 移除空连接用户
            if (emitterMap.isEmpty()) {
                toRemoveUsers.add(userId);
            }
        });

        // 循环结束后统一清理空用户，避免并发修改异常
        toRemoveUsers.forEach(USER_TOKEN_EMITTERS::remove);
    }

    /**
     * 订阅 SSE 广播主题消息。
     *
     * @param consumer 处理SSE消息的消费者函数
     */
    @Override
    public void subscribeMessage(Consumer<PushDTO> consumer) {
        RedisUtils.subscribe(MessageConstants.MESSAGE_TOPIC, PushDTO.class, consumer);
    }

    /**
     * 向指定用户的全部本地 SSE 会话发送消息。
     *
     * @param userId  要发送消息的用户id
     * @param message 要发送的消息内容
     */
    public void sendMessage(Long userId, String message) {
        Map<String, SseEmitter> emitters = USER_TOKEN_EMITTERS.get(userId);
        if (MapUtil.isNotEmpty(emitters)) {
            for (Map.Entry<String, SseEmitter> entry : emitters.entrySet()) {
                try {
                    entry.getValue().send(SseEmitter.event()
                        .name("message")
                        .data(message));
                } catch (Exception e) {
                    SseEmitter remove = emitters.remove(entry.getKey());
                    if (remove != null) {
                        remove.complete();
                    }
                }
            }
        } else {
            USER_TOKEN_EMITTERS.remove(userId);
        }
    }

    /**
     * 向指定用户的全部本地 SSE 会话发送统一 JSON 消息。
     *
     * @param userId  要发送消息的用户id
     * @param payload 要发送的消息体
     */
    @Override
    public void sendMessage(Long userId, PushPayloadDTO payload) {
        if (payload == null) {
            return;
        }
        sendMessage(userId, JsonUtils.toJsonString(payload));
    }

    /**
     * 向指定用户的全部本地 SSE 会话发送统一 JSON 消息。
     *
     * @param userId        要发送消息的用户id
     * @param pushDTO 要发送的消息内容
     */
    public void sendMessage(Long userId, PushDTO pushDTO) {
        if (pushDTO == null) {
            return;
        }
        sendMessage(userId, pushDTO.getPayload());
    }

    /**
     * 向当前节点所有 SSE 会话发送消息。
     *
     * @param message 要发送的消息内容
     */
    public void sendMessage(String message) {
        List<Long> userIds = new ArrayList<>(USER_TOKEN_EMITTERS.keySet());
        Runnable[] sendTasks = userIds.stream()
            .map(userId -> (Runnable) () -> sendMessage(userId, message))
            .toArray(Runnable[]::new);
        ThreadUtils.virtualInvokeAll(sendTasks);
    }

    /**
     * 向当前节点所有 SSE 会话发送统一 JSON 消息。
     *
     * @param payload 要发送的消息体
     */
    @Override
    public void sendMessage(PushPayloadDTO payload) {
        if (payload == null) {
            return;
        }
        sendMessage(JsonUtils.toJsonString(payload));
    }

    /**
     * 发布 SSE 订阅消息。
     *
     * @param pushDTO 要发布的SSE消息对象
     */
    @Override
    public void publishMessage(PushDTO pushDTO) {
        if (pushDTO == null || pushDTO.getPayload() == null) {
            return;
        }
        RedisUtils.publish(MessageConstants.MESSAGE_TOPIC, pushDTO, consumer -> log.info(
            "发送主题订阅消息topic:{} userIds:{} message:{}",
            MessageConstants.MESSAGE_TOPIC,
            pushDTO.getUserIds(),
            pushDTO.getPayload() == null ? null : pushDTO.getPayload().getMessage()
        ));
    }

    /**
     * 发布 SSE 广播消息。
     *
     * @param message 要发布的消息内容
     */
    public void publishAll(String message) {
        publishAll(PushPayloadDTO.of("message", "backend", message, null));
    }

    /**
     * 发布 SSE 广播 JSON 消息。
     *
     * @param payload 要发布的消息体
     */
    @Override
    public void publishAll(PushPayloadDTO payload) {
        publishMessage(PushDTO.broadcast(payload));
    }
}
