package peer

import (
	"log"
	"sync"
	"sync/atomic"
)

type relayHandler struct {
	lock   sync.Mutex
	relays map[byte]*Remote
	relay  atomic.Pointer[Remote]
}

func newRelayHandler() *relayHandler {
	return &relayHandler{
		relays: make(map[byte]*Remote, 256),
	}
}

func (h *relayHandler) Add(r *Remote) {
	h.lock.Lock()
	defer h.lock.Unlock()

	h.relays[r.RemotePeerIP] = r

	if h.relay.Load() == nil {
		log.Printf("Setting Relay: %v", r.conf().Peer.Name)
		h.relay.Store(r)
	}
}

func (h *relayHandler) Remove(r *Remote) {
	h.lock.Lock()
	defer h.lock.Unlock()

	log.Printf("Removing relay %d...", r.RemotePeerIP)
	delete(h.relays, r.RemotePeerIP)

	if h.relay.Load() == r {
		// Remove current relay.
		h.relay.Store(nil)

		// Find new relay.
		for _, r := range h.relays {
			h.relay.Store(r)
			break
		}
	}
}

func (h *relayHandler) Load() *Remote {
	return h.relay.Load()
}