package llm

import (
	"sort"
	"strings"
)

// Price is per-million-token list pricing.
type Price struct {
	InputPerMTok  float64
	OutputPerMTok float64
}

// Cache multipliers relative to the input rate.
const (
	cacheReadMultiplier  = 0.1
	cacheWriteMultiplier = 1.25
)

// prices covers the models this harness is likely to be pointed at. A model
// that is not listed reports no cost rather than a wrong one.
var prices = map[string]Price{
	"claude-opus-5":    {InputPerMTok: 5, OutputPerMTok: 25},
	"claude-opus-4-8":  {InputPerMTok: 5, OutputPerMTok: 25},
	"claude-opus-4-7":  {InputPerMTok: 5, OutputPerMTok: 25},
	"claude-sonnet-5":  {InputPerMTok: 2, OutputPerMTok: 10},
	"claude-haiku-4-5": {InputPerMTok: 1, OutputPerMTok: 5},
	"claude-fable-5-1": {InputPerMTok: 10, OutputPerMTok: 50},
	"claude-fable-5":   {InputPerMTok: 10, OutputPerMTok: 50},
}

// Cost prices one usage record. The boolean reports whether the model's
// pricing is known; a made-up number in METRICS.md is worse than a blank.
func (u Usage) Cost(model string) (float64, bool) {
	p, ok := prices[basePrice(model)]
	if !ok {
		return 0, false
	}
	return (float64(u.InputTokens)*p.InputPerMTok +
		float64(u.OutputTokens)*p.OutputPerMTok +
		float64(u.CacheReadInputTokens)*p.InputPerMTok*cacheReadMultiplier +
		float64(u.CacheCreationInputTokens)*p.InputPerMTok*cacheWriteMultiplier) / 1_000_000, true
}

// PriceOf reports the list pricing for a model, if this harness knows it.
// The picker shows what it knows and says nothing about what it does not.
func PriceOf(model string) (Price, bool) {
	p, ok := prices[basePrice(model)]
	return p, ok
}

// basePrice strips the decoration the API puts on a model id — a bracketed
// variant tag, a -latest or -YYYYMMDD suffix — leaving the name the price is
// listed under. A dated id and a context-window variant are the same model
// billed at the same rate, and /v1/models hands out both.
func basePrice(model string) string {
	// Gateways decorate the id: a routing prefix (openrouter/anthropic/...),
	// a Bedrock namespace (anthropic.claude-opus-5), a Vertex version suffix
	// (claude-opus-5@20260401). They all name the same model, and a lookup
	// that cannot see past them prices nothing and, worse, answers a
	// capability question wrongly for every request through a proxy.
	if i := strings.LastIndexByte(model, '/'); i >= 0 {
		model = model[i+1:]
	}
	model = strings.TrimPrefix(model, "anthropic.")
	if i := strings.IndexByte(model, '@'); i >= 0 {
		model = model[:i]
	}
	if i := strings.IndexByte(model, '['); i >= 0 {
		model = model[:i]
	}
	model = strings.TrimSuffix(model, "-latest")
	if i := strings.LastIndexByte(model, '-'); i >= 0 {
		if suffix := model[i+1:]; len(suffix) == 8 && allDigits(suffix) {
			model = model[:i]
		}
	}
	return model
}

func allDigits(s string) bool {
	for _, r := range s {
		if r < '0' || r > '9' {
			return false
		}
	}
	return len(s) > 0
}

// PricedModels are the ids this build has list pricing for, flagship first and
// cheapest last within a family.
//
// It exists for the setup wizard, which has to offer a usable list even when
// the endpoint will not serve /v1/models — a gateway that only proxies
// /v1/messages is the common case, and a first run cannot be left with nothing
// to pick. The set it can price is the set it can fully describe, which is
// what makes this the honest fallback rather than a second hardcoded list.
func PricedModels() []string {
	out := make([]string, 0, len(prices))
	for id := range prices {
		out = append(out, id)
	}
	// Map order is random; a picker has to look the same every time.
	sort.Slice(out, func(i, j int) bool {
		if ri, rj := familyRank(out[i]), familyRank(out[j]); ri != rj {
			return ri < rj
		}
		return out[i] > out[j] // newer version first within a family
	})
	return out
}

// familyRank orders the families the way someone choosing between them does:
// the flagship first, the cheap one last, anything unrecognised after both.
func familyRank(id string) int {
	for i, family := range []string{"claude-opus", "claude-sonnet", "claude-haiku", "claude-fable"} {
		if strings.HasPrefix(id, family) {
			return i
		}
	}
	// Unrecognised: after every known family, and left to the id tiebreak to
	// order among themselves.
	return len(prices)
}