package more

import (
	"sync"

	uuid "github.com/satori/go.uuid"
)

// 同一个key对应多个连接
// 2023年6月29日13:45:25

// 场景：
// 1. 同时支持多个用户在线
// 1. 一个用户对应多个连接
// 2. 在其中一个用户连接中发送消息，同用户其他连接也能收到

type connPool struct {
	pool map[string]map[string]PoolValue
	idx  map[string]string // 索引
	mu   sync.RWMutex
}

type PoolValue interface {
	Close() error
}

// 初始化连接池
func NewConnPool() *connPool {
	return &connPool{
		pool: make(map[string]map[string]PoolValue),
		idx:  make(map[string]string),
		mu:   sync.RWMutex{},
	}
}

// 添加连接, 返回connId
func (c *connPool) Store(key string, value PoolValue) string {
	c.mu.Lock()
	defer c.mu.Unlock()
	val, _ := c.pool[key]

	connId := uuid.NewV4().String()
	val[connId] = value
	c.pool[key] = val
	c.idx[connId] = key
	return connId
}

// 获取一个用户的所有连接
func (c *connPool) LoadByKey(key string) (conns map[string]PoolValue, ok bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	valMap, ok := c.pool[key]

	return valMap, ok
}

// 根据连接ID获取连接
func (c *connPool) LoadByConnId(connId string) (*PoolValue, bool) {
	c.mu.RLock()
	defer c.mu.Unlock()

	key, ok := c.idx[connId]
	if !ok {
		return nil, false
	}

	valMap, ok := c.pool[key]
	if !ok {
		return nil, false
	}

	val, ok := valMap[connId]
	if !ok {
		return nil, false
	}
	return &val, true
}

// 关闭连接
func (c *connPool) Close(connId string) {
	c.mu.Lock()
	defer c.mu.Unlock()

	key, ok := c.idx[connId]
	if !ok {
		return
	}
	valMap, ok := c.pool[key]
	if !ok {
		return
	}
	val, ok := valMap[connId]
	if !ok {
		return
	}
	delete(valMap, connId)

	// 异步关闭
	go val.Close()
}

// 遍历所有的连接
func (c *connPool) Range(f func(key string, connId string, values PoolValue)) {
	c.mu.RLock()
	defer c.mu.Unlock()

	for key, val := range c.pool {
		val := val
		for keyPoo, valPoo := range val {
			go f(key, keyPoo, valPoo)
		}
	}
}