package stats

import (
	"encoding/json"
	"time"
)

// Usage is the subset of the OpenAI chat-completion usage object the recorder
// persists. Numbers arrive from JSON unmarshal as float64.
type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
	ReasoningTokens  int `json:"reasoning_tokens"`
	CachedTokens     int `json:"cached_tokens"`

	// PromptTokensDetails.CachedTokens is emitted by providers that support
	// prompt caching (OpenAI/DeepSeek/Zhipu litellm gateways). Some upstreams
	// put cached_tokens at the top level instead.
	PromptTokensDetails struct {
		CachedTokens int `json:"cached_tokens"`
	} `json:"prompt_tokens_details"`

	// CompletionTokensDetails.ReasoningTokens is emitted by reasoning models;
	// some upstreams put reasoning_tokens at top level instead.
	CompletionTokensDetails struct {
		ReasoningTokens int `json:"reasoning_tokens"`
	} `json:"completion_tokens_details"`
}

// ExtractUsage decodes a raw usage value (as produced by encoding/json into an
// any) into token counts. It accepts both full usage maps and raw JSON bytes.
// Missing fields default to 0; a nil v yields zero usage.
func ExtractUsage(v any) Usage {
	var u Usage
	if v == nil {
		return u
	}
	switch t := v.(type) {
	case []byte:
		_ = json.Unmarshal(t, &u)
	case json.RawMessage:
		_ = json.Unmarshal(t, &u)
	case map[string]any:
		// Re-marshal + unmarshal is the simplest robust path for nested
		// *_tokens_details; usage payloads are tiny.
		if b, err := json.Marshal(t); err == nil {
			_ = json.Unmarshal(b, &u)
		}
	}
	if u.ReasoningTokens == 0 {
		u.ReasoningTokens = u.CompletionTokensDetails.ReasoningTokens
	}
	if u.CachedTokens == 0 {
		u.CachedTokens = u.PromptTokensDetails.CachedTokens
	}
	return u
}

// RecordFromUsage builds a Record from a captured usage value plus context.
func RecordFromUsage(model string, stream bool, usage any, status string, start time.Time) Record {
	u := ExtractUsage(usage)
	if u.TotalTokens == 0 && (u.PromptTokens != 0 || u.CompletionTokens != 0) {
		u.TotalTokens = u.PromptTokens + u.CompletionTokens
	}
	return Record{
		Ts:               time.Now(),
		Model:            model,
		Stream:           stream,
		PromptTokens:     u.PromptTokens,
		CompletionTokens: u.CompletionTokens,
		TotalTokens:      u.TotalTokens,
		ReasoningTokens:  u.ReasoningTokens,
		CachedTokens:     u.CachedTokens,
		Status:           status,
		LatencyMs:        time.Since(start).Milliseconds(),
	}
}