package models

import (
	"sync"

	"github.com/gorilla/websocket"
)

type WebSocketManager struct {
	connections map[*websocket.Conn]string // 连接指针:任务ID
	mutex       sync.RWMutex
}

// 创建一个新的连接池
func NewWebSocketManager() *WebSocketManager {
	return &WebSocketManager{
		connections: make(map[*websocket.Conn]string),
		mutex:       sync.RWMutex{},
	}
}

// 向连接池加入一个新连接
func (mgr *WebSocketManager) AddConnection(conn *websocket.Conn, task string) {
	mgr.mutex.Lock()
	defer mgr.mutex.Unlock()
	mgr.connections[conn] = task
}

// 从连接池中移除一个连接
func (mgr *WebSocketManager) RemoveConnection(conn *websocket.Conn) {
	mgr.mutex.Lock()
	defer mgr.mutex.Unlock()
	delete(mgr.connections, conn)
}

// 任务状态变化时, 向监听此任务的所有连接发送消息
func (mgr *WebSocketManager) NotifyTaskChange(task string, data interface{}) {
	mgr.mutex.Lock()
	defer mgr.mutex.Unlock()

	for conn, value := range mgr.connections {
		if value == task {
			conn.WriteJSON(data)
		}
	}
}