// Package bench measures what a replica can actually do, so the numbers the
// router depends on stop being guesses (plan §6.2.1).
//
// Three of them matter: the weight a replica gets when the scheduler compares
// it to others, the queue depth admission treats as healthy, and the KV usage
// it treats as full. All three are currently hand-written in configuration, and
// a hand-written capacity number is a hand-written opinion about hardware
// nobody measured.
//
// The measurement is a load ramp. Concurrency climbs until the SLO breaks, and
// the last level that still met it is the capacity — not peak throughput, which
// a GPU will happily deliver long after latency has become unusable.
package bench

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/rand/v2"
	"net/http"
	"sort"
	"strings"
	"sync"
	"time"
)

// SLO is the latency target capacity is measured against.
type SLO struct {
	TTFTp99 time.Duration
	TPOTp99 time.Duration
}

// Met reports whether a step's latencies satisfy the target.
func (s SLO) Met(ttft, tpot time.Duration) bool {
	if s.TTFTp99 > 0 && ttft > s.TTFTp99 {
		return false
	}
	if s.TPOTp99 > 0 && tpot > s.TPOTp99 {
		return false
	}
	return true
}

// Config describes a run.
type Config struct {
	// Target is the base URL: the gateway, or one replica directly. Pointing at
	// a replica measures that replica; pointing at the gateway measures the
	// whole path including scheduling, which is a different question.
	Target string
	// MetricsURL is an engine's /metrics, sampled during the run so queue depth
	// and KV usage can be reported alongside latency. Only meaningful when
	// Target is a single replica.
	MetricsURL string
	// Engine names which metric vocabulary MetricsURL speaks. Getting it wrong
	// is silent: nothing matches, every sample reads zero, and the suggested
	// thresholds are built on nothing.
	Engine string
	// MetricsKey is the engine key MetricsURL is read with. Not APIKey: that
	// is whatever Target takes, which for the gateway is a client's key, and
	// an engine that guards its metrics (llama.cpp, TensorRT-LLM) takes only
	// its own.
	MetricsKey string
	Model      string
	APIKey     string

	// Concurrency is the ramp. Each level runs for Duration.
	Concurrency []int
	Duration    time.Duration

	PromptTokens int
	MaxTokens    int

	// SharedPrefix sends every request the same prompt.
	//
	// Off by default, and that default is the whole point: with a shared prefix
	// the engine's cache answers most of the prefill and the measurement
	// describes the cache rather than the hardware. Turning it on measures the
	// cached case deliberately, which is a fair question to ask separately.
	SharedPrefix bool

	SLO SLO
}

// DefaultConfig returns a ramp that covers most single-replica hardware.
func DefaultConfig() Config {
	return Config{
		Concurrency:  []int{1, 2, 4, 8, 16, 32, 64, 128},
		Duration:     20 * time.Second,
		PromptTokens: 512,
		MaxTokens:    128,
		SLO:          SLO{TTFTp99: 2 * time.Second, TPOTp99: 50 * time.Millisecond},
	}
}

func (c *Config) normalize() error {
	d := DefaultConfig()
	if c.Target == "" {
		return errors.New("bench: a target is required")
	}
	if c.Model == "" {
		return errors.New("bench: a model is required")
	}
	c.Target = strings.TrimRight(c.Target, "/")
	if len(c.Concurrency) == 0 {
		c.Concurrency = d.Concurrency
	}
	if c.Duration <= 0 {
		c.Duration = d.Duration
	}
	if c.PromptTokens <= 0 {
		c.PromptTokens = d.PromptTokens
	}
	if c.MaxTokens <= 0 {
		c.MaxTokens = d.MaxTokens
	}
	sort.Ints(c.Concurrency)
	return nil
}

// Percentiles summarizes a latency sample.
type Percentiles struct {
	P50, P90, P99, Max time.Duration
}

// Step is one concurrency level's result.
type Step struct {
	Concurrency int
	Completed   int
	Failed      int
	// Requests and output tokens per second, over the step's wall clock.
	RequestsPerSec float64
	TokensPerSec   float64
	TTFT           Percentiles
	// TPOT is time per output token after the first, which is what a reader
	// experiences as the speed of the stream.
	TPOT     Percentiles
	Waiting  float64
	Running  float64
	KVUsage  float64
	MeetsSLO bool
	// FirstError is kept so a step that failed says why rather than just
	// reporting zeroes.
	FirstError string
	// MetricsUnread says why a step that was to read the engine's metrics
	// has none, so its zero queue depth and KV usage are not taken for an
	// idle engine.
	MetricsUnread string `json:",omitempty"`
}

