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 UDPRClient struct {
	conn       *net.UDPConn
	ackManager *ackManager
}

func DialUDPR(network string, laddr, raddr *net.UDPAddr) (*UDPRClient, error) {
	conn, err := net.DialUDP(network, laddr, raddr)
	if err != nil {
		return nil, err
	}
	return &UDPRClient{
		conn:       conn,
		ackManager: newAckManager(),
	}, nil
}

func (u *UDPRClient) Write(data []byte) 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 tries := 0; tries < maxRetries; tries++ {
		if _, err := u.conn.Write(dg); 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 *UDPRClient) Read() ([]byte, error) {
	buf := make([]byte, maxDatagramSize)
	for {
		n, err := u.conn.Read(buf)
		if err != nil {
			return 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.Write(ack)
			return buf[5:n], nil
		case ackType:
			u.ackManager.clearPending(ackNum)
			continue
		default:
			continue
		}
	}
}

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