package server

import (
	"context"
	"sync"

	"net-tunnel/pkg/proto"
)

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

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

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

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

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

// NewConnectionManager 创建新的连接管理器
func NewConnectionManager() *ConnectionManager {
	return &ConnectionManager{
		connections: sync.Map{},
	}
}

// AddConnection 添加新的控制连接
func (m *ConnectionManager) AddConnection(id string, conn *ControlConnection) {
	m.connections.Store(id, conn)
}

// RemoveConnection 移除控制连接
func (m *ConnectionManager) RemoveConnection(id string) {
	m.connections.Delete(id)
}

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