package models

import (
	"sync"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

type WebSocketManager struct {
	connections map[string]*websocket.Conn
	listeners   map[string]map[chan struct{}]struct{}
	mutex       sync.RWMutex
}

func NewWebSocketManager() *WebSocketManager {
	return &WebSocketManager{
		connections: make(map[string]*websocket.Conn),
		mutex:       sync.RWMutex{},
	}
}

func (mgr *WebSocketManager) AddConnection(conn *websocket.Conn) string {
	mgr.mutex.Lock()
	defer mgr.mutex.Unlock()

	id := uuid.New().String() // 为每个连接生成一个唯一的 ID
	mgr.connections[id] = conn

	return id
}

func (mgr *WebSocketManager) RemoveConnection(id string) {
	mgr.mutex.Lock()
	defer mgr.mutex.Unlock()
	delete(mgr.connections, id)
}

func (mgr *WebSocketManager) ListenForChanges(target string, callback func()) {
	notifications := make(chan struct{})
	mgr.mutex.Lock()
	defer mgr.mutex.Unlock()

	if _, ok := mgr.listeners[target]; !ok {
		mgr.listeners[target] = make(map[chan struct{}]struct{})
	}
	mgr.listeners[target][notifications] = struct{}{}

	go func() {
		for {
			callback()
			for listener := range mgr.listeners[target] {
				select {
				case listener <- struct{}{}:
				default:
					delete(mgr.listeners[target], listener)
				}
			}
		}
	}()
}