package hub

import (
	"sync"

	"github.com/gorilla/websocket"
)

// Connection 包装单个 WebSocket 连接及其异步发送通道
type Connection struct {
	UserID string
	conn   *websocket.Conn
	send   chan []byte
}

// Send 将消息投递到该连接的发送通道（非阻塞，通道满则丢弃）
func (c *Connection) Send(message []byte) {
	select {
	case c.send <- message:
	default:
		// 通道满，丢弃消息避免阻塞 hub
	}
}

// Outgoing 返回该连接的发送通道，供写协程消费
func (c *Connection) Outgoing() <-chan []byte {
	return c.send
}

// Hub 管理所有在线 WebSocket 连接，按 userID 索引
type Hub struct {
	mu      sync.RWMutex
	clients map[string]map[*Connection]bool
}

// NewHub 创建 Hub 实例
func NewHub() *Hub {
	return &Hub{
		clients: make(map[string]map[*Connection]bool),
	}
}

// ClientCount 返回当前在线连接总数（供 /readyz 探针使用）
func (h *Hub) ClientCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	count := 0
	for _, conns := range h.clients {
		count += len(conns)
	}
	return count
}

// Register 注册一个用户连接，返回 Connection 供调用方持有
func (h *Hub) Register(userID string, conn *websocket.Conn) *Connection {
	c := &Connection{
		UserID: userID,
		conn:   conn,
		send:   make(chan []byte, 64),
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.clients[userID] == nil {
		h.clients[userID] = make(map[*Connection]bool)
	}
	h.clients[userID][c] = true
	return c
}

// Unregister 注销一个用户连接并关闭其发送通道
func (h *Hub) Unregister(c *Connection) {
	h.mu.Lock()
	defer h.mu.Unlock()
	conns, ok := h.clients[c.UserID]
	if !ok {
		return
	}
	if _, exists := conns[c]; exists {
		delete(conns, c)
		close(c.send)
	}
	if len(conns) == 0 {
		delete(h.clients, c.UserID)
	}
}

// SendToUser 向指定用户的所有在线连接推送消息
func (h *Hub) SendToUser(userID string, message []byte) error {
	h.mu.RLock()
	defer h.mu.RUnlock()
	conns, ok := h.clients[userID]
	if !ok {
		return nil
	}
	for c := range conns {
		c.Send(message)
	}
	return nil
}

// Broadcast 向所有在线连接广播消息
func (h *Hub) Broadcast(message []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, conns := range h.clients {
		for c := range conns {
			c.Send(message)
		}
	}
}