package client

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestParseBandwidth(t *testing.T) {
	cases := []struct {
		in   string
		want int64
	}{
		{"", 0},
		{"20mbps", 2_500_000},
		{"20Mbps", 2_500_000},
		{"20 mbps", 2_500_000},
		{"1.5mbit", 187_500},
		{"512kbps", 64_000},
		{"1gbps", 125_000_000},
		{"2MB/s", 2 << 20},
		{"500kb/s", 500 << 10},
		{"1GB/s", 1 << 30},
		{"8bps", 1},
		{"4096b/s", 4096},
		{"1500000", 1_500_000}, // bare number is already bytes/sec
	}
	for _, c := range cases {
		got, err := parseBandwidth(c.in)
		if err != nil {
			t.Errorf("parseBandwidth(%q): unexpected error %v", c.in, err)
			continue
		}
		if got != c.want {
			t.Errorf("parseBandwidth(%q) = %d, want %d", c.in, got, c.want)
		}
	}

	for _, bad := range []string{"fast", "20megabits", "-5mbps", "0", "0mbps", "mbps", "20 mb ps"} {
		if _, err := parseBandwidth(bad); err == nil {
			t.Errorf("parseBandwidth(%q): expected an error", bad)
		}
	}
}

// A nil shaper is the "unlimited" case and must be safe on every path, because
// call sites deliberately do not branch on it.
func TestNilShaperIsUnlimited(t *testing.T) {
	var sh *Shaper
	if sh = NewShaper(0); sh != nil {
		t.Fatal("NewShaper(0) should return nil")
	}
	if got := sh.chunkSize(); got != DataChunkSize {
		t.Errorf("nil chunkSize = %d, want %d", got, DataChunkSize)
	}
	if !sh.Acquire(&shaperShare{}, 1<<20, nil) {
		t.Error("nil Acquire should always succeed")
	}
	sh.Stop() // must not panic
}

// The virtual-time bookkeeping is what makes the shaper fair, so assert it
// directly. The rate is high enough that tokens never bind, leaving only the
// stamping under test — no timing, no flakiness.
func TestShaperIdleStreamCannotHoardCredit(t *testing.T) {
	sh := NewShaper(1 << 30)
	defer sh.Stop()

	var heavy, light shaperShare
	for i := 0; i < 10; i++ {
		if !sh.Acquire(&heavy, 1000, nil) {
			t.Fatal("acquire failed")
		}
	}
	if heavy.vfinish != 10000 {
		t.Errorf("heavy.vfinish = %v, want 10000", heavy.vfinish)
	}
	sh.mu.Lock()
	vclock := sh.vclock
	sh.mu.Unlock()
	if vclock != 9000 {
		t.Errorf("vclock = %v, want 9000 (the stamp of the last request served)", vclock)
	}

	// light was idle for all of it. Its stale vfinish of 0 must be clamped up to
	// the current clock: it may not bank the virtual time it never spent, which
	// is what would let it starve heavy on return.
	if !sh.Acquire(&light, 1000, nil) {
		t.Fatal("acquire failed")
	}
	if light.vfinish != vclock+1000 {
		t.Errorf("light.vfinish = %v, want %v (clamped to the clock, not 1000)", light.vfinish, vclock+1000)
	}
}

func TestShaperEnforcesRate(t *testing.T) {
	const rate = 1 << 20 // 1 MiB/s
	sh := NewShaper(rate)
	defer sh.Stop()

	var share shaperShare
	const total = 512 << 10
	const chunk = 8 << 10

	start := time.Now()
	for sent := 0; sent < total; sent += chunk {
		if !sh.Acquire(&share, chunk, nil) {
			t.Fatal("acquire failed")
		}
	}
	elapsed := time.Since(start)

	// The bucket starts full, so the burst is free and only the remainder is
	// paced: (512 KiB - 200 KiB) / 1 MiB/s ≈ 300 ms. Bounds are wide on purpose.
	if elapsed < 200*time.Millisecond {
		t.Errorf("sent %d bytes at %d B/s in only %v; the cap is not being enforced", total, rate, elapsed)
	}
	if elapsed > time.Second {
		t.Errorf("took %v, far longer than the ~300ms the rate implies", elapsed)
	}
}

