package ws

import cmap "github.com/orcaman/concurrent-map/v2"

var wsMgr = newManager()

type wsManager struct {
	wsConnAll cmap.ConcurrentMap[uint32, *wsConnect]
}

func newManager() *wsManager {
	return &wsManager{
		wsConnAll: cmap.NewWithCustomShardingFunction[uint32, *wsConnect](func(key uint32) uint32 {
			return key
		}),
	}
}

func (m *wsManager) Add(conn *wsConnect) {
	m.wsConnAll.Set(conn.id, conn)
}

func (m *wsManager) SetUserId(connId uint32, userId int64) {
	userMgr.Add(connId, userId)
}

func (m *wsManager) Remove(conn *wsConnect) {
	if conn.UserId() > 0 {
		userMgr.Remove(conn.UserId())
	}
	m.wsConnAll.Remove(conn.id)
}

func (m *wsManager) Get(connId uint32) (*wsConnect, bool) {
	return m.wsConnAll.Get(connId)
}

func (m *wsManager) FindByUserId(userId int64) (*wsConnect, bool) {
	connId := userMgr.GetConnId(userId)
	return m.wsConnAll.Get(connId)
}

func (m *wsManager) Count() int {
	return m.wsConnAll.Count()
}