package udphop

import (
	"context"
	"errors"
	"math/rand"
	"net"
	"sync"
	"syscall"
	"time"
)

const (
	packetQueueSize = 1024
	udpBufferSize   = 2048 // QUIC packets are at most 1500 bytes long, so 2k should be more than enough

	defaultHopInterval = 30 * time.Second
)

type udpHopPacketConn struct {
	Addr        net.Addr
	Addrs       []net.Addr
	HopInterval time.Duration
	dialFunc    contextDialFunc
	ctx         context.Context
	cancel      context.CancelFunc

	connMutex   sync.RWMutex
	prevConn    net.PacketConn
	currentAddr net.Addr
	currentConn net.PacketConn

	readBufferSize  int
	writeBufferSize int

	recvQueue chan *udpPacket
	closeChan chan struct{}
	closed    bool

	bufPool sync.Pool
}

type udpPacket struct {
	Buf  []byte
	N    int
	Addr net.Addr
	Err  error
}

type dialFunc = func(addr net.Addr) (net.PacketConn, error)
type contextDialFunc = func(ctx context.Context, addr net.Addr) (net.PacketConn, error)

func NewUDPHopPacketConn(addr *UDPHopAddr, hopInterval time.Duration, dialFunc dialFunc) (net.PacketConn, error) {
	return NewUDPHopPacketConnContext(context.Background(), addr, hopInterval, func(_ context.Context, addr net.Addr) (net.PacketConn, error) {
		return dialFunc(addr)
	})
}

func NewUDPHopPacketConnContext(ctx context.Context, addr *UDPHopAddr, hopInterval time.Duration, dialFunc contextDialFunc) (net.PacketConn, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	if hopInterval == 0 {
		hopInterval = defaultHopInterval
	} else if hopInterval < 5*time.Second {
		return nil, errors.New("hop interval must be at least 5 seconds")
	}
	addrs, err := addr.addrs()
	if err != nil {
		return nil, err
	}

	newAddrIndex := rand.Intn(len(addrs))
	curAddr := addrs[newAddrIndex]
	curConn, err := dialFunc(ctx, curAddr)
	if err != nil {
		return nil, err
	}
	if actualAddr := remoteAddrOfPacketConn(curConn); actualAddr != nil {
		if frozenAddrs, frozenCurrentAddr, ok := freezeAddrsForResolvedIP(addr, actualAddr); ok {
			addrs = frozenAddrs
			curAddr = frozenCurrentAddr
		} else {
			curAddr = actualAddr
		}
	}
	hopCtx, cancel := context.WithCancel(context.Background())
	hConn := &udpHopPacketConn{
		Addr:        addr,
		Addrs:       addrs,
		HopInterval: hopInterval,
		dialFunc:    dialFunc,
		ctx:         hopCtx,
		cancel:      cancel,
		prevConn:    nil,
		currentAddr: curAddr,
		currentConn: curConn,
		recvQueue:   make(chan *udpPacket, packetQueueSize),
		closeChan:   make(chan struct{}),
		bufPool: sync.Pool{
			New: func() interface{} {
				return make([]byte, udpBufferSize)
			},
		},
	}
	go hConn.recvLoop(curConn)
	go hConn.hopLoop()
	return hConn, nil
}

func (u *udpHopPacketConn) recvLoop(conn net.PacketConn) {
	for {
		buf := u.bufPool.Get().([]byte)
		n, addr, err := conn.ReadFrom(buf)
		if err != nil {
			u.bufPool.Put(buf) // nolint:staticcheck
			var netErr net.Error
			if errors.As(err, &netErr) && netErr.Timeout() {
				// Only pass through timeout errors here, not permanent errors
				// like connection closed. Connection close is normal as we close
				// the old connection to exit this loop every time we hop.
				select {
				case u.recvQueue <- &udpPacket{nil, 0, nil, netErr}:
				case <-u.closeChan:
				default:
					// If the consumer is already backlogged, dropping the timeout
					// notification is better than pinning this goroutine forever.
				}
			}
			return
		}
		select {
		case u.recvQueue <- &udpPacket{buf, n, addr, nil}:
			// Packet successfully queued
		default:
			// Queue is full, drop the packet
			u.bufPool.Put(buf) // nolint:staticcheck
		}
	}
}

