package bench

import (
	"context"
	"strings"
	"testing"
	"time"
)

func TestCapacityIsTheLastLevelThatMetTheSLONotThePeak(t *testing.T) {
	// A GPU keeps delivering more tokens per second long after latency has
	// become unusable. Taking peak throughput as capacity would tell the router
	// a replica can serve a load at which nobody would want the answers.
	r := Result{SLO: SLO{TTFTp99: time.Second}}
	r.Steps = []Step{
		{Concurrency: 1, Completed: 10, TokensPerSec: 100, MeetsSLO: true,
			TTFT: Percentiles{P99: 100 * time.Millisecond}},
		{Concurrency: 8, Completed: 80, TokensPerSec: 700, MeetsSLO: true,
			TTFT: Percentiles{P99: 900 * time.Millisecond}, Waiting: 3.4, KVUsage: 0.71},
		// Faster in tokens, far too slow in latency.
		{Concurrency: 32, Completed: 200, TokensPerSec: 1500, MeetsSLO: false,
			TTFT: Percentiles{P99: 9 * time.Second}},
	}
	r.conclude()

	if r.MaxConcurrencyAtSLO != 8 {
		t.Errorf("capacity = %d, want 8", r.MaxConcurrencyAtSLO)
	}
	if r.ThroughputTokPerSec != 700 {
		t.Errorf("throughput = %v; it must come from the level that met the SLO", r.ThroughputTokPerSec)
	}
	if r.SuggestedTargetWaiting != 3 {
		t.Errorf("target_waiting = %d, want 3", r.SuggestedTargetWaiting)
	}
	if r.SuggestedKVCeiling != 0.70 {
		t.Errorf("kv_ceiling = %v, want 0.70", r.SuggestedKVCeiling)
	}
}

func TestAStepWithFailuresHasNotMetAnything(t *testing.T) {
	// Latency percentiles over the requests that survived say nothing useful
	// when the rest were refused.
	r := Result{SLO: SLO{TTFTp99: time.Second}}
	r.Steps = []Step{
		{Concurrency: 4, Completed: 10, Failed: 90, TokensPerSec: 50, MeetsSLO: false,
			TTFT: Percentiles{P99: 50 * time.Millisecond}},
	}
	r.conclude()
	if r.MaxConcurrencyAtSLO != 0 {
		t.Errorf("a step that failed 90 of 100 requests was counted as capacity")
	}
}

func TestSLOIsMissedWhenEitherHalfIs(t *testing.T) {
	s := SLO{TTFTp99: time.Second, TPOTp99: 50 * time.Millisecond}
	cases := []struct {
		ttft, tpot time.Duration
		want       bool
	}{
		{500 * time.Millisecond, 20 * time.Millisecond, true},
		{2 * time.Second, 20 * time.Millisecond, false},
		{500 * time.Millisecond, 80 * time.Millisecond, false},
		{2 * time.Second, 80 * time.Millisecond, false},
	}
	for _, c := range cases {
		if got := s.Met(c.ttft, c.tpot); got != c.want {
			t.Errorf("Met(%s, %s) = %v", c.ttft, c.tpot, got)
		}
	}

	// A target of zero is no target, not a target of zero.
	none := SLO{}
	if !none.Met(time.Hour, time.Hour) {
		t.Error("an unset SLO refused everything")
	}
}

func TestPercentiles(t *testing.T) {
	var v []time.Duration
	for i := 1; i <= 100; i++ {
		v = append(v, time.Duration(i)*time.Millisecond)
	}
	p := percentiles(v)
	if p.P50 != 50*time.Millisecond || p.P90 != 90*time.Millisecond {
		t.Errorf("p50=%s p90=%s", p.P50, p.P90)
	}
	if p.Max != 100*time.Millisecond {
		t.Errorf("max=%s", p.Max)
	}
	if got := percentiles(nil); got != (Percentiles{}) {
		t.Errorf("empty sample gave %+v", got)
	}
}

func TestPromptsAreUniqueUnlessAskedNotToBe(t *testing.T) {
	// The default matters more than it looks. With one repeated prompt the
	// engine's cache answers most of the prefill, and the measurement describes
	// the cache rather than the hardware.
	cfg := Config{PromptTokens: 64}
	seen := map[string]bool{}
	for range 100 {
		seen[prompt(cfg)] = true
	}
	if len(seen) < 100 {
		t.Errorf("%d distinct prompts out of 100; the cache would answer the rest", len(seen))
	}

	cfg.SharedPrefix = true
	first := prompt(cfg)
	for range 20 {
		if got := prompt(cfg); !strings.HasPrefix(got, first[:200]) {
			t.Error("--shared-prefix did not produce a shared prefix")
		}
	}
}

func TestPromptSizeFollowsTheRequest(t *testing.T) {
	small := prompt(Config{PromptTokens: 32})
	large := prompt(Config{PromptTokens: 512})
	if len(large) < len(small)*4 {
		t.Errorf("a 512-token prompt is %d chars against %d for 32", len(large), len(small))
	}
}

func TestConfigIsValidatedBeforeAnyLoadIsGenerated(t *testing.T) {
	// The failure has to come before the traffic, not halfway through a ramp.
	if _, err := Run(context.Background(), Config{Model: "m"}, nil); err == nil {
		t.Error("a run with no target started")
	}
	if _, err := Run(context.Background(), Config{Target: "http://x"}, nil); err == nil {
		t.Error("a run with no model started")
	}
}

func TestTheRampStopsOnceTheSLOBreaks(t *testing.T) {
	// Everything above a missed level is slower by construction, and finding
	// that out costs real GPU time.
	r := Result{}
	r.Steps = []Step{
		{Concurrency: 1, Completed: 1, MeetsSLO: true},
		{Concurrency: 2, Completed: 1, MeetsSLO: false},
	}
	r.conclude()
	if r.MaxConcurrencyAtSLO != 1 {
		t.Errorf("capacity = %d", r.MaxConcurrencyAtSLO)
	}
}

func TestRoundTo(t *testing.T) {
	cases := map[float64]float64{0.71: 0.70, 0.73: 0.75, 0.88: 0.90, 0.0: 0.0}
	for in, want := range cases {
		if got := roundTo(in, 0.05); got != want {
			t.Errorf("roundTo(%v) = %v, want %v", in, got, want)
		}
	}
}