// Result is the whole ramp plus what it concluded.
type Result struct {
	Model   string
	Target  string
	Steps   []Step
	SLO     SLO
	RanAt   time.Time
	Elapsed time.Duration

	// MaxConcurrencyAtSLO is the capacity: the highest level that still met the
	// target. Zero means even one request did not.
	MaxConcurrencyAtSLO int
	// ThroughputTokPerSec and the latencies are taken at that level.
	ThroughputTokPerSec float64
	TTFTp99, TPOTp99    time.Duration
	// SuggestedTargetWaiting and SuggestedKVCeiling are what admission should
	// treat as healthy, observed rather than guessed.
	SuggestedTargetWaiting int
	SuggestedKVCeiling     float64
}

// Run executes the ramp.
func Run(ctx context.Context, cfg Config, progress func(Step)) (Result, error) {
	if err := cfg.normalize(); err != nil {
		return Result{}, err
	}
	res := Result{Model: cfg.Model, Target: cfg.Target, SLO: cfg.SLO, RanAt: time.Now()}

	// Read once before any load. Metrics that cannot be read — refused for
	// want of the engine's key, or written in another engine's names — fail
	// every sample alike, and finding that out after the ramp has cost its
	// GPU time would leave a run with no thresholds to suggest.
	if cfg.MetricsURL != "" {
		if err := newSampler(cfg.MetricsURL, cfg.Engine, cfg.MetricsKey).read(ctx); err != nil {
			return Result{}, fmt.Errorf("bench: the engine's metrics cannot be read: %w", err)
		}
	}

	start := time.Now()

	client := &http.Client{Timeout: 10 * time.Minute}
	for _, c := range cfg.Concurrency {
		if ctx.Err() != nil {
			break
		}
		step := runStep(ctx, client, cfg, c)
		res.Steps = append(res.Steps, step)
		if progress != nil {
			progress(step)
		}
		// Stop climbing once the target is missed. Everything above is slower
		// by construction, and the cost of finding out is real GPU time.
		if !step.MeetsSLO {
			break
		}
	}
	res.Elapsed = time.Since(start)
	res.conclude()
	return res, nil
}

func (r *Result) conclude() {
	for _, s := range r.Steps {
		if !s.MeetsSLO || s.Completed == 0 {
			continue
		}
		if s.Concurrency <= r.MaxConcurrencyAtSLO {
			continue
		}
		r.MaxConcurrencyAtSLO = s.Concurrency
		r.ThroughputTokPerSec = s.TokensPerSec
		r.TTFTp99, r.TPOTp99 = s.TTFT.P99, s.TPOT.P99
		// The queue depth and cache pressure seen while still meeting the SLO
		// are exactly what admission should treat as healthy.
		r.SuggestedTargetWaiting = int(s.Waiting + 0.5)
		r.SuggestedKVCeiling = roundTo(s.KVUsage, 0.05)
	}
}

func runStep(ctx context.Context, client *http.Client, cfg Config, concurrency int) Step {
	step := Step{Concurrency: concurrency}
	ctx, cancel := context.WithTimeout(ctx, cfg.Duration)
	defer cancel()

	var (
		mu       sync.Mutex
		ttfts    []time.Duration
		tpots    []time.Duration
		tokens   int64
		failed   int
		firstErr string
	)

	sampler := newSampler(cfg.MetricsURL, cfg.Engine, cfg.MetricsKey)
	sampler.start(ctx)

	start := time.Now()
	var wg sync.WaitGroup
	for range concurrency {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for ctx.Err() == nil {
				r, err := one(ctx, client, cfg)
				mu.Lock()
				switch {
				case err != nil:
					// A cancelled context is the step ending, not a failure.
					if ctx.Err() == nil {
						failed++
						if firstErr == "" {
							firstErr = err.Error()
						}
					}
				default:
					ttfts = append(ttfts, r.ttft)
					if r.tpot > 0 {
						tpots = append(tpots, r.tpot)
					}
					tokens += r.tokens
				}
				mu.Unlock()
			}
		}()
	}
	wg.Wait()
	elapsed := time.Since(start)
	sampler.stop()

	step.Completed = len(ttfts)
	step.Failed = failed
	step.FirstError = firstErr
	step.TTFT = percentiles(ttfts)
	step.TPOT = percentiles(tpots)
	if elapsed > 0 {
		step.RequestsPerSec = float64(step.Completed) / elapsed.Seconds()
		step.TokensPerSec = float64(tokens) / elapsed.Seconds()
	}
	step.Waiting, step.Running, step.KVUsage = sampler.averages()
	step.MetricsUnread = sampler.unread()

	// A step that could not complete anything has not met anything.
	step.MeetsSLO = step.Completed > 0 && step.Failed == 0 &&
		cfg.SLO.Met(step.TTFT.P99, step.TPOT.P99)
	return step
}

