package com.market.websocket;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketSession;

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

/**
 * 订阅关系管理器
 * <ul>
 *   <li>维护 topic -> sessions 正向映射（推送时使用）</li>
 *   <li>维护 sessionId -> topics 反向映射（断线清理时使用）</li>
 *   <li>缓存每个 topic 的最新数据（req 时返回）</li>
 * </ul>
 */
@Slf4j
@Component
public class SubscriptionManager {

    /** topic -> 订阅该 topic 的所有 session */
    private final ConcurrentHashMap<String, CopyOnWriteArraySet<WebSocketSession>> topicSessions
            = new ConcurrentHashMap<>();

    /** sessionId -> 该 session 订阅的所有 topic */
    private final ConcurrentHashMap<String, Set<String>> sessionTopics
            = new ConcurrentHashMap<>();

    /** topic -> 最新数据（用于 req） */
    private final ConcurrentHashMap<String, Object> latestData
            = new ConcurrentHashMap<>();

    public void subscribe(WebSocketSession session, String topic) {
        topicSessions.computeIfAbsent(topic, k -> new CopyOnWriteArraySet<>()).add(session);
        sessionTopics.computeIfAbsent(session.getId(), k -> ConcurrentHashMap.newKeySet()).add(topic);
    }

    public void unsubscribe(WebSocketSession session, String topic) {
        CopyOnWriteArraySet<WebSocketSession> sessions = topicSessions.get(topic);
        if (sessions != null) sessions.remove(session);
        Set<String> topics = sessionTopics.get(session.getId());
        if (topics != null) topics.remove(topic);
    }

    /** 连接断开时清理该 session 所有订阅 */
    public void removeSession(WebSocketSession session) {
        Set<String> topics = sessionTopics.remove(session.getId());
        if (topics != null) {
            for (String topic : topics) {
                CopyOnWriteArraySet<WebSocketSession> sessions = topicSessions.get(topic);
                if (sessions != null) sessions.remove(session);
            }
        }
    }

    /** 获取订阅该 topic 的所有 session */
    public Set<WebSocketSession> getSessions(String topic) {
        CopyOnWriteArraySet<WebSocketSession> sessions = topicSessions.get(topic);
        return sessions != null ? sessions : Collections.emptySet();
    }

    public void cacheLatest(String topic, Object data) {
        latestData.put(topic, data);
    }

    public Object getLatest(String topic) {
        return latestData.get(topic);
    }

    /** 返回所有 topic 以 {@code prefix} 开头的最新缓存数据 */
    public java.util.Map<String, Object> getLatestsByPrefix(String prefix) {
        java.util.Map<String, Object> result = new java.util.HashMap<>();
        for (java.util.Map.Entry<String, Object> entry : latestData.entrySet()) {
            if (entry.getKey().startsWith(prefix)) {
                result.put(entry.getKey(), entry.getValue());
            }
        }
        return result;
    }
}
