package proxy

import "encoding/json"

// usageEnvelope mirrors just the usage block of an OpenAI response.
// Everything else in the payload is passed through untouched, so only these
// fields are ever decoded.
type usageEnvelope struct {
	Usage *struct {
		PromptTokens        int `json:"prompt_tokens"`
		CompletionTokens    int `json:"completion_tokens"`
		TotalTokens         int `json:"total_tokens"`
		PromptTokensDetails *struct {
			CachedTokens int `json:"cached_tokens"`
		} `json:"prompt_tokens_details"`
	} `json:"usage"`
}

// usageFromJSON extracts token usage from a response or an SSE chunk.
// It returns nil when the payload carries no usage block, which is the normal
// case for every chunk but the last.
func usageFromJSON(b []byte) *Usage {
	var env usageEnvelope
	if err := json.Unmarshal(b, &env); err != nil || env.Usage == nil {
		return nil
	}
	u := &Usage{
		PromptTokens:     env.Usage.PromptTokens,
		CompletionTokens: env.Usage.CompletionTokens,
		TotalTokens:      env.Usage.TotalTokens,
	}
	if d := env.Usage.PromptTokensDetails; d != nil {
		u.CachedTokens = d.CachedTokens
	}
	if u.TotalTokens == 0 {
		u.TotalTokens = u.PromptTokens + u.CompletionTokens
	}
	return u
}