package websocket

import "sync"

type Hub struct {
	Clients    map[string]*Client
	Register   chan *Client
	Unregister chan *Client
	mu         sync.RWMutex
}

func NewHub() *Hub {
	return &Hub{
		Clients:    make(map[string]*Client),
		Register:   make(chan *Client),
		Unregister: make(chan *Client),
	}
}

func (h *Hub) Run() {
	for {
		select {
		case client := <-h.Register:
			h.mu.Lock()
			h.Clients[client.ID] = client
			h.mu.Unlock()
		case client := <-h.Unregister:
			h.mu.Lock()
			if _, ok := h.Clients[client.ID]; ok {
				delete(h.Clients, client.ID)
				close(client.Send)
			}
			h.mu.Unlock()
		}
	}
}

func (h *Hub) SendToClient(id string, msg []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	if c, ok := h.Clients[id]; ok {
		c.Send <- msg
	}
}

func (h *Hub) SendToAll(ids []string, msg []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	for _, id := range ids {
		if c, ok := h.Clients[id]; ok {
			c.Send <- msg
		}
	}
}

func (h *Hub) Broadcast(msg []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, client := range h.Clients {
		client.Send <- msg
	}
}