codex-helper/backend/internal/cliproxy/client.go

package cliproxy

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const (
	usageURL        = "https://chatgpt.com/backend-api/wham/usage"
	profileURL      = "https://chatgpt.com/backend-api/wham/profiles/me"
	resetCreditsURL = "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits"
	codexUserAgent  = "codex_cli_rs/0.76.0 (Debian 13.0.0; x86_64) WindowsTerminal"
	maxResponseBody = 4 << 20
)

var ErrNotConfigured = errors.New("CLIProxyAPI Management API 未配置")

type Client struct {
	baseURL       string
	managementKey string
	httpClient    *http.Client
	now           func() time.Time
}

type AuthFile struct {
	AuthIndex string
	Label     string
	Name      string
	Email     *string
	Status    string
	AccountID string
	PlanType  *string
}

type Limit struct {
	LimitID               string
	LimitName             *string
	WindowType            string
	UsedPercent           float64
	WindowDurationMinutes int
	ResetsAt              int64
	PlanType              *string
}

type UsageSummary struct {
	LifetimeTokens        *int64
	PeakDailyTokens       *int64
	LongestRunningTurnSec *int64
	CurrentStreakDays     *int
	LongestStreakDays     *int
}

type UsagePoint struct {
	Date        string
	TotalTokens int64
}

type ResetCredits struct {
	AvailableCount int
	ExpiresAt      []int64
}

type Snapshot struct {
	Auth             AuthFile
	Limits           []Limit
	Summary          UsageSummary
	Usage            []UsagePoint
	ProfileAvailable bool
	UsageAvailable   bool
	ResetCredits     *ResetCredits
	FetchedAt        time.Time
}

func New(baseURL, managementKey string) *Client {
	return NewWithHTTPClient(baseURL, managementKey, &http.Client{Timeout: 15 * time.Second})
}

func NewWithHTTPClient(baseURL, managementKey string, httpClient *http.Client) *Client {
	if httpClient == nil {
		httpClient = &http.Client{Timeout: 15 * time.Second}
	}
	return &Client{
		baseURL:       strings.TrimRight(strings.TrimSpace(baseURL), "/"),
		managementKey: managementKey,
		httpClient:    httpClient,
		now:           time.Now,
	}
}

func (c *Client) Configured() bool {
	return c.baseURL != "" && c.managementKey != ""
}

func (c *Client) Auth(ctx context.Context, authIndex string) (AuthFile, error) {
	if !c.Configured() {
		return AuthFile{}, ErrNotConfigured
	}
	authIndex = strings.TrimSpace(authIndex)
	if authIndex == "" {
		return AuthFile{}, errors.New("authIndex 不能为空")
	}
	u, err := url.Parse(c.baseURL + "/v0/management/auth-files")
	if err != nil {
		return AuthFile{}, errors.New("CLIProxyAPI 地址无效")
	}
	q := u.Query()
	q.Set("auth_index", authIndex)
	u.RawQuery = q.Encode()
	body, err := c.managementRequest(ctx, http.MethodGet, u.String(), nil)
	if err != nil {
		return AuthFile{}, err
	}
	var raw any
	if err := json.Unmarshal(body, &raw); err != nil {
		return AuthFile{}, errors.New("CLIProxyAPI auth-files 响应格式错误")
	}
	items := authItems(raw)
	matches := make([]AuthFile, 0, 1)
	for _, item := range items {
		auth, ok := parseAuthFile(item, authIndex)
		if ok {
			matches = append(matches, auth)
		}
	}
	if len(matches) == 0 {
		return AuthFile{}, errors.New("未找到可用的 Codex auth")
	}
	if len(matches) != 1 {
		return AuthFile{}, errors.New("Codex auth 匹配结果不唯一")
	}
	return matches[0], nil
}

