package bench

import (
	"context"
	"net/http"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/scrape"
)

// sampler reads an engine's metrics during a step.
//
// It is what turns a latency measurement into something admission can use: the
// queue depth and cache pressure observed while the SLO was still being met are
// the values target_waiting and kv_ceiling should have, rather than the round
// numbers somebody typed.
//
// Optional, because a run pointed at the gateway has no single engine to read.
type sampler struct {
	url    string
	apiKey string
	engine string
	client *http.Client
	cancel context.CancelFunc

	mu      sync.Mutex
	wg      sync.WaitGroup
	waiting float64
	running float64
	kv      float64
	n       int
}

func newSampler(metricsURL, engine string) *sampler {
	return &sampler{
		url:    metricsURL,
		engine: engine,
		client: &http.Client{Timeout: 2 * time.Second},
	}
}

func (s *sampler) start(ctx context.Context) {
	if s.url == "" {
		return
	}
	ctx, s.cancel = context.WithCancel(ctx)
	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		t := time.NewTicker(time.Second)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				s.sample(ctx)
			}
		}
	}()
}

func (s *sampler) sample(ctx context.Context) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.url, nil)
	if err != nil {
		return
	}
	if s.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+s.apiKey)
	}
	resp, err := s.client.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return
	}

	// The same mapping the router reads with, so a bench result cannot disagree
	// with what the scheduler sees. The engine has to be named: reading a
	// llama.cpp endpoint with vLLM's names finds nothing, and the run would
	// report a queue depth of zero and suggest thresholds built on it.
	stats, err := scrape.Parse(resp.Body, scrape.MappingFor(s.engine))
	if err != nil {
		return
	}
	s.mu.Lock()
	s.waiting += float64(stats.Waiting)
	s.running += float64(stats.Running)
	s.kv += stats.KVCacheUsage
	s.n++
	s.mu.Unlock()
}

func (s *sampler) stop() {
	if s.cancel != nil {
		s.cancel()
	}
	s.wg.Wait()
}

func (s *sampler) averages() (waiting, running, kv float64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.n == 0 {
		return 0, 0, 0
	}
	n := float64(s.n)
	return s.waiting / n, s.running / n, s.kv / n
}