package convert

import (
	"encoding/json"
)

// TokenUsage 从上游响应提取的 token 用量。
// 三种协议的字段名不同，此处统一为：input / output / cache_read / cache_creation，
// 供用量记录与计费使用。
type TokenUsage struct {
	InputTokens         int
	OutputTokens        int
	CacheReadTokens     int
	CacheCreationTokens int
}

// has 判断是否真的拿到了非零用量（过滤掉没有 usage 字段的响应）。
func (u *TokenUsage) has() bool {
	return u.InputTokens > 0 || u.OutputTokens > 0 ||
		u.CacheReadTokens > 0 || u.CacheCreationTokens > 0
}

// mergeJSON 把一张 usage 对象并入累计值。proto 决定字段名（chat/responses 与 messages 不同）。
func (u *TokenUsage) mergeJSON(raw map[string]any, proto string) {
	switch proto {
	case ProtoChat, ProtoResponses:
		in, _ := raw["prompt_tokens"].(float64)
		out, _ := raw["completion_tokens"].(float64)
		if in == 0 && out == 0 {
			in, _ = raw["input_tokens"].(float64)
			out, _ = raw["output_tokens"].(float64)
		}
		u.InputTokens += int(in)
		u.OutputTokens += int(out)
		if d, ok := raw["prompt_tokens_details"].(map[string]any); ok {
			if c, _ := d["cached_tokens"].(float64); c > 0 {
				u.CacheReadTokens += int(c)
			}
		}
		if d, ok := raw["input_tokens_details"].(map[string]any); ok {
			if c, _ := d["cached_tokens"].(float64); c > 0 {
				u.CacheReadTokens += int(c)
			}
		}
	case ProtoMessages:
		in, _ := raw["input_tokens"].(float64)
		out, _ := raw["output_tokens"].(float64)
		u.InputTokens += int(in)
		u.OutputTokens += int(out)
		if c, _ := raw["cache_read_input_tokens"].(float64); c > 0 {
			u.CacheReadTokens += int(c)
		}
		if c, _ := raw["cache_creation_input_tokens"].(float64); c > 0 {
			u.CacheCreationTokens += int(c)
		}
	}
}

// ExtractUsageJSON 从完整非流式响应体中提取用量。proto 为上游协议。
// 返回 (用量, 是否有效)。
func ExtractUsageJSON(body []byte, proto string) (TokenUsage, bool) {
	var top map[string]any
	if err := json.Unmarshal(body, &top); err != nil {
		return TokenUsage{}, false
	}
	var u TokenUsage
	if usage, ok := top["usage"].(map[string]any); ok {
		u.mergeJSON(usage, proto)
	}
	return u, u.has()
}

// StreamUsageAccum 流式用量累计器。逐行喂入上游 SSE 的 data 载荷，
// 按协议分别取各事件里的 usage 字段（各事件只会携带一部分字段，取最大值合并）。
type StreamUsageAccum struct {
	u TokenUsage
}

// NewStreamUsageAccum 创建一个流式用量累计器。
func NewStreamUsageAccum() *StreamUsageAccum {
	return &StreamUsageAccum{}
}

// Feed 喂入一行 SSE data 载荷（不含 "data:" 前缀与换行）。
func (a *StreamUsageAccum) Feed(payload []byte, proto string) {
	var top map[string]any
	if json.Unmarshal(payload, &top) != nil {
		return
	}
	var t TokenUsage
	switch proto {
	case ProtoChat:
		if usage, ok := top["usage"].(map[string]any); ok {
			t.mergeJSON(usage, proto)
		}
	case ProtoResponses:
		// response.completed 事件把用量放在 response.usage 下。
		if resp, ok := top["response"].(map[string]any); ok {
			if usage, ok := resp["usage"].(map[string]any); ok {
				t.mergeJSON(usage, proto)
			}
		}
	case ProtoMessages:
		// message_start: {message: {usage: {input_tokens, cache_*}}}
		// message_delta: {usage: {output_tokens}}
		if msg, ok := top["message"].(map[string]any); ok {
			if usage, ok := msg["usage"].(map[string]any); ok {
				t.mergeJSON(usage, proto)
			}
		}
		if usage, ok := top["usage"].(map[string]any); ok {
			var t2 TokenUsage
			t2.mergeJSON(usage, proto)
			t.InputTokens = max(t.InputTokens, t2.InputTokens)
			t.OutputTokens = max(t.OutputTokens, t2.OutputTokens)
			t.CacheReadTokens = max(t.CacheReadTokens, t2.CacheReadTokens)
			t.CacheCreationTokens = max(t.CacheCreationTokens, t2.CacheCreationTokens)
		}
	default:
		return
	}
	a.u.InputTokens = max(a.u.InputTokens, t.InputTokens)
	a.u.OutputTokens = max(a.u.OutputTokens, t.OutputTokens)
	a.u.CacheReadTokens = max(a.u.CacheReadTokens, t.CacheReadTokens)
	a.u.CacheCreationTokens = max(a.u.CacheCreationTokens, t.CacheCreationTokens)
}

// Usage 返回当前累计用量。
func (a *StreamUsageAccum) Usage() TokenUsage {
	return a.u
}