func (c *Client) Snapshot(ctx context.Context, authIndex string) (Snapshot, error) {
	auth, err := c.Auth(ctx, authIndex)
	if err != nil {
		return Snapshot{}, err
	}
	receivedAt := c.now().UTC()
	usageBody, err := c.apiCall(ctx, auth, usageURL)
	if err != nil {
		return Snapshot{}, fmt.Errorf("获取用量限额失败: %w", err)
	}
	limits, availableFallback, effectivePlan, err := parseUsage(usageBody, auth.PlanType, receivedAt)
	if err != nil {
		return Snapshot{}, err
	}
	auth.PlanType = effectivePlan
	var summary UsageSummary
	usage := []UsagePoint{}
	profileAvailable := false
	usageAvailable := false
	if profileBody, profileErr := c.apiCall(ctx, auth, profileURL); profileErr == nil {
		if parsedSummary, parsedUsage, parsedUsageAvailable, parseErr := parseProfile(profileBody); parseErr == nil {
			summary = parsedSummary
			usage = parsedUsage
			usageAvailable = parsedUsageAvailable
			profileAvailable = usageAvailable || usageSummaryAvailable(summary)
		}
	}
	var resetCredits *ResetCredits
	resetCtx, cancelReset := context.WithTimeout(ctx, 5*time.Second)
	resetBody, resetErr := c.apiCall(resetCtx, auth, resetCreditsURL)
	cancelReset()
	if resetErr == nil {
		resetCredits = parseResetCredits(resetBody, availableFallback)
	} else if availableFallback > 0 {
		resetCredits = &ResetCredits{AvailableCount: availableFallback, ExpiresAt: []int64{}}
	}
	return Snapshot{
		Auth:             auth,
		Limits:           limits,
		Summary:          summary,
		Usage:            usage,
		ProfileAvailable: profileAvailable,
		UsageAvailable:   usageAvailable,
		ResetCredits:     resetCredits,
		FetchedAt:        receivedAt,
	}, nil
}

func (c *Client) managementRequest(ctx context.Context, method, endpoint string, body io.Reader) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, method, endpoint, body)
	if err != nil {
		return nil, errors.New("创建 CLIProxyAPI 请求失败")
	}
	req.Header.Set("Authorization", "Bearer "+c.managementKey)
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	httpClient := *c.httpClient
	httpClient.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
		return http.ErrUseLastResponse
	}
	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, errors.New("CLIProxyAPI 请求失败")
	}
	defer resp.Body.Close()
	payload, err := readLimited(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("CLIProxyAPI Management API 返回状态码 %d", resp.StatusCode)
	}
	return payload, nil
}

func (c *Client) apiCall(ctx context.Context, auth AuthFile, upstreamURL string) ([]byte, error) {
	requestBody := struct {
		AuthIndex string            `json:"auth_index"`
		Method    string            `json:"method"`
		URL       string            `json:"url"`
		Header    map[string]string `json:"header"`
	}{
		AuthIndex: auth.AuthIndex,
		Method:    http.MethodGet,
		URL:       upstreamURL,
		Header: map[string]string{
			"Authorization":      "Bearer $TOKEN$",
			"Chatgpt-Account-Id": auth.AccountID,
			"Accept":             "application/json",
			"Content-Type":       "application/json",
			"User-Agent":         codexUserAgent,
		},
	}
	encoded, err := json.Marshal(requestBody)
	if err != nil {
		return nil, errors.New("创建 CLIProxyAPI api-call 请求失败")
	}
	payload, err := c.managementRequest(ctx, http.MethodPost, c.baseURL+"/v0/management/api-call", strings.NewReader(string(encoded)))
	if err != nil {
		return nil, err
	}
	var envelope struct {
		StatusCode int             `json:"status_code"`
		Body       json.RawMessage `json:"body"`
	}
	if err := json.Unmarshal(payload, &envelope); err != nil || envelope.StatusCode == 0 || len(envelope.Body) == 0 {
		return nil, errors.New("CLIProxyAPI api-call 响应格式错误")
	}
	if envelope.StatusCode < 200 || envelope.StatusCode >= 300 {
		if errorType := safeUpstreamErrorType(envelope.Body); errorType != "" {
			return nil, fmt.Errorf("CLIProxyAPI 上游请求返回状态码 %d (%s)", envelope.StatusCode, errorType)
		}
		return nil, fmt.Errorf("CLIProxyAPI 上游请求返回状态码 %d", envelope.StatusCode)
	}
	var bodyString string
	if len(envelope.Body) > 0 && envelope.Body[0] == '"' {
		if err := json.Unmarshal(envelope.Body, &bodyString); err != nil {
			return nil, errors.New("CLIProxyAPI api-call body 格式错误")
		}
		if len(bodyString) > maxResponseBody {
			return nil, errors.New("CLIProxyAPI 响应体过大")
		}
		return []byte(bodyString), nil
	}
	if len(envelope.Body) > maxResponseBody {
		return nil, errors.New("CLIProxyAPI 响应体过大")
	}
	return envelope.Body, nil
}

