package com.njzscloud.common.ws.support;

import cn.hutool.core.collection.CollUtil;
import lombok.extern.slf4j.Slf4j;

import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.function.Consumer;
import java.util.stream.Collectors;

/**
 * websocket 工具
 */
@Slf4j
public class Websocket {

    public static final String ACTION_PUBLISH = "publish";
    public static final String ACTION_PING = "ping";
    public static final String ACTION_PONG = "pong";
    public static final String ACTION_SUBSCRIBE = "subscribe";
    public static final String ACTION_UNSUBSCRIBE = "unsubscribe";

    /**
     * 客户端事件订阅列表
     */
    static final ConcurrentHashMap<String, List<Wsdoor>> EVENT_CACHE = new ConcurrentHashMap<>();
    static final ReentrantReadWriteLock EVENT_CACHE_LOCK = new ReentrantReadWriteLock();
    static final ReentrantReadWriteLock.ReadLock EVENT_CACHE_READ_LOCK = EVENT_CACHE_LOCK.readLock();
    static final ReentrantReadWriteLock.WriteLock EVENT_CACHE_WRITE_LOCK = EVENT_CACHE_LOCK.writeLock();

    static final ConcurrentHashMap<String, List<Consumer<WsMsg>>> LISTENER_CACHE = new ConcurrentHashMap<>();
    static final ReentrantReadWriteLock LISTENER_CACHE_LOCK = new ReentrantReadWriteLock();
    static final ReentrantReadWriteLock.ReadLock LISTENER_CACHE_READ_LOCK = LISTENER_CACHE_LOCK.readLock();
    static final ReentrantReadWriteLock.WriteLock LISTENER_CACHE_WRITE_LOCK = LISTENER_CACHE_LOCK.writeLock();

    public static boolean publish(WsMsg wsMsg) {
        wsMsg.setAction(ACTION_PUBLISH);
        // 总数，成功数
        Integer[] success = {0, 0};
        String event = wsMsg.getEvent();
        List<WsMsg.User> to = wsMsg.getTo();
        Set<Long> uids;
        Set<String> cids;
        if (CollUtil.isNotEmpty(to)) {
            uids = to.stream().map(WsMsg.User::getUid).collect(Collectors.toSet());
            cids = to.stream().map(WsMsg.User::getCid).collect(Collectors.toSet());
        } else {
            uids = Collections.emptySet();
            cids = Collections.emptySet();
        }
        EVENT_CACHE_READ_LOCK.lock();
        try {
            EVENT_CACHE.computeIfPresent(event, (k, v) -> {
                if (CollUtil.isEmpty(v)) return v;
                List<Wsdoor> wsdoors = v;
                if (!uids.isEmpty() || !cids.isEmpty()) {
                    wsdoors = v.stream().filter(it -> {
                        Long uid = it.getUid();
                        String cid = it.getCid();
                        return cids.contains(cid) || uids.contains(uid);
                    }).collect(Collectors.toList());
                }
                success[1] = wsdoors.size();
                for (Wsdoor wsdoor : wsdoors) {
                    boolean send = wsdoor.send(wsMsg);
                    if (send) success[1] = success[1] + 1;
                }
                return v;
            });
        } finally {
            EVENT_CACHE_READ_LOCK.unlock();
        }

        LISTENER_CACHE_READ_LOCK.lock();
        try {
            LISTENER_CACHE.computeIfPresent(event, (k, v) -> {
                if (CollUtil.isEmpty(v)) return v;
                success[1] = success[1] + v.size();
                for (Consumer<WsMsg> listener : v) {
                    try {
                        listener.accept(wsMsg);
                        success[1] = success[1] + 1;
                    } catch (Exception e) {
                        log.error("事件 {} 监听方法执行异常", event, e);
                    }
                }
                return v;
            });
        } finally {
            LISTENER_CACHE_READ_LOCK.unlock();
        }

        return success[1] > 0;
    }

    public static boolean publish(String event, Object data) {
        return publish(new WsMsg()
            .setAction(ACTION_PUBLISH)
            .setEvent(event)
            .setData(data));
    }

    public static void subscribe(String event, Wsdoor wsdoor) {
        EVENT_CACHE_READ_LOCK.lock();
        try {
            EVENT_CACHE.computeIfAbsent(event, k -> new LinkedList<>());

            EVENT_CACHE.computeIfPresent(event, (k, v) -> {
                v.add(wsdoor);
                return v;
            });
        } finally {
            EVENT_CACHE_READ_LOCK.unlock();
        }
    }

    public static void unsubscribe(Wsdoor wsdoor) {
        EVENT_CACHE_WRITE_LOCK.lock();
        try {
            EVENT_CACHE.forEach((k, v) -> {
                v.removeIf(it -> it.equals(wsdoor));
            });
            EVENT_CACHE.entrySet().removeIf(it -> CollUtil.isEmpty(it.getValue()));
        } finally {
            EVENT_CACHE_WRITE_LOCK.unlock();
        }
    }

    public static void unsubscribe(String event, Wsdoor wsdoor) {
        EVENT_CACHE_WRITE_LOCK.lock();
        try {
            EVENT_CACHE.computeIfPresent(event, (k, v) -> {
                v.removeIf(it -> it.equals(wsdoor));
                return v;
            });
            EVENT_CACHE.entrySet().removeIf(it -> CollUtil.isEmpty(it.getValue()));
        } finally {
            EVENT_CACHE_WRITE_LOCK.unlock();
        }
    }

    /**
     * 系统内部专用订阅, 用于系统事件通知<br/>
     *
     * @param event 事件名称
     */
    public static void subscribe(String event, Consumer<WsMsg> listener) {
        LISTENER_CACHE_READ_LOCK.lock();
        try {
            LISTENER_CACHE.computeIfAbsent(event, k -> new LinkedList<>());
            LISTENER_CACHE.computeIfPresent(event, (k, v) -> {
                v.add(listener);
                return v;
            });
        } finally {
            LISTENER_CACHE_READ_LOCK.unlock();
        }
    }

    public static void unsubscribe(String event, Consumer<WsMsg> listener) {
        LISTENER_CACHE_READ_LOCK.lock();
        try {
            LISTENER_CACHE.computeIfPresent(event, (k, v) -> {
                v.removeIf(it -> it == listener);
                return v;
            });
        } finally {
            LISTENER_CACHE_READ_LOCK.unlock();
        }
    }
}
