package udpr

import (
	"encoding/binary"
	"errors"
	"math/rand"
	"net"
	"time"
)

const (
	dataType       = 0x00
	ackType        = 0x01
	headerSize     = 5
	maxDatagramSize = 508
	maxRetries     = 5
	initialDelay   = 1 * time.Second
	maxDelay       = 16 * time.Second
)

var errFailedAck = errors.New("failed to receive ACK")

type UDPRConn struct {
	conn       *net.UDPConn
	ackManager *ackManager
}

func ListenUDPR(network string, laddr *net.UDPAddr) (*UDPRConn, error) {
	conn, err := net.ListenUDP(network, laddr)
	if err != nil {
		return nil, err
	}
	return &UDPRConn{
		conn:       conn,
		ackManager: newAckManager(),
	}, nil
}

func (u *UDPRConn) WriteToUDPR(data []byte, dest *net.UDPAddr) error {
	ackNum := rand.Uint32()
	dg := make([]byte, headerSize+len(data))
	dg[0] = dataType
	binary.BigEndian.PutUint32(dg[1:5], ackNum)
	copy(dg[5:], data)

	u.ackManager.setPending(ackNum)

	delay := initialDelay
	for retries := 0; retries < maxRetries; retries++ {
		if _, err := u.conn.WriteToUDP(dg, dest); err != nil {
			return err
		}
		time.Sleep(delay)
		if !u.ackManager.isPending(ackNum) {
			return nil
		}
		if delay < maxDelay {
			delay *= 2
		}
	}
	u.ackManager.clearPending(ackNum)
	return errFailedAck
}

func (u *UDPRConn) ReadFromUDPR() ([]byte, *net.UDPAddr, error) {
	buf := make([]byte, maxDatagramSize)
	for {
		n, addr, err := u.conn.ReadFromUDP(buf)
		if err != nil {
			return nil, nil, err
		}
		if n < headerSize {
			continue
		}

		dgType := buf[0]
		ackNum := binary.BigEndian.Uint32(buf[1:5])
		switch dgType {
		case dataType:
			ack := make([]byte, headerSize)
			ack[0] = ackType
			binary.BigEndian.PutUint32(ack[1:5], ackNum)
			_, _ = u.conn.WriteToUDP(ack, addr)
			return buf[5:n], addr, nil
		case ackType:
			u.ackManager.clearPending(ackNum)
			continue
		default:
			continue
		}
	}
}

func (u *UDPRConn) Close() {
	u.conn.Close()
}