func safeUpstreamErrorType(raw json.RawMessage) string {
	body := []byte(raw)
	if len(body) > 0 && body[0] == '"' {
		var text string
		if json.Unmarshal(body, &text) != nil {
			return ""
		}
		body = []byte(text)
	}
	var payload map[string]any
	if json.Unmarshal(body, &payload) != nil {
		return ""
	}
	candidates := []string{stringValue(payload, "type", "code")}
	if nested, ok := objectValue(payload, "error"); ok {
		candidates = append([]string{stringValue(nested, "type", "code")}, candidates...)
	}
	for _, candidate := range candidates {
		candidate = strings.TrimSpace(candidate)
		if candidate == "" || len(candidate) > 80 {
			continue
		}
		safe := true
		for _, r := range candidate {
			if !(r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '_' || r == '-' || r == '.') {
				safe = false
				break
			}
		}
		if safe {
			return candidate
		}
	}
	return ""
}

func readLimited(r io.Reader) ([]byte, error) {
	payload, err := io.ReadAll(io.LimitReader(r, maxResponseBody+1))
	if err != nil {
		return nil, errors.New("读取 CLIProxyAPI 响应失败")
	}
	if len(payload) > maxResponseBody {
		return nil, errors.New("CLIProxyAPI 响应体过大")
	}
	return payload, nil
}

func authItems(raw any) []map[string]any {
	switch value := raw.(type) {
	case []any:
		out := make([]map[string]any, 0, len(value))
		for _, item := range value {
			if object, ok := item.(map[string]any); ok {
				out = append(out, object)
			}
		}
		return out
	case map[string]any:
		for _, key := range []string{"auth_files", "authFiles", "files", "data", "items"} {
			if nested, ok := value[key]; ok {
				if out := authItems(nested); len(out) > 0 {
					return out
				}
			}
		}
		if stringValue(value, "auth_index", "authIndex") != "" {
			return []map[string]any{value}
		}
	}
	return []map[string]any{}
}

func parseAuthFile(item map[string]any, expectedIndex string) (AuthFile, bool) {
	index := strings.TrimSpace(stringValue(item, "auth_index", "authIndex"))
	if index != expectedIndex {
		return AuthFile{}, false
	}
	provider := strings.ToLower(strings.TrimSpace(stringValue(item, "provider")))
	typ := strings.ToLower(strings.TrimSpace(stringValue(item, "type")))
	if provider != "" && provider != "codex" || typ != "" && typ != "codex" || provider == "" && typ == "" {
		return AuthFile{}, false
	}
	status := strings.TrimSpace(stringValue(item, "status"))
	statusLower := strings.ToLower(status)
	// Do not reject CPA's transient unavailable/error state: quota exhaustion
	// itself can set it, and this dashboard still needs to read the reset time.
	if boolValue(item, "disabled") || statusLower == "disabled" {
		return AuthFile{}, false
	}
	claims := claimsFrom(item)
	authClaims, _ := objectValue(claims, "https://api.openai.com/auth")
	accountID := strings.TrimSpace(stringValue(claims, "chatgpt_account_id", "chatgptAccountId"))
	if accountID == "" {
		accountID = strings.TrimSpace(stringValue(authClaims, "chatgpt_account_id", "chatgptAccountId"))
	}
	if accountID == "" {
		return AuthFile{}, false
	}
	var email *string
	if value := strings.TrimSpace(stringValue(item, "email")); value != "" {
		email = &value
	} else if value := strings.TrimSpace(stringValue(claims, "email")); value != "" {
		email = &value
	}
	var planType *string
	value := strings.TrimSpace(stringValue(claims, "plan_type", "planType"))
	if value == "" {
		value = strings.TrimSpace(stringValue(authClaims, "chatgpt_plan_type", "plan_type", "planType"))
	}
	if value != "" {
		planType = &value
	}
	return AuthFile{
		AuthIndex: index,
		Label:     strings.TrimSpace(stringValue(item, "label")),
		Name:      strings.TrimSpace(stringValue(item, "name")),
		Email:     email,
		Status:    status,
		AccountID: accountID,
		PlanType:  planType,
	}, true
}

