package ws

import "sync"

// Hub 管理房间及连接。它只负责单进程广播；跨实例广播可在此接入 Redis Pub/Sub。
type Hub struct {
	mu    sync.RWMutex
	rooms map[string]map[*Client]struct{}
}

func NewHub() *Hub {
	return &Hub{rooms: make(map[string]map[*Client]struct{})}
}

func (h *Hub) Register(client *Client) int {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.rooms[client.room] == nil {
		h.rooms[client.room] = make(map[*Client]struct{})
	}
	h.rooms[client.room][client] = struct{}{}
	return len(h.rooms[client.room])
}

func (h *Hub) Unregister(client *Client) int {
	h.mu.Lock()
	defer h.mu.Unlock()
	clients, ok := h.rooms[client.room]
	if !ok {
		return 0
	}
	if _, ok := clients[client]; ok {
		delete(clients, client)
		close(client.send)
	}
	remaining := len(clients)
	if remaining == 0 {
		delete(h.rooms, client.room)
	}
	return remaining
}

func (h *Hub) Broadcast(room string, message OutgoingMessage) {
	h.mu.RLock()
	clients := h.rooms[room]
	stale := make([]*Client, 0)
	for client := range clients {
		select {
		case client.send <- message:
		default:
			stale = append(stale, client)
		}
	}
	h.mu.RUnlock()

	for _, client := range stale {
		h.Unregister(client)
		_ = client.conn.Close()
	}
}

func (h *Hub) Count(room string) int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.rooms[room])
}