func (u *udpHopPacketConn) hopLoop() {
	ticker := time.NewTicker(u.HopInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			u.hop()
		case <-u.closeChan:
			return
		}
	}
}

func (u *udpHopPacketConn) hop() {
	u.connMutex.RLock()
	if u.closed || len(u.Addrs) == 0 {
		u.connMutex.RUnlock()
		return
	}
	addrs := append([]net.Addr(nil), u.Addrs...)
	dialFunc := u.dialFunc
	u.connMutex.RUnlock()

	newAddrIndex := rand.Intn(len(addrs))
	newAddr := addrs[newAddrIndex]
	newConn, err := dialFunc(u.context(), newAddr)
	if err != nil {
		// Could be temporary, just skip this hop
		return
	}
	if actualAddr := remoteAddrOfPacketConn(newConn); actualAddr != nil {
		newAddr = actualAddr
	}

	u.connMutex.Lock()
	defer u.connMutex.Unlock()
	if u.closed {
		_ = newConn.Close()
		return
	}
	// We need to keep receiving packets from the previous connection,
	// because otherwise there will be packet loss due to the time gap
	// between we hop to a new port and the server acknowledges this change.
	// So we do the following:
	// Close prevConn,
	// move currentConn to prevConn,
	// set newConn as currentConn,
	// start recvLoop on newConn.
	if u.prevConn != nil {
		_ = u.prevConn.Close() // recvLoop for this conn will exit
	}
	u.prevConn = u.currentConn
	u.currentAddr = newAddr
	u.currentConn = newConn
	// Set buffer sizes if previously set
	if u.readBufferSize > 0 {
		_ = trySetReadBuffer(u.currentConn, u.readBufferSize)
	}
	if u.writeBufferSize > 0 {
		_ = trySetWriteBuffer(u.currentConn, u.writeBufferSize)
	}
	go u.recvLoop(newConn)
}

func (u *udpHopPacketConn) context() context.Context {
	if u.ctx == nil {
		return context.Background()
	}
	return u.ctx
}

func (u *udpHopPacketConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
	for {
		select {
		case p := <-u.recvQueue:
			if p.Err != nil {
				return 0, nil, p.Err
			}
			// Preserve the actual packet source so callers see the concrete
			// hop address instead of the logical port-union descriptor.
			n := copy(b, p.Buf[:p.N])
			u.bufPool.Put(p.Buf) // nolint:staticcheck
			addr = p.Addr
			if addr == nil {
				u.connMutex.RLock()
				addr = u.currentAddr
				u.connMutex.RUnlock()
			}
			return n, addr, nil
		case <-u.closeChan:
			return 0, nil, net.ErrClosed
		}
	}
}

func (u *udpHopPacketConn) WriteTo(b []byte, _ net.Addr) (n int, err error) {
	u.connMutex.RLock()
	defer u.connMutex.RUnlock()
	if u.closed {
		return 0, net.ErrClosed
	}
	return u.currentConn.WriteTo(b, u.currentAddr)
}

func (u *udpHopPacketConn) Close() error {
	if u.cancel != nil {
		u.cancel()
	}
	u.connMutex.Lock()
	defer u.connMutex.Unlock()
	if u.closed {
		return nil
	}
	// Close prevConn and currentConn
	// Close closeChan to unblock ReadFrom & hopLoop
	// Set closed flag to true to prevent double close
	if u.prevConn != nil {
		_ = u.prevConn.Close()
	}
	err := u.currentConn.Close()
	close(u.closeChan)
	u.closed = true
	u.Addrs = nil // For GC
	u.currentAddr = nil
	return err
}

func (u *udpHopPacketConn) RemoteAddr() net.Addr {
	u.connMutex.RLock()
	defer u.connMutex.RUnlock()
	return u.currentAddr
}

