package steamcache

import (
	"bytes"
	"context"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func TestBandwidthFairShareRates(t *testing.T) {
	b := newClientBandwidthLimiter(1000, 0)
	lim1, rel1 := b.acquire("1.1.1.1")
	defer rel1()
	lim2, rel2 := b.acquire("2.2.2.2")
	defer rel2()
	if lim1 == nil || lim2 == nil {
		t.Fatal("expected limiters")
	}
	// With 2 active clients, each should get ~500 bytes/sec.
	got1 := float64(lim1.Limit())
	got2 := float64(lim2.Limit())
	if got1 < 400 || got1 > 600 || got2 < 400 || got2 > 600 {
		t.Fatalf("fair-share rates = %v,%v want ~500", got1, got2)
	}
	rel2()
	// After release, sole client should get full uplink.
	lim1b, rel1b := b.acquire("1.1.1.1")
	defer rel1b()
	if float64(lim1b.Limit()) < 900 {
		t.Fatalf("after release limit=%v want ~1000", lim1b.Limit())
	}
}

func TestBandwidthAbsoluteCap(t *testing.T) {
	b := newClientBandwidthLimiter(0, 250)
	lim, rel := b.acquire("9.9.9.9")
	defer rel()
	if lim == nil {
		t.Fatal("expected limiter")
	}
	if float64(lim.Limit()) != 250 {
		t.Fatalf("limit=%v want 250", lim.Limit())
	}
}

func TestBandwidthDisabled(t *testing.T) {
	b := newClientBandwidthLimiter(0, 0)
	lim, rel := b.acquire("9.9.9.9")
	defer rel()
	if lim != nil {
		t.Fatal("expected nil limiter when disabled")
	}
}

func TestLimitedResponseWriterShapes(t *testing.T) {
	var buf bytes.Buffer
	rec := httptest.NewRecorder()
	// Use a custom writer sink via ResponseRecorder is fine; WaitN will delay.
	lim := newClientBandwidthLimiter(0, 2000) // 2KB/s
	l, rel := lim.acquire("127.0.0.1")
	defer rel()
	w := &limitedResponseWriter{ResponseWriter: rec, lim: l, ctx: context.Background()}
	payload := bytes.Repeat([]byte("x"), 4000)
	start := time.Now()
	n, err := w.Write(payload)
	elapsed := time.Since(start)
	if err != nil {
		t.Fatal(err)
	}
	if n != len(payload) {
		t.Fatalf("wrote %d want %d", n, len(payload))
	}
	_ = buf
	// 4000 bytes at 2000 B/s should take ~2s (allow slack for CI).
	if elapsed < 1500*time.Millisecond {
		t.Fatalf("elapsed %v too fast for 2KB/s shaping of 4KB", elapsed)
	}
	if elapsed > 8*time.Second {
		t.Fatalf("elapsed %v unexpectedly slow", elapsed)
	}
}

func TestServeHTTPBandwidthCap(t *testing.T) {
	body := bytes.Repeat([]byte("a"), 3000)
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/octet-stream")
		w.Header().Set("Content-Length", "3000")
		w.WriteHeader(200)
		_, _ = w.Write(body)
	}))
	t.Cleanup(upstream.Close)

	sc, err := NewWithOptions(Options{
		Address:                 "127.0.0.1:0",
		MemorySize:              "1MB",
		DiskSize:                "0",
		Upstream:                upstream.URL,
		MemoryGC:                "lru",
		DiskGC:                  "lru",
		MaxConcurrentRequests:   20,
		MaxRequestsPerClient:    10,
		MaxObjectSize:           "0",
		MaxBytesPerClientPerSec: 1500, // 1.5KB/s
	})
	if err != nil {
		t.Fatalf("NewWithOptions: %v", err)
	}
	t.Cleanup(func() { sc.Shutdown() })

	req := httptest.NewRequest(http.MethodGet, "/depot/bw/chunk", nil)
	req.Header.Set("User-Agent", "Valve/Steam HTTP Client 1.0")
	rr := httptest.NewRecorder()
	start := time.Now()
	sc.ServeHTTP(rr, req)
	elapsed := time.Since(start)
	if rr.Code != 200 {
		t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
	}
	got := rr.Body.Bytes()
	if len(got) != len(body) {
		t.Fatalf("body len=%d want %d", len(got), len(body))
	}
	if elapsed < time.Second {
		t.Fatalf("elapsed %v too fast for shaping", elapsed)
	}
}