// An idle stream must be able to spend the banked burst at once, otherwise a
// player joining pays for the cap in visible chunk-loading latency.
func TestShaperAllowsBurst(t *testing.T) {
	sh := NewShaper(1 << 20)
	defer sh.Stop()

	var share shaperShare
	start := time.Now()
	for i := 0; i < 6; i++ {
		if !sh.Acquire(&share, 32<<10, nil) { // 192 KiB, inside the 200 KiB bucket
			t.Fatal("acquire failed")
		}
	}
	if elapsed := time.Since(start); elapsed > 100*time.Millisecond {
		t.Errorf("burst of 192 KiB took %v; the bucket should have covered it instantly", elapsed)
	}
}

// The point of the whole exercise: a stream that never stops asking must not
// crowd another one out.
func TestShaperSharesFairlyBetweenStreams(t *testing.T) {
	sh := NewShaper(1 << 20)
	defer sh.Stop()

	stop := make(chan struct{})
	var counts [2]atomic.Int64
	var wg sync.WaitGroup
	for i := range counts {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			var share shaperShare
			for {
				if !sh.Acquire(&share, 4<<10, stop) {
					return
				}
				counts[i].Add(4 << 10)
			}
		}(i)
	}

	// Spend the token bucket before measuring, the way the idle-credit test
	// raises the rate so tokens never bind: isolate the property under test.
	//
	// While the bucket has tokens there is no queue to arbitrate — every request
	// is granted the moment it arrives, and being the stream that is owed service
	// only helps when both are enqueued at the same instant. The burst is
	// therefore first-come-first-served by construction, and at 0.2s of
	// transmission it is a quarter of a 600ms window, enough to swamp the result:
	// a run where one goroutine happened to win the bucket landed at 520192 vs
	// 315392, which is exactly "one took the whole burst, then the two split the
	// remainder evenly".
	//
	// Fairness here is a steady-state property, and that is what matters in
	// practice — the bucket is empty whenever the link is actually busy.
	time.Sleep(250 * time.Millisecond)
	counts[0].Store(0)
	counts[1].Store(0)

	time.Sleep(600 * time.Millisecond)
	close(stop)
	wg.Wait()

	a, b := counts[0].Load(), counts[1].Load()
	if a == 0 || b == 0 {
		t.Fatalf("one stream was starved entirely: %d vs %d", a, b)
	}
	lo, hi := min(a, b), max(a, b)
	t.Logf("steady-state split: %d vs %d bytes (%.3fx)", a, b, float64(hi)/float64(lo))
	if float64(hi) > 1.35*float64(lo) {
		t.Errorf("unfair split: %d vs %d bytes (>35%% apart)", a, b)
	}
}

// A stream torn down while parked must release immediately and leave no trace
// in the queue, or its goroutine (and the Stream it closes over) leaks.
func TestShaperAcquireCancels(t *testing.T) {
	sh := NewShaper(MinBandwidth) // 8 KiB/s: a parked request would wait seconds
	defer sh.Stop()

	var share shaperShare
	if !sh.Acquire(&share, MinShaperBurst, nil) { // drain the bucket
		t.Fatal("acquire failed")
	}

	cancel := make(chan struct{})
	result := make(chan bool, 1)
	go func() { result <- sh.Acquire(&share, 32<<10, cancel) }()

	time.Sleep(50 * time.Millisecond)
	close(cancel)

	select {
	case ok := <-result:
		if ok {
			t.Error("Acquire returned true after cancellation")
		}
	case <-time.After(time.Second):
		t.Fatal("Acquire did not return after its cancel channel closed")
	}

	sh.mu.Lock()
	n := len(sh.waiting)
	sh.mu.Unlock()
	if n != 0 {
		t.Errorf("%d cancelled request(s) left in the queue", n)
	}
}