package ws

import (
	"log"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

type WebSocketPool struct {
	conns map[string][]*websocket.Conn
	sync.RWMutex
}

// 创建新的连接池
func NewWebSocketPool() *WebSocketPool {
	return &WebSocketPool{
		conns: make(map[string][]*websocket.Conn),
	}
}

// 添加一个 WebSocket 连接到指定 appID 下
func (p *WebSocketPool) Add(appID string, conn *websocket.Conn) {
	p.Lock()
	defer p.Unlock()
	p.conns[appID] = append(p.conns[appID], conn)
}

// 从连接池中移除一个连接
func (p *WebSocketPool) Remove(appID string, conn *websocket.Conn) {
	p.Lock()
	defer p.Unlock()
	if conns, ok := p.conns[appID]; ok {
		for i, c := range conns {
			if c == conn {
				p.conns[appID] = append(conns[:i], conns[i+1:]...)
				break
			}
		}
		if len(p.conns[appID]) == 0 {
			delete(p.conns, appID)
		}
	}
}

// 向指定 appID 的所有连接广播消息
func (p *WebSocketPool) Broadcast(appID string, message string) {
	p.RLock()
	defer p.RUnlock()
	if conns, ok := p.conns[appID]; ok {
		for _, conn := range conns {
			if err := conn.WriteMessage(websocket.TextMessage, []byte(message)); err != nil {
				log.Printf("[WebSocketPool] failed to send to appID=%s: %v", appID, err)
			}
		}
	}
}

// 清理掉线或不可用的连接（可定期调用）
func (p *WebSocketPool) CleanInactive(appID string) {
	p.Lock()
	defer p.Unlock()
	if conns, ok := p.conns[appID]; ok {
		active := make([]*websocket.Conn, 0, len(conns))
		for _, conn := range conns {
			if err := conn.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(1*time.Second)); err == nil {
				active = append(active, conn)
			} else {
				_ = conn.Close()
			}
		}
		if len(active) > 0 {
			p.conns[appID] = active
		} else {
			delete(p.conns, appID)
		}
	}
}