package core

import (
	"github.com/gorilla/websocket"
)

// WebSocketConnManager 管理WebSocket连接
type WebSocketConnManager struct {
	ConnMap map[string]*websocket.Conn
}

// 创建一个新的WebSocketConnManager
func NewWebSocketConnManager() *WebSocketConnManager {
	return &WebSocketConnManager{
		ConnMap: make(map[string]*websocket.Conn),
	}
}

// 添加连接
func (ws *WebSocketConnManager) AddConn(id string, conn *websocket.Conn) {
	ws.ConnMap[id] = conn
}

// 根据ID移除连接
func (ws *WebSocketConnManager) RemoveConn(id string) {
	delete(ws.ConnMap, id)
}

// 根据ID获取连接
func (ws *WebSocketConnManager) GetConn(id string) (*websocket.Conn, bool) {
	conn, ok := ws.ConnMap[id]
	return conn, ok
}

// 向所有连接发送消息
func (ws *WebSocketConnManager) SendMessageToAll(message []byte) {
	for _, conn := range ws.ConnMap {
		err := conn.WriteMessage(websocket.TextMessage, message)
		if err != nil {
			continue
		}
	}
}

// 向特定用户发送消息
func (ws *WebSocketConnManager) SendMessageToUser(id string, message []byte) {
	conn, ok := ws.GetConn(id)
	if !ok {
		return
	}
	err := conn.WriteMessage(websocket.TextMessage, message)
	if err != nil {
		return
	}
}

// 向多个用户发送消息
func (ws *WebSocketConnManager) SendMessageToUsers(ids []string, message []byte) {
	for _, id := range ids {
		conn, ok := ws.GetConn(id)
		if !ok {
			continue
		}
		err := conn.WriteMessage(websocket.TextMessage, message)
		if err != nil {
			continue
		}
	}
}