package peer

import (
	"bytes"
	"net"
	"net/netip"
	"sync"
)

type TestPacket struct {
	Addr netip.AddrPort
	Data []byte
}

type TestNetwork struct {
	lock    sync.Mutex
	packets map[netip.AddrPort]chan TestPacket
}

func NewTestNetwork() *TestNetwork {
	return &TestNetwork{packets: map[netip.AddrPort]chan TestPacket{}}
}

func (n *TestNetwork) NewUDPConn(localAddr netip.AddrPort) *TestUDPConn {
	n.lock.Lock()
	defer n.lock.Unlock()
	if _, ok := n.packets[localAddr]; !ok {
		n.packets[localAddr] = make(chan TestPacket, 1024)
	}
	return &TestUDPConn{
		addr:    localAddr,
		n:       n,
		packets: n.packets[localAddr],
	}
}

func (n *TestNetwork) write(b []byte, from, to netip.AddrPort) {
	n.lock.Lock()
	defer n.lock.Unlock()
	if _, ok := n.packets[to]; !ok {
		n.packets[to] = make(chan TestPacket, 1024)
	}
	n.packets[to] <- TestPacket{
		Addr: from,
		Data: bytes.Clone(b),
	}
}

type TestUDPConn struct {
	addr    netip.AddrPort
	n       *TestNetwork
	packets chan TestPacket
}

func (c *TestUDPConn) WriteToUDPAddrPort(b []byte, addr netip.AddrPort) (int, error) {
	c.n.write(b, c.addr, addr)
	return len(b), nil
}

func (c *TestUDPConn) WriteToUDP(b []byte, addr *net.UDPAddr) (int, error) {
	return c.WriteToUDPAddrPort(b, addr.AddrPort())
}

func (c *TestUDPConn) ReadFromUDPAddrPort(b []byte) (n int, addr netip.AddrPort, err error) {
	pkt := <-c.packets
	b = b[:len(pkt.Data)]
	copy(b, pkt.Data)
	return len(b), pkt.Addr, nil
}

func (c *TestUDPConn) Packets() (out []TestPacket) {
	for {
		select {
		case pkt := <-c.packets:
			out = append(out, pkt)
		default:
			return
		}
	}
}