package bwtest

import (
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/vishvananda/netlink"
)

const udpFeedbackVersion = "receiver-v1"
const udpProbeType uint8 = 4

// UDP without a requested limit probes capacity using receiver feedback instead
// of enqueuing an unbounded amount of traffic ahead of the reverse direction.
type udpPacer struct {
	mu                                sync.Mutex
	limiter                           *RateLimiter
	automatic                         bool
	rate                              int64
	started, lastFeedback, lastChange time.Time
	delivered                         int64
}

func newUDPPacer(limit int64) *udpPacer {
	now := time.Now()
	automatic := limit <= 0
	if automatic {
		limit = 2_000_000
	}
	return &udpPacer{limiter: NewRateLimiter(limit), automatic: automatic, rate: limit, started: now, lastFeedback: now, lastChange: now}
}
func (p *udpPacer) wait(ctx context.Context, n int) bool {
	// Take only a small token-accounting lock; feedback and other streams can
	// update the rate while this stream waits for its next datagram.
	for {
		p.mu.Lock()
		if p.automatic && time.Since(p.lastFeedback) > 2*time.Second && time.Since(p.lastChange) > time.Second {
			p.setRate(max(int64(128000), p.rate/2))
			p.lastChange = time.Now()
		}
		rl := p.limiter
		now := time.Now()
		rl.tokens += now.Sub(rl.lastCheckTime).Seconds() * float64(p.rate) / 8
		rl.lastCheckTime = now
		cap := max(float64(n), float64(p.rate)/8*.02)
		if rl.tokens > cap {
			rl.tokens = cap
		}
		if rl.tokens >= float64(n) {
			rl.tokens -= float64(n)
			p.mu.Unlock()
			return ctx.Err() == nil
		}
		delay := time.Duration((float64(n) - rl.tokens) / (float64(p.rate) / 8) * float64(time.Second))
		p.mu.Unlock()
		if delay > 20*time.Millisecond {
			delay = 20 * time.Millisecond
		}
		timer := time.NewTimer(delay)
		select {
		case <-ctx.Done():
			timer.Stop()
			return false
		case <-timer.C:
		}
	}
}
func (p *udpPacer) setRate(rate int64) { p.rate = rate; p.limiter.bps = rate }
func (p *udpPacer) observe(bytes int64, now time.Time) {
	p.mu.Lock()
	defer p.mu.Unlock()
	dt := now.Sub(p.lastFeedback).Seconds()
	delta := bytes - p.delivered
	if dt <= 0 || delta < 0 {
		return
	}
	p.lastFeedback = now
	p.delivered = bytes
	if !p.automatic || now.Sub(p.started) < time.Second || now.Sub(p.lastChange) < 450*time.Millisecond {
		return
	}
	deliveredRate := int64(float64(delta) * 8 / dt)
	switch {
	case deliveredRate < p.rate*7/10:
		p.setRate(max(int64(128000), deliveredRate*11/10))
	case deliveredRate >= p.rate*85/100:
		p.setRate(min(int64(10_000_000_000), p.rate*5/4))
	}
	p.lastChange = now
}
func udpPayloadSize(requested int, ip net.IP) (int, error) {
	routes, err := netlink.RouteGet(ip)
	if err != nil {
		return 0, fmt.Errorf("bandwidth target route: %w", err)
	}
	if len(routes) == 0 {
		return 0, fmt.Errorf("no route to bandwidth target")
	}
	link, err := netlink.LinkByIndex(routes[0].LinkIndex)
	if err != nil {
		return 0, err
	}
	mtu := link.Attrs().MTU
	if routes[0].MTU > 0 {
		mtu = min(mtu, routes[0].MTU)
	}
	overhead := 28
	if ip.To4() == nil {
		overhead = 48
	}
	size := min(requested, mtu-overhead)
	if size < 64 {
		return 0, fmt.Errorf("target MTU is too small for bandwidth datagrams")
	}
	return size, nil
}
func udpProbe(hash uint64) []byte {
	b := make([]byte, HeaderSize)
	EncodeHeader(b, PacketHeader{Magic: MagicBytes, Version: ProtoVer, PacketType: udpProbeType, SessionHash: hash})
	return b
}

// Each UDP stream has its own sequence space. Late packets within the bounded
// reorder window undo inferred losses, and duplicate datagrams are not counted.
type udpSequenceState struct {
	max  uint64
	seen [1024]uint64
}
type udpTracker struct {
	streams map[uint8]*udpSequenceState
	lost    int64
}

func (t *udpTracker) accept(stream uint8, seq uint64) (bool, int64) {
	if seq == 0 {
		return false, t.lost
	}
	if t.streams == nil {
		t.streams = make(map[uint8]*udpSequenceState)
	}
	s := t.streams[stream]
	if s == nil {
		s = &udpSequenceState{}
		t.streams[stream] = s
	}
	if seq > s.max {
		t.lost += int64(seq - s.max - 1)
		s.max = seq
	} else {
		if s.max-seq >= 1024 || s.seen[seq%1024] == seq {
			return false, t.lost
		}
		if t.lost > 0 {
			t.lost--
		}
	}
	s.seen[seq%1024] = seq
	return true, t.lost
}