package transport

import (
	"net"
	"sync"

	"go.uber.org/zap"
)

// ConnManager 连接管理器
// 维护 peer_key → net.Conn 的映射关系
type ConnManager struct {
	conns  map[string]net.Conn
	mu     sync.RWMutex
	logger *zap.Logger
}

// NewConnManager 创建连接管理器
func NewConnManager(logger *zap.Logger) *ConnManager {
	return &ConnManager{
		conns:  make(map[string]net.Conn),
		logger: logger,
	}
}

// Add 添加连接
func (m *ConnManager) Add(peerKey string, conn net.Conn) {
	m.mu.Lock()
	defer m.mu.Unlock()

	// 如果已存在，先关闭旧连接
	if oldConn, ok := m.conns[peerKey]; ok {
		oldConn.Close()
		m.logger.Debug("关闭旧连接", zap.String("peer_key", peerKey))
	}

	m.conns[peerKey] = conn
	m.logger.Info("添加新连接",
		zap.String("peer_key", peerKey),
		zap.String("remote_addr", conn.RemoteAddr().String()))
}

// Get 获取连接
func (m *ConnManager) Get(peerKey string) (net.Conn, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	conn, ok := m.conns[peerKey]
	return conn, ok
}

// Remove 移除连接
func (m *ConnManager) Remove(peerKey string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if conn, ok := m.conns[peerKey]; ok {
		conn.Close()
		delete(m.conns, peerKey)
		m.logger.Info("移除连接", zap.String("peer_key", peerKey))
	}
}

// Count 获取连接数量
func (m *ConnManager) Count() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.conns)
}

// CloseAll 关闭所有连接
func (m *ConnManager) CloseAll() {
	m.mu.Lock()
	defer m.mu.Unlock()

	for peerKey, conn := range m.conns {
		conn.Close()
		m.logger.Debug("关闭连接", zap.String("peer_key", peerKey))
	}

	m.conns = make(map[string]net.Conn)
}

// List 列出所有连接（返回副本）
func (m *ConnManager) List() map[string]net.Conn {
	m.mu.RLock()
	defer m.mu.RUnlock()

	result := make(map[string]net.Conn)
	for k, v := range m.conns {
		result[k] = v
	}
	return result
}

// GetAll 获取所有连接（同 List，为了兼容）
func (m *ConnManager) GetAll() map[string]net.Conn {
	return m.List()
}