package com.zksy.web.websocket;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket 订阅管理器
 * <p>
 * 维护 WebSocket 会话与设备编码的双向映射关系，支持：
 * <ul>
 *   <li>sessionId → Set&lt;deviceCode&gt; 正向映射</li>
 *   <li>deviceCode → Set&lt;sessionId&gt; 反向映射（快速查找需要推送的会话）</li>
 * </ul>
 * <p>
 * 线程安全：使用 ConcurrentHashMap + ConcurrentHashMap.newKeySet() 保证并发安全。
 *
 * @author zksy
 */
@Slf4j
@Component
public class SubscriptionManager {

    /**
     * 正向映射：会话ID → 订阅的设备编码集合
     */
    private final ConcurrentHashMap<String, Set<String>> sessionToDevices = new ConcurrentHashMap<>();

    /**
     * 反向映射：设备编码 → 订阅了该设备的会话ID集合
     */
    private final ConcurrentHashMap<String, Set<String>> deviceToSessions = new ConcurrentHashMap<>();

    /**
     * 批量订阅设备
     *
     * @param sessionId   WebSocket 会话ID
     * @param deviceCodes 要订阅的设备编码集合
     * @return 当前会话订阅的所有设备编码
     */
    public Set<String> subscribe(String sessionId, Set<String> deviceCodes) {
        if (deviceCodes == null || deviceCodes.isEmpty()) {
            return getSubscribedDevices(sessionId);
        }

        // 更新正向映射
        sessionToDevices.compute(sessionId, (k, existing) -> {
            if (existing == null) {
                existing = ConcurrentHashMap.newKeySet();
            }
            existing.addAll(deviceCodes);
            return existing;
        });

        // 更新反向映射
        for (String deviceCode : deviceCodes) {
            deviceToSessions.compute(deviceCode, (k, existing) -> {
                if (existing == null) {
                    existing = ConcurrentHashMap.newKeySet();
                }
                existing.add(sessionId);
                return existing;
            });
        }

        log.info("订阅成功: sessionId={}, deviceCodes={}, 当前订阅总数={}",
                sessionId, deviceCodes, getSubscribedDevices(sessionId).size());
        return getSubscribedDevices(sessionId);
    }

    /**
     * 取消订阅设备
     *
     * @param sessionId   WebSocket 会话ID
     * @param deviceCodes 要取消订阅的设备编码集合
     * @return 当前会话剩余的订阅设备编码
     */
    public Set<String> unsubscribe(String sessionId, Set<String> deviceCodes) {
        if (deviceCodes == null || deviceCodes.isEmpty()) {
            return getSubscribedDevices(sessionId);
        }

        // 更新正向映射
        sessionToDevices.computeIfPresent(sessionId, (k, existing) -> {
            existing.removeAll(deviceCodes);
            return existing.isEmpty() ? null : existing;
        });

        // 更新反向映射
        for (String deviceCode : deviceCodes) {
            deviceToSessions.computeIfPresent(deviceCode, (k, existing) -> {
                existing.remove(sessionId);
                return existing.isEmpty() ? null : existing;
            });
        }

        log.info("取消订阅: sessionId={}, deviceCodes={}, 当前订阅总数={}",
                sessionId, deviceCodes, getSubscribedDevices(sessionId).size());
        return getSubscribedDevices(sessionId);
    }

    /**
     * 移除会话的所有订阅（连接断开时调用）
     *
     * @param sessionId WebSocket 会话ID
     */
    public void removeSession(String sessionId) {
        Set<String> deviceCodes = sessionToDevices.remove(sessionId);
        if (deviceCodes == null || deviceCodes.isEmpty()) {
            return;
        }

        for (String deviceCode : deviceCodes) {
            deviceToSessions.computeIfPresent(deviceCode, (k, existing) -> {
                existing.remove(sessionId);
                return existing.isEmpty() ? null : existing;
            });
        }

        log.info("会话移除，清理订阅: sessionId={}, 清理设备数={}", sessionId, deviceCodes.size());
    }

    /**
     * 获取订阅了指定设备的所有会话ID
     *
     * @param deviceCode 设备编码
     * @return 会话ID集合（不可变），未找到返回空集合
     */
    public Set<String> getSubscribedSessions(String deviceCode) {
        Set<String> sessions = deviceToSessions.get(deviceCode);
        return sessions != null ? Collections.unmodifiableSet(sessions) : Collections.emptySet();
    }

    /**
     * 获取指定会话订阅的所有设备编码
     *
     * @param sessionId WebSocket 会话ID
     * @return 设备编码集合（不可变），未找到返回空集合
     */
    public Set<String> getSubscribedDevices(String sessionId) {
        Set<String> devices = sessionToDevices.get(sessionId);
        return devices != null ? Collections.unmodifiableSet(devices) : Collections.emptySet();
    }

    /**
     * 获取当前订阅统计信息
     */
    public String getStats() {
        return String.format("活跃会话数=%d, 被订阅设备数=%d",
                sessionToDevices.size(), deviceToSessions.size());
    }
}