package pricing

import "strings"

// 价格单位：美元 / 1M tokens
type ModelPrice struct {
	InputPer1M  float64
	OutputPer1M float64
}

var catalog = map[string]ModelPrice{
	"gpt-4o":              {InputPer1M: 2.5, OutputPer1M: 10},
	"gpt-4o-mini":         {InputPer1M: 0.15, OutputPer1M: 0.6},
	"gpt-4-turbo":         {InputPer1M: 10, OutputPer1M: 30},
	"gpt-3.5-turbo":       {InputPer1M: 0.5, OutputPer1M: 1.5},
	"claude-3-5-sonnet":   {InputPer1M: 3, OutputPer1M: 15},
	"claude-3-5-haiku":    {InputPer1M: 0.8, OutputPer1M: 4},
	"claude-3-opus":       {InputPer1M: 15, OutputPer1M: 75},
	"text-embedding-3-small": {InputPer1M: 0.02, OutputPer1M: 0},
	"text-embedding-3-large": {InputPer1M: 0.13, OutputPer1M: 0},
	"text-embedding-ada-002": {InputPer1M: 0.1, OutputPer1M: 0},
	"bge-reranker-v2-m3":    {InputPer1M: 0.1, OutputPer1M: 0},
	"rerank-english-v3.0":   {InputPer1M: 2.0, OutputPer1M: 0},
	"rerank-multilingual-v3.0": {InputPer1M: 2.0, OutputPer1M: 0},
}

func Estimate(model string, tokensIn, tokensOut int, endpoint string) float64 {
	if strings.Contains(endpoint, "embeddings") || strings.Contains(endpoint, "rerank") {
		tokensOut = 0
	}
	price, ok := lookup(model)
	if !ok {
		// fallback: average small model pricing
		price = ModelPrice{InputPer1M: 0.5, OutputPer1M: 1.5}
	}
	inCost := float64(tokensIn) / 1_000_000 * price.InputPer1M
	outCost := float64(tokensOut) / 1_000_000 * price.OutputPer1M
	return inCost + outCost
}

func lookup(model string) (ModelPrice, bool) {
	if p, ok := catalog[model]; ok {
		return p, true
	}
	// strip provider prefix
	if i := strings.LastIndex(model, "/"); i >= 0 {
		if p, ok := catalog[model[i+1:]]; ok {
			return p, true
		}
	}
	// prefix match
	for k, v := range catalog {
		if strings.HasPrefix(model, k) {
			return v, true
		}
	}
	return ModelPrice{}, false
}

func ListCatalog() map[string]ModelPrice {
	out := make(map[string]ModelPrice, len(catalog))
	for k, v := range catalog {
		out[k] = v
	}
	return out
}