type outcome struct {
	ttft   time.Duration
	tpot   time.Duration
	tokens int64
}

// one sends a single streaming completion and times it.
func one(ctx context.Context, client *http.Client, cfg Config) (outcome, error) {
	body, err := json.Marshal(map[string]any{
		"model":      cfg.Model,
		"stream":     true,
		"max_tokens": cfg.MaxTokens,
		"messages": []map[string]string{
			{"role": "user", "content": prompt(cfg)},
		},
	})
	if err != nil {
		return outcome{}, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		cfg.Target+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		return outcome{}, err
	}
	req.Header.Set("Content-Type", "application/json")
	if cfg.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	}

	start := time.Now()
	resp, err := client.Do(req)
	if err != nil {
		return outcome{}, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
		return outcome{}, fmt.Errorf("%s: %s", resp.Status, strings.TrimSpace(string(b)))
	}

	var out outcome
	var firstAt, lastAt time.Time
	sc := bufio.NewScanner(resp.Body)
	sc.Buffer(make([]byte, 0, 64*1024), 1<<20)
	for sc.Scan() {
		line := sc.Text()
		if !strings.HasPrefix(line, "data: ") || strings.HasSuffix(line, "[DONE]") {
			continue
		}
		// Only chunks that carry content count: a role-only first chunk is not
		// a token, and counting it would flatter the time to first token.
		if !hasContent(line[6:]) {
			continue
		}
		now := time.Now()
		if out.tokens == 0 {
			firstAt = now
			out.ttft = now.Sub(start)
		}
		lastAt = now
		out.tokens++
	}
	if err := sc.Err(); err != nil {
		return outcome{}, err
	}
	if out.tokens == 0 {
		return outcome{}, errors.New("the response carried no content")
	}
	if out.tokens > 1 {
		out.tpot = lastAt.Sub(firstAt) / time.Duration(out.tokens-1)
	}
	return out, nil
}

func hasContent(payload string) bool {
	var chunk struct {
		Choices []struct {
			Delta struct {
				Content string `json:"content"`
			} `json:"delta"`
		} `json:"choices"`
	}
	if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
		return false
	}
	for _, c := range chunk.Choices {
		if c.Delta.Content != "" {
			return true
		}
	}
	return false
}

// prompt builds a request body of roughly the configured size.
//
// Unique per request unless a shared prefix was asked for: repeating one prompt
// lets the engine's cache answer most of the prefill, and the result then
// describes the cache rather than the hardware.
//
// "Shared" has to mean byte-identical. Dropping only the random header while
// still picking words at random would leave every prompt different after the
// first few tokens, so the cache would miss and the flag would quietly measure
// the thing it exists to avoid measuring.
func prompt(cfg Config) string {
	var b strings.Builder
	if cfg.SharedPrefix {
		for i := 0; b.Len() < cfg.PromptTokens*4; i++ {
			b.WriteString(words[i%len(words)])
			b.WriteByte(' ')
		}
	} else {
		fmt.Fprintf(&b, "%016x %016x ", rand.Uint64(), rand.Uint64()) //nolint:gosec // not security
		// Roughly four characters to a token, the same crude ratio the router
		// uses; exactness is not what this number is for.
		for b.Len() < cfg.PromptTokens*4 {
			b.WriteString(words[rand.IntN(len(words))]) //nolint:gosec // not security
			b.WriteByte(' ')
		}
	}
	b.WriteString("\n\nSummarize the text above.")
	return b.String()
}

var words = strings.Fields(`the quick brown fox jumps over lazy dog while distant
thunder rolls across an empty valley and somewhere a door closes softly in the
dark house where nobody has lived for years although the garden still produces
apples every autumn which fall and rot beneath the trees unnoticed by anyone
except the birds that come each morning to pick at them before the frost`)

func percentiles(v []time.Duration) Percentiles {
	if len(v) == 0 {
		return Percentiles{}
	}
	s := make([]time.Duration, len(v))
	copy(s, v)
	sort.Slice(s, func(i, j int) bool { return s[i] < s[j] })
	at := func(p float64) time.Duration {
		i := int(float64(len(s)-1) * p)
		return s[i]
	}
	return Percentiles{P50: at(0.50), P90: at(0.90), P99: at(0.99), Max: s[len(s)-1]}
}

// roundTo rounds to a multiple of step, through integers.
//
// The obvious float arithmetic returns 0.7000000000000001 for what should be
// 0.70, which is harmless on a terminal and ugly in the JSON someone pastes
// into a configuration file.
func roundTo(v, step float64) float64 {
	if step <= 0 {
		return v
	}
	const scale = 10000
	units := int64(v/step + 0.5)
	return float64(int64(float64(units)*step*scale+0.5)) / scale
}