func claimsFrom(item map[string]any) map[string]any {
	for _, key := range []string{"id_token_claims", "idTokenClaims", "id_token", "idToken"} {
		value, ok := item[key]
		if !ok {
			continue
		}
		if object, ok := value.(map[string]any); ok {
			if nested, ok := object["claims"].(map[string]any); ok {
				return nested
			}
			return object
		}
		if encoded, ok := value.(string); ok {
			if claims := decodeClaimsString(encoded); claims != nil {
				return claims
			}
		}
	}
	return map[string]any{}
}

func decodeClaimsString(value string) map[string]any {
	var claims map[string]any
	if json.Unmarshal([]byte(value), &claims) == nil {
		return claims
	}
	parts := strings.Split(value, ".")
	if len(parts) < 2 {
		return nil
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil || json.Unmarshal(payload, &claims) != nil {
		return nil
	}
	return claims
}

func parseUsage(body []byte, fallbackPlan *string, receivedAt time.Time) ([]Limit, int, *string, error) {
	var root map[string]any
	if err := json.Unmarshal(body, &root); err != nil {
		return nil, 0, nil, errors.New("wham usage 响应格式错误")
	}
	if data, ok := objectValue(root, "data"); ok {
		root = data
	}
	planType := fallbackPlan
	if value := strings.TrimSpace(stringValue(root, "plan_type", "planType")); value != "" {
		planType = &value
	}
	limits := make([]Limit, 0)
	if raw, ok := objectValue(root, "rate_limit", "rateLimit"); ok {
		limits = append(limits, parseLimit("codex", nil, raw, planType, receivedAt)...)
	}
	if raw, ok := objectValue(root, "code_review_rate_limit", "codeReviewRateLimit"); ok {
		limits = append(limits, parseLimit("code_review", nil, raw, planType, receivedAt)...)
	}
	for _, raw := range arrayValue(root, "additional_rate_limits", "additionalRateLimits") {
		object, ok := raw.(map[string]any)
		if !ok {
			continue
		}
		limitID := strings.TrimSpace(stringValue(object, "metered_feature", "meteredFeature"))
		if limitID == "" {
			continue
		}
		var name *string
		if value := strings.TrimSpace(stringValue(object, "limit_name", "limitName")); value != "" {
			name = &value
		}
		rateLimit := object
		if nested, ok := objectValue(object, "rate_limit", "rateLimit"); ok {
			rateLimit = nested
		}
		limits = append(limits, parseLimit(limitID, name, rateLimit, planType, receivedAt)...)
	}
	availableCount := 0
	if credits, ok := objectValue(root, "rate_limit_reset_credits", "rateLimitResetCredits"); ok {
		availableCount = int(numberValue(credits, "available_count", "availableCount"))
	}
	return limits, availableCount, planType, nil
}

func parseLimit(limitID string, name *string, raw map[string]any, planType *string, receivedAt time.Time) []Limit {
	out := make([]Limit, 0, 2)
	for _, window := range []struct {
		kind string
		keys []string
	}{{"primary", []string{"primary_window", "primaryWindow"}}, {"secondary", []string{"secondary_window", "secondaryWindow"}}} {
		object, ok := objectValue(raw, window.keys...)
		if !ok {
			continue
		}
		used, usedOK := numericValue(object, "used_percent", "usedPercent")
		if !usedOK {
			continue
		}
		seconds := int64(numberValue(object, "limit_window_seconds", "limitWindowSeconds"))
		resetsAt := unixTimeValue(object, "reset_at", "resetAt")
		if resetsAt == 0 {
			resetAfter := int64(numberValue(object, "reset_after_seconds", "resetAfterSeconds"))
			if resetAfter > 0 {
				resetsAt = receivedAt.Unix() + resetAfter
			}
		}
		used = math.Max(0, math.Min(100, used))
		out = append(out, Limit{
			LimitID:               limitID,
			LimitName:             name,
			WindowType:            window.kind,
			UsedPercent:           used,
			WindowDurationMinutes: int((seconds + 59) / 60),
			ResetsAt:              resetsAt,
			PlanType:              planType,
		})
	}
	return out
}

func parseProfile(body []byte) (UsageSummary, []UsagePoint, bool, error) {
	var root map[string]any
	if err := json.Unmarshal(body, &root); err != nil {
		return UsageSummary{}, nil, false, errors.New("wham profile 响应格式错误")
	}
	stats := root
	nestedStats := false
	if nested, ok := objectValue(root, "stats", "usage_stats", "usageStats"); ok {
		stats = nested
		nestedStats = true
	}
	summary := UsageSummary{
		LifetimeTokens:        nonNegativeInt64Pointer(stats, "lifetime_tokens", "lifetimeTokens"),
		PeakDailyTokens:       nonNegativeInt64Pointer(stats, "peak_daily_tokens", "peakDailyTokens"),
		LongestRunningTurnSec: nonNegativeInt64Pointer(stats, "longest_running_turn_sec", "longestRunningTurnSec"),
		CurrentStreakDays:     nonNegativeIntPointer(stats, "current_streak_days", "currentStreakDays"),
		LongestStreakDays:     nonNegativeIntPointer(stats, "longest_streak_days", "longestStreakDays"),
	}
	buckets, usageAvailable := arrayField(root, "daily_usage_buckets", "dailyUsageBuckets")
	if !usageAvailable && nestedStats {
		buckets, usageAvailable = arrayField(stats, "daily_usage_buckets", "dailyUsageBuckets")
	}
	usage := make([]UsagePoint, 0, len(buckets))
	for _, raw := range buckets {
		bucket, ok := raw.(map[string]any)
		if !ok {
			continue
		}
		date := strings.TrimSpace(stringValue(bucket, "date", "start_date", "startDate"))
		if len(date) >= 10 {
			date = date[:10]
		}
		if _, err := time.Parse("2006-01-02", date); err != nil {
			continue
		}
		tokens, ok := numericValue(bucket, "total_tokens", "totalTokens", "tokens")
		if !ok || tokens < 0 {
			continue
		}
		usage = append(usage, UsagePoint{Date: date, TotalTokens: int64(tokens)})
	}
	return summary, usage, usageAvailable, nil
}

func usageSummaryAvailable(summary UsageSummary) bool {
	return summary.LifetimeTokens != nil || summary.PeakDailyTokens != nil || summary.LongestRunningTurnSec != nil || summary.CurrentStreakDays != nil || summary.LongestStreakDays != nil
}

func parseResetCredits(body []byte, fallback int) *ResetCredits {
	var root map[string]any
	if json.Unmarshal(body, &root) != nil {
		if fallback > 0 {
			return &ResetCredits{AvailableCount: fallback, ExpiresAt: []int64{}}
		}
		return nil
	}
	if data, ok := objectValue(root, "data"); ok {
		root = data
	}
	countPresent := hasValue(root, "available_count", "availableCount")
	count := int(numberValue(root, "available_count", "availableCount"))
	if countPresent && count <= 0 {
		return nil
	}
	if !countPresent {
		count = fallback
	}
	if count <= 0 {
		return nil
	}
	expiresAt := make([]int64, 0)
	for _, raw := range arrayValue(root, "credits", "available_credits", "availableCredits") {
		credit, ok := raw.(map[string]any)
		if !ok {
			continue
		}
		if status := strings.ToLower(strings.TrimSpace(stringValue(credit, "status"))); status != "" && status != "available" {
			continue
		}
		if value := unixTimeValue(credit, "expires_at", "expiresAt"); value > 0 {
			expiresAt = append(expiresAt, value)
		}
	}
	return &ResetCredits{AvailableCount: count, ExpiresAt: uniqueSorted(expiresAt)}
}

func uniqueSorted(values []int64) []int64 {
	for i := 0; i < len(values); i++ {
		for j := i + 1; j < len(values); j++ {
			if values[j] < values[i] {
				values[i], values[j] = values[j], values[i]
			}
		}
	}
	out := values[:0]
	for _, value := range values {
		if len(out) == 0 || out[len(out)-1] != value {
			out = append(out, value)
		}
	}
	return out
}

func objectValue(object map[string]any, keys ...string) (map[string]any, bool) {
	for _, key := range keys {
		if value, ok := object[key].(map[string]any); ok {
			return value, true
		}
	}
	return nil, false
}

func arrayValue(object map[string]any, keys ...string) []any {
	value, _ := arrayField(object, keys...)
	return value
}

func arrayField(object map[string]any, keys ...string) ([]any, bool) {
	for _, key := range keys {
		if value, ok := object[key].([]any); ok {
			return value, true
		}
	}
	return []any{}, false
}

func stringValue(object map[string]any, keys ...string) string {
	for _, key := range keys {
		value, ok := object[key]
		if !ok || value == nil {
			continue
		}
		switch typed := value.(type) {
		case string:
			return typed
		case json.Number:
			return typed.String()
		case float64:
			if typed == math.Trunc(typed) {
				return strconv.FormatInt(int64(typed), 10)
			}
			return strconv.FormatFloat(typed, 'f', -1, 64)
		}
	}
	return ""
}

func hasValue(object map[string]any, keys ...string) bool {
	for _, key := range keys {
		if value, ok := object[key]; ok && value != nil {
			return true
		}
	}
	return false
}

func boolValue(object map[string]any, keys ...string) bool {
	for _, key := range keys {
		if value, ok := object[key].(bool); ok {
			return value
		}
	}
	return false
}

func numberValue(object map[string]any, keys ...string) float64 {
	value, _ := numericValue(object, keys...)
	return value
}

func numericValue(object map[string]any, keys ...string) (float64, bool) {
	for _, key := range keys {
		value, ok := object[key]
		if !ok || value == nil {
			continue
		}
		var parsed float64
		var err error
		switch typed := value.(type) {
		case float64:
			parsed = typed
		case string:
			parsed, err = strconv.ParseFloat(strings.TrimSpace(typed), 64)
		case json.Number:
			parsed, err = typed.Float64()
		default:
			continue
		}
		if err == nil && !math.IsNaN(parsed) && !math.IsInf(parsed, 0) {
			return parsed, true
		}
	}
	return 0, false
}

func nonNegativeInt64Pointer(object map[string]any, keys ...string) *int64 {
	if value, ok := numericValue(object, keys...); ok && value >= 0 {
		parsed := int64(value)
		return &parsed
	}
	return nil
}

func nonNegativeIntPointer(object map[string]any, keys ...string) *int {
	value := nonNegativeInt64Pointer(object, keys...)
	if value == nil {
		return nil
	}
	parsed := int(*value)
	return &parsed
}

func unixTimeValue(object map[string]any, keys ...string) int64 {
	for _, key := range keys {
		value, ok := object[key]
		if !ok || value == nil {
			continue
		}
		if text, ok := value.(string); ok {
			if parsed, err := time.Parse(time.RFC3339, text); err == nil {
				return parsed.Unix()
			}
			if parsed, err := strconv.ParseInt(text, 10, 64); err == nil {
				return parsed
			}
		}
		return int64(numberValue(object, key))
	}
	return 0
}