package server

import (
	"context"
	"sync"

	"net-tunnel/pkg/proto"
)

// ClientConnection 表示客户端和服务端之间的控制连接
type ClientConnection struct {
	stream proto.TunnelService_ConnectServer
	cancel context.CancelFunc
	ctx    context.Context
}

// NewClientConnection 创建新的控制连接
func NewClientConnection(stream proto.TunnelService_ConnectServer) *ClientConnection {
	ctx, cancel := context.WithCancel(context.Background())
	return &ClientConnection{
		stream: stream,
		ctx:    ctx,
		cancel: cancel,
	}
}

// Send 发送消息到客户端
func (c *ClientConnection) Send(msg *proto.Message) error {
	return c.stream.Send(msg)
}

// Close 关闭连接
func (c *ClientConnection) Close() {
	c.cancel()
}

// ClientManager 管理所有客户端的控制连接
type ClientManager struct {
	clients sync.Map // map[clientID]*ControlConnection
}

// NewClientManager 创建新的连接管理器
func NewClientManager() *ClientManager {
	return &ClientManager{
		clients: sync.Map{},
	}
}

// Add 添加新的控制连接
func (m *ClientManager) Add(id string, conn *ClientConnection) {
	m.clients.Store(id, conn)
}

// Remove 移除控制连接
func (m *ClientManager) Remove(id string) {
	m.clients.Delete(id)
}

// Get 获取指定客户端的控制连接
func (m *ClientManager) Get(id string) (*ClientConnection, bool) {
	conn, ok := m.clients.Load(id)
	if !ok {
		return nil, false
	}
	return conn.(*ClientConnection), true
}