package udp

import (
	"net"
	"sync"

	mapset "github.com/deckarep/golang-set/v2"
)

// ... existing code ...

type RoutingTable struct {
	mu     sync.RWMutex
	routes map[string]mapset.Set[*net.UDPAddr]
}

// NewRoutingTable crée une table vide.
func NewRoutingTable() *RoutingTable {
	return &RoutingTable{
		routes: make(map[string]mapset.Set[*net.UDPAddr]),
	}
}

// Add enregistre un client dans un channel.
func (rt *RoutingTable) Add(channel string, addr *net.UDPAddr) {
	rt.mu.Lock()
	defer rt.mu.Unlock()

	set, ok := rt.routes[channel]
	if !ok {
		set = mapset.NewSet[*net.UDPAddr]()
		rt.routes[channel] = set
	}
	set.Add(addr)
}

// Remove supprime un client d'un channel.
func (rt *RoutingTable) Remove(channel string, addr *net.UDPAddr) {
	rt.mu.Lock()
	defer rt.mu.Unlock()

	set, ok := rt.routes[channel]
	if !ok {
		return
	}
	set.Remove(addr)
	if set.Cardinality() == 0 {
		delete(rt.routes, channel)
	}
}

// GetAddrs renvoie une copie de la liste des clients d'un channel.
func (rt *RoutingTable) GetAddrs(channel string) []*net.UDPAddr {
	rt.mu.RLock()
	defer rt.mu.RUnlock()

	set, ok := rt.routes[channel]
	if !ok {
		return nil
	}

	addrs := make([]*net.UDPAddr, 0, set.Cardinality())
	for addr := range set.Iter() {
		addrs = append(addrs, addr)
	}
	return addrs
}