func (u *udpHopPacketConn) LocalAddr() net.Addr {
	u.connMutex.RLock()
	defer u.connMutex.RUnlock()
	return u.currentConn.LocalAddr()
}

func (u *udpHopPacketConn) SetDeadline(t time.Time) error {
	u.connMutex.RLock()
	defer u.connMutex.RUnlock()
	if u.prevConn != nil {
		_ = u.prevConn.SetDeadline(t)
	}
	return u.currentConn.SetDeadline(t)
}

func (u *udpHopPacketConn) SetReadDeadline(t time.Time) error {
	u.connMutex.RLock()
	defer u.connMutex.RUnlock()
	if u.prevConn != nil {
		_ = u.prevConn.SetReadDeadline(t)
	}
	return u.currentConn.SetReadDeadline(t)
}

func (u *udpHopPacketConn) SetWriteDeadline(t time.Time) error {
	u.connMutex.RLock()
	defer u.connMutex.RUnlock()
	if u.prevConn != nil {
		_ = u.prevConn.SetWriteDeadline(t)
	}
	return u.currentConn.SetWriteDeadline(t)
}

// UDP-specific methods below

func (u *udpHopPacketConn) SetReadBuffer(bytes int) error {
	u.connMutex.Lock()
	defer u.connMutex.Unlock()
	u.readBufferSize = bytes
	if u.prevConn != nil {
		_ = trySetReadBuffer(u.prevConn, bytes)
	}
	return trySetReadBuffer(u.currentConn, bytes)
}

func (u *udpHopPacketConn) SetWriteBuffer(bytes int) error {
	u.connMutex.Lock()
	defer u.connMutex.Unlock()
	u.writeBufferSize = bytes
	if u.prevConn != nil {
		_ = trySetWriteBuffer(u.prevConn, bytes)
	}
	return trySetWriteBuffer(u.currentConn, bytes)
}

func (u *udpHopPacketConn) SyscallConn() (syscall.RawConn, error) {
	u.connMutex.RLock()
	defer u.connMutex.RUnlock()
	sc, ok := u.currentConn.(syscall.Conn)
	if !ok {
		return nil, errors.New("not supported")
	}
	return sc.SyscallConn()
}

func trySetReadBuffer(pc net.PacketConn, bytes int) error {
	sc, ok := pc.(interface {
		SetReadBuffer(bytes int) error
	})
	if ok {
		return sc.SetReadBuffer(bytes)
	}
	return nil
}

func trySetWriteBuffer(pc net.PacketConn, bytes int) error {
	sc, ok := pc.(interface {
		SetWriteBuffer(bytes int) error
	})
	if ok {
		return sc.SetWriteBuffer(bytes)
	}
	return nil
}

func remoteAddrOfPacketConn(conn net.PacketConn) net.Addr {
	type remoteAddrConn interface {
		RemoteAddr() net.Addr
	}
	if c, ok := conn.(remoteAddrConn); ok {
		if addr := c.RemoteAddr(); addr != nil {
			return addr
		}
	}
	return nil
}

func freezeAddrsForResolvedIP(addr *UDPHopAddr, actualAddr net.Addr) ([]net.Addr, net.Addr, bool) {
	if addr == nil || len(addr.Ports) == 0 {
		return nil, nil, false
	}
	udpAddr, ok := actualAddr.(*net.UDPAddr)
	if !ok || udpAddr == nil || udpAddr.IP == nil {
		return nil, nil, false
	}
	ip := append(net.IP(nil), udpAddr.IP...)
	addrs := make([]net.Addr, 0, len(addr.Ports))
	for _, port := range addr.Ports {
		addrs = append(addrs, &net.UDPAddr{
			IP:   ip,
			Port: int(port),
			Zone: udpAddr.Zone,
		})
	}
	currentAddr := &net.UDPAddr{
		IP:   append(net.IP(nil), udpAddr.IP...),
		Port: udpAddr.Port,
		Zone: udpAddr.Zone,
	}
	return addrs, currentAddr, true
}
