package udp

import (
	"net"
	"sync"

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

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

func NewRoutingTable() *RoutingTable {
	return &RoutingTable{
		routes: make(map[string]mapset.Set[*net.UDPAddr]),
	}
}

func (rt *RoutingTable) AddClient(channelID string, addr *net.UDPAddr) {
	rt.mu.Lock()
	defer rt.mu.Unlock()

	if rt.routes[channelID] == nil {
		rt.routes[channelID] = mapset.NewSet[*net.UDPAddr]()
	}
	rt.routes[channelID].Add(addr)
}

func (rt *RoutingTable) RemoveClient(channelID string, addr *net.UDPAddr) {
	rt.mu.Lock()
	defer rt.mu.Unlock()

	if clients, exists := rt.routes[channelID]; exists {
		clients.Remove(addr)
		if clients.Cardinality() == 0 {
			delete(rt.routes, channelID)
		}
	}
}

// GetClients returns the clients connected to the given channelID
// don't modify the returned set!
func (rt *RoutingTable) GetClients(channelID string) mapset.Set[*net.UDPAddr] {
	rt.mu.RLock()
	defer rt.mu.RUnlock()

	clients, exists := rt.routes[channelID]
	if !exists {
		return nil
	}

	return clients
}