package web

import (
	"context"
	"sort"
	"time"
)

// modelEfficiency is what a model gets out of the GPUs it holds, over the
// last quarter hour.
type modelEfficiency struct {
	Model string
	// TokensPerSec is output tokens a second; GPUs the cards it holds now.
	TokensPerSec float64
	GPUs         float64
	// Watts is what those cards draw; Util their average utilisation.
	Watts, Util float64
	// PerGPU and PerWh are the two ways of saying how much work the cards
	// do: output tokens a second per card, and per watt-hour drawn.
	PerGPU, PerWh float64
}

// The queries behind the table. Each model's cards are found by joining the
// GPU series to llmgw_gpu_assigned on the card, which carries the model.
var efficiencyQueries = []struct {
	expr string
	into func(*modelEfficiency, float64)
}{
	{`sum by (model) (rate(llmgw_tokens_total{direction="completion"}[15m]))`,
		func(e *modelEfficiency, v float64) { e.TokensPerSec = v }},
	{`count by (model) (llmgw_gpu_assigned)`,
		func(e *modelEfficiency, v float64) { e.GPUs = v }},
	{`sum by (model) (llmgw_gpu_power_watts * on (node, gpu) group_left (model) llmgw_gpu_assigned)`,
		func(e *modelEfficiency, v float64) { e.Watts = v }},
	{`avg by (model) (llmgw_gpu_utilization_percent * on (node, gpu) group_left (model) llmgw_gpu_assigned)`,
		func(e *modelEfficiency, v float64) { e.Util = v }},
}

// loadEfficiency is the table, best per card first; nil when the metrics
// store cannot be read or no model holds a GPU.
func loadEfficiency(ctx context.Context, h History) []modelEfficiency {
	if h == nil {
		return nil
	}
	byModel := map[string]*modelEfficiency{}
	for _, q := range efficiencyQueries {
		// The last of a short range: what History offers, and the one point
		// the table needs.
		series, err := h.RangeQuery(ctx, q.expr, "model", 10*time.Minute, 5*time.Minute)
		if err != nil {
			return nil
		}
		for _, s := range series {
			if len(s.Points) == 0 {
				continue
			}
			e := byModel[s.Name]
			if e == nil {
				e = &modelEfficiency{Model: s.Name}
				byModel[s.Name] = e
			}
			q.into(e, s.Last())
		}
	}
	var out []modelEfficiency
	for _, e := range byModel {
		if e.GPUs == 0 {
			continue // a model on the CPU; this table is about cards
		}
		e.PerGPU = e.TokensPerSec / e.GPUs
		if e.Watts > 0 {
			e.PerWh = e.TokensPerSec * 3600 / e.Watts
		}
		out = append(out, *e)
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].PerGPU != out[j].PerGPU {
			return out[i].PerGPU > out[j].PerGPU
		}
		return out[i].Model < out[j].Model
	})
	return out
}