sub2api/backend/internal/service/gateway_service.go

package service

import (
	"bufio"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"sub2api/internal/config"
	"sub2api/internal/model"
	"sub2api/internal/repository"

	"github.com/gin-gonic/gin"
	"github.com/redis/go-redis/v9"
)

const (
	claudeAPIURL        = "https://api.anthropic.com/v1/messages?beta=true"
	stickySessionPrefix = "sticky_session:"
	stickySessionTTL    = time.Hour // 粘性会话TTL
	tokenRefreshBuffer  = 5 * 60    // 提前5分钟刷新token
)

// allowedHeaders 白名单headers（参考CRS项目）
var allowedHeaders = map[string]bool{
	"accept":                                  true,
	"x-stainless-retry-count":                 true,
	"x-stainless-timeout":                     true,
	"x-stainless-lang":                        true,
	"x-stainless-package-version":             true,
	"x-stainless-os":                          true,
	"x-stainless-arch":                        true,
	"x-stainless-runtime":                     true,
	"x-stainless-runtime-version":             true,
	"x-stainless-helper-method":               true,
	"anthropic-dangerous-direct-browser-access": true,
	"anthropic-version":                       true,
	"x-app":                                   true,
	"anthropic-beta":                          true,
	"accept-language":                         true,
	"sec-fetch-mode":                          true,
	"accept-encoding":                         true,
	"user-agent":                              true,
	"content-type":                            true,
}

// ClaudeUsage 表示Claude API返回的usage信息
type ClaudeUsage struct {
	InputTokens              int `json:"input_tokens"`
	OutputTokens             int `json:"output_tokens"`
	CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
	CacheReadInputTokens     int `json:"cache_read_input_tokens"`
}

// ForwardResult 转发结果
type ForwardResult struct {
	RequestID    string
	Usage        ClaudeUsage
	Model        string
	Stream       bool
	Duration     time.Duration
	FirstTokenMs *int // 首字时间（流式请求）
}

// GatewayService handles API gateway operations
type GatewayService struct {
	repos               *repository.Repositories
	rdb                 *redis.Client
	cfg                 *config.Config
	oauthService        *OAuthService
	billingService      *BillingService
	rateLimitService    *RateLimitService
	billingCacheService *BillingCacheService
	identityService     *IdentityService
	httpClient          *http.Client
}

// NewGatewayService creates a new GatewayService
func NewGatewayService(repos *repository.Repositories, rdb *redis.Client, cfg *config.Config, oauthService *OAuthService, billingService *BillingService, rateLimitService *RateLimitService, billingCacheService *BillingCacheService, identityService *IdentityService) *GatewayService {
	// 计算响应头超时时间
	responseHeaderTimeout := time.Duration(cfg.Gateway.ResponseHeaderTimeout) * time.Second
	if responseHeaderTimeout == 0 {
		responseHeaderTimeout = 300 * time.Second // 默认5分钟，LLM高负载时可能排队较久
	}

	transport := &http.Transport{
		MaxIdleConns:          100,
		MaxIdleConnsPerHost:   10,
		IdleConnTimeout:       90 * time.Second,
		ResponseHeaderTimeout: responseHeaderTimeout, // 等待上游响应头的超时
		// 注意：不设置整体 Timeout，让流式响应可以无限时间传输
	}
	return &GatewayService{
		repos:               repos,
		rdb:                 rdb,
		cfg:                 cfg,
		oauthService:        oauthService,
		billingService:      billingService,
		rateLimitService:    rateLimitService,
		billingCacheService: billingCacheService,
		identityService:     identityService,
		httpClient: &http.Client{
			Transport: transport,
			// 不设置 Timeout：流式请求可能持续十几分钟
			// 超时控制由 Transport.ResponseHeaderTimeout 负责（只控制等待响应头）
		},
	}
}

// GenerateSessionHash 从请求体计算粘性会话hash
func (s *GatewayService) GenerateSessionHash(body []byte) string {
	var req map[string]interface{}
	if err := json.Unmarshal(body, &req); err != nil {
		return ""
	}

	// 1. 最高优先级：从metadata.user_id提取session_xxx
	if metadata, ok := req["metadata"].(map[string]interface{}); ok {
		if userID, ok := metadata["user_id"].(string); ok {
			re := regexp.MustCompile(`session_([a-f0-9-]{36})`)
			if match := re.FindStringSubmatch(userID); len(match) > 1 {
				return match[1]
			}
		}
	}

	// 2. 提取带cache_control: {type: "ephemeral"}的内容
	cacheableContent := s.extractCacheableContent(req)
	if cacheableContent != "" {
		return s.hashContent(cacheableContent)
	}

	// 3. Fallback: 使用system内容
	if system := req["system"]; system != nil {
		systemText := s.extractTextFromSystem(system)
		if systemText != "" {
			return s.hashContent(systemText)
		}
	}

	// 4. 最后fallback: 使用第一条消息
	if messages, ok := req["messages"].([]interface{}); ok && len(messages) > 0 {
		if firstMsg, ok := messages[0].(map[string]interface{}); ok {
			msgText := s.extractTextFromContent(firstMsg["content"])
			if msgText != "" {
				return s.hashContent(msgText)
			}
		}
	}

	return ""
}

func (s *GatewayService) extractCacheableContent(req map[string]interface{}) string {
	var content string

	// 检查system中的cacheable内容
	if system, ok := req["system"].([]interface{}); ok {
		for _, part := range system {
			if partMap, ok := part.(map[string]interface{}); ok {
				if cc, ok := partMap["cache_control"].(map[string]interface{}); ok {
					if cc["type"] == "ephemeral" {
						if text, ok := partMap["text"].(string); ok {
							content += text
						}
					}
				}
			}
		}
	}

	// 检查messages中的cacheable内容
	if messages, ok := req["messages"].([]interface{}); ok {
		for _, msg := range messages {
			if msgMap, ok := msg.(map[string]interface{}); ok {
				if msgContent, ok := msgMap["content"].([]interface{}); ok {
					for _, part := range msgContent {
						if partMap, ok := part.(map[string]interface{}); ok {
							if cc, ok := partMap["cache_control"].(map[string]interface{}); ok {
								if cc["type"] == "ephemeral" {
									// 找到cacheable内容，提取第一条消息的文本
									return s.extractTextFromContent(msgMap["content"])
								}
							}
						}
					}
				}
			}
		}
	}

	return content
}

func (s *GatewayService) extractTextFromSystem(system interface{}) string {
	switch v := system.(type) {
	case string:
		return v
	case []interface{}:
		var texts []string
		for _, part := range v {
			if partMap, ok := part.(map[string]interface{}); ok {
				if text, ok := partMap["text"].(string); ok {
					texts = append(texts, text)
				}
			}
		}
		return strings.Join(texts, "")
	}
	return ""
}

func (s *GatewayService) extractTextFromContent(content interface{}) string {
	switch v := content.(type) {
	case string:
		return v
	case []interface{}:
		var texts []string
		for _, part := range v {
			if partMap, ok := part.(map[string]interface{}); ok {
				if partMap["type"] == "text" {
					if text, ok := partMap["text"].(string); ok {
						texts = append(texts, text)
					}
				}
			}
		}
		return strings.Join(texts, "")
	}
	return ""
}

func (s *GatewayService) hashContent(content string) string {
	hash := sha256.Sum256([]byte(content))
	return hex.EncodeToString(hash[:16]) // 32字符
}

// replaceModelInBody 替换请求体中的model字段
func (s *GatewayService) replaceModelInBody(body []byte, newModel string) []byte {
	var req map[string]interface{}
	if err := json.Unmarshal(body, &req); err != nil {
		return body
	}
	req["model"] = newModel
	newBody, err := json.Marshal(req)
	if err != nil {
		return body
	}
	return newBody
}

// SelectAccount 选择账号（粘性会话+优先级）
func (s *GatewayService) SelectAccount(ctx context.Context, groupID *int64, sessionHash string) (*model.Account, error) {
	return s.SelectAccountForModel(ctx, groupID, sessionHash, "")
}

// SelectAccountForModel 选择支持指定模型的账号（粘性会话+优先级+模型映射）
func (s *GatewayService) SelectAccountForModel(ctx context.Context, groupID *int64, sessionHash string, requestedModel string) (*model.Account, error) {
	// 1. 查询粘性会话
	if sessionHash != "" {
		accountID, err := s.rdb.Get(ctx, stickySessionPrefix+sessionHash).Int64()
		if err == nil && accountID > 0 {
			account, err := s.repos.Account.GetByID(ctx, accountID)
			// 使用IsSchedulable代替IsActive，确保限流/过载账号不会被选中
			// 同时检查模型支持
			if err == nil && account.IsSchedulable() && (requestedModel == "" || account.IsModelSupported(requestedModel)) {
				// 续期粘性会话
				s.rdb.Expire(ctx, stickySessionPrefix+sessionHash, stickySessionTTL)
				return account, nil
			}
		}
	}

	// 2. 获取可调度账号列表（排除限流和过载的账号）
	var accounts []model.Account
	var err error
	if groupID != nil {
		accounts, err = s.repos.Account.ListSchedulableByGroupID(ctx, *groupID)
	} else {
		accounts, err = s.repos.Account.ListSchedulable(ctx)
	}
	if err != nil {
		return nil, fmt.Errorf("query accounts failed: %w", err)
	}

	// 3. 按优先级+最久未用选择（考虑模型支持）
	var selected *model.Account
	for i := range accounts {
		acc := &accounts[i]
		// 检查模型支持
		if requestedModel != "" && !acc.IsModelSupported(requestedModel) {
			continue
		}
		if selected == nil {
			selected = acc
			continue
		}
		// 优先选择priority值更小的（priority值越小优先级越高）
		if acc.Priority < selected.Priority {
			selected = acc
		} else if acc.Priority == selected.Priority {
			// 优先级相同时，选最久未用的
			if acc.LastUsedAt == nil || (selected.LastUsedAt != nil && acc.LastUsedAt.Before(*selected.LastUsedAt)) {
				selected = acc
			}
		}
	}

	if selected == nil {
		if requestedModel != "" {
			return nil, fmt.Errorf("no available accounts supporting model: %s", requestedModel)
		}
		return nil, errors.New("no available accounts")
	}

	// 4. 建立粘性绑定
	if sessionHash != "" {
		s.rdb.Set(ctx, stickySessionPrefix+sessionHash, selected.ID, stickySessionTTL)
	}

	return selected, nil
}

// GetAccessToken 获取账号凭证
func (s *GatewayService) GetAccessToken(ctx context.Context, account *model.Account) (string, string, error) {
	switch account.Type {
	case model.AccountTypeOAuth, model.AccountTypeSetupToken:
		// Both oauth and setup-token use OAuth token flow
		return s.getOAuthToken(ctx, account)
	case model.AccountTypeApiKey:
		apiKey := account.GetCredential("api_key")
		if apiKey == "" {
			return "", "", errors.New("api_key not found in credentials")
		}
		return apiKey, "apikey", nil
	default:
		return "", "", fmt.Errorf("unsupported account type: %s", account.Type)
	}
}

func (s *GatewayService) getOAuthToken(ctx context.Context, account *model.Account) (string, string, error) {
	accessToken := account.GetCredential("access_token")
	expiresAtStr := account.GetCredential("expires_at")

	// 检查是否需要刷新
	needRefresh := false
	if expiresAtStr != "" {
		expiresAt, err := strconv.ParseInt(expiresAtStr, 10, 64)
		if err == nil && time.Now().Unix()+tokenRefreshBuffer > expiresAt {
			needRefresh = true
		}
	}

	if needRefresh || accessToken == "" {
		tokenInfo, err := s.oauthService.RefreshAccountToken(ctx, account)
		if err != nil {
			return "", "", fmt.Errorf("refresh token failed: %w", err)
		}

		// 更新账号凭证
		account.Credentials["access_token"] = tokenInfo.AccessToken
		account.Credentials["expires_at"] = strconv.FormatInt(tokenInfo.ExpiresAt, 10)
		if tokenInfo.RefreshToken != "" {
			account.Credentials["refresh_token"] = tokenInfo.RefreshToken
		}

		if err := s.repos.Account.Update(ctx, account); err != nil {
			log.Printf("Failed to update account credentials: %v", err)
		}

		return tokenInfo.AccessToken, "oauth", nil
	}

	return accessToken, "oauth", nil
}

// Forward 转发请求到Claude API
func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *model.Account, body []byte) (*ForwardResult, error) {
	startTime := time.Now()

	// 解析请求获取model和stream
	var req struct {
		Model  string `json:"model"`
		Stream bool   `json:"stream"`
	}
	if err := json.Unmarshal(body, &req); err != nil {
		return nil, fmt.Errorf("parse request: %w", err)
	}

	// 应用模型映射（仅对apikey类型账号）
	originalModel := req.Model
	if account.Type == model.AccountTypeApiKey {
		mappedModel := account.GetMappedModel(req.Model)
		if mappedModel != req.Model {
			// 替换请求体中的模型名
			body = s.replaceModelInBody(body, mappedModel)
			req.Model = mappedModel
			log.Printf("Model mapping applied: %s -> %s (account: %s)", originalModel, mappedModel, account.Name)
		}
	}

	// 获取凭证
	token, tokenType, err := s.GetAccessToken(ctx, account)
	if err != nil {
		return nil, err
	}

	// 构建上游请求
	upstreamReq, err := s.buildUpstreamRequest(ctx, c, account, body, token, tokenType)
	if err != nil {
		return nil, err
	}

	// 发送请求
	resp, err := s.httpClient.Do(upstreamReq)
	if err != nil {
		return nil, fmt.Errorf("upstream request failed: %w", err)
	}
	defer resp.Body.Close()

	// 处理401错误：刷新token重试
	if resp.StatusCode == http.StatusUnauthorized && tokenType == "oauth" {
		resp.Body.Close()
		token, tokenType, err = s.forceRefreshToken(ctx, account)
		if err != nil {
			return nil, fmt.Errorf("token refresh failed: %w", err)
		}
		upstreamReq, err = s.buildUpstreamRequest(ctx, c, account, body, token, tokenType)
		if err != nil {
			return nil, err
		}
		resp, err = s.httpClient.Do(upstreamReq)
		if err != nil {
			return nil, fmt.Errorf("retry request failed: %w", err)
		}
		defer resp.Body.Close()
	}

	// 处理错误响应
	if resp.StatusCode >= 400 {
		return s.handleErrorResponse(ctx, resp, c, account)
	}

	// 处理正常响应
	var usage *ClaudeUsage
	var firstTokenMs *int
	if req.Stream {
		streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, req.Model)
		if err != nil {
			return nil, err
		}
		usage = streamResult.usage
		firstTokenMs = streamResult.firstTokenMs
	} else {
		usage, err = s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, req.Model)
		if err != nil {
			return nil, err
		}
	}

	return &ForwardResult{
		RequestID:    resp.Header.Get("x-request-id"),
		Usage:        *usage,
		Model:        originalModel, // 使用原始模型用于计费和日志
		Stream:       req.Stream,
		Duration:     time.Since(startTime),
		FirstTokenMs: firstTokenMs,
	}, nil
}

func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *model.Account, body []byte, token, tokenType string) (*http.Request, error) {
	// 确定目标URL
	targetURL := claudeAPIURL
	if account.Type == model.AccountTypeApiKey {
		baseURL := account.GetBaseURL()
		targetURL = baseURL + "/v1/messages"
	}

	// OAuth账号：应用统一指纹
	var fingerprint *Fingerprint
	if account.IsOAuth() && s.identityService != nil {
		// 1. 获取或创建指纹（包含随机生成的ClientID）
		fp, err := s.identityService.GetOrCreateFingerprint(ctx, account.ID, c.Request.Header)
		if err != nil {
			log.Printf("Warning: failed to get fingerprint for account %d: %v", account.ID, err)
			// 失败时降级为透传原始headers
		} else {
			fingerprint = fp

			// 2. 重写metadata.user_id（需要指纹中的ClientID和账号的account_uuid）
			accountUUID := account.GetExtraString("account_uuid")
			if accountUUID != "" && fp.ClientID != "" {
				if newBody, err := s.identityService.RewriteUserID(body, account.ID, accountUUID, fp.ClientID); err == nil && len(newBody) > 0 {
					body = newBody
				}
			}
		}
	}

	req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}

	// 设置认证头
	if tokenType == "oauth" {
		req.Header.Set("Authorization", "Bearer "+token)
	} else {
		req.Header.Set("x-api-key", token)
	}

	// 白名单透传headers
	for key, values := range c.Request.Header {
		lowerKey := strings.ToLower(key)
		if allowedHeaders[lowerKey] {
			for _, v := range values {
				req.Header.Add(key, v)
			}
		}
	}

	// OAuth账号：应用缓存的指纹到请求头（覆盖白名单透传的头）
	if fingerprint != nil {
		s.identityService.ApplyFingerprint(req, fingerprint)
	}

	// 确保必要的headers存在
	if req.Header.Get("Content-Type") == "" {
		req.Header.Set("Content-Type", "application/json")
	}
	if req.Header.Get("anthropic-version") == "" {
		req.Header.Set("anthropic-version", "2023-06-01")
	}

	// 处理anthropic-beta header（OAuth账号需要特殊处理）
	if tokenType == "oauth" {
		req.Header.Set("anthropic-beta", s.getBetaHeader(body, c.GetHeader("anthropic-beta")))
	}

	// 配置代理
	if account.ProxyID != nil && account.Proxy != nil {
		proxyURL := account.Proxy.URL()
		if proxyURL != "" {
			if parsedURL, err := url.Parse(proxyURL); err == nil {
				// 计算响应头超时时间（与默认 Transport 保持一致）
				responseHeaderTimeout := time.Duration(s.cfg.Gateway.ResponseHeaderTimeout) * time.Second
				if responseHeaderTimeout == 0 {
					responseHeaderTimeout = 300 * time.Second
				}
				transport := &http.Transport{
					Proxy:                 http.ProxyURL(parsedURL),
					MaxIdleConns:          100,
					MaxIdleConnsPerHost:   10,
					IdleConnTimeout:       90 * time.Second,
					ResponseHeaderTimeout: responseHeaderTimeout,
				}
				s.httpClient.Transport = transport
			}
		}
	}

	return req, nil
}

// getBetaHeader 处理anthropic-beta header
// 对于OAuth账号，需要确保包含oauth-2025-04-20
func (s *GatewayService) getBetaHeader(body []byte, clientBetaHeader string) string {
	const oauthBeta = "oauth-2025-04-20"
	const claudeCodeBeta = "claude-code-20250219"

	// 如果客户端传了anthropic-beta
	if clientBetaHeader != "" {
		// 已包含oauth beta则直接返回
		if strings.Contains(clientBetaHeader, oauthBeta) {
			return clientBetaHeader
		}

		// 需要添加oauth beta
		parts := strings.Split(clientBetaHeader, ",")
		for i, p := range parts {
			parts[i] = strings.TrimSpace(p)
		}

		// 在claude-code-20250219后面插入oauth beta
		claudeCodeIdx := -1
		for i, p := range parts {
			if p == claudeCodeBeta {
				claudeCodeIdx = i
				break
			}
		}

		if claudeCodeIdx >= 0 {
			// 在claude-code后面插入
			newParts := make([]string, 0, len(parts)+1)
			newParts = append(newParts, parts[:claudeCodeIdx+1]...)
			newParts = append(newParts, oauthBeta)
			newParts = append(newParts, parts[claudeCodeIdx+1:]...)
			return strings.Join(newParts, ",")
		}

		// 没有claude-code，放在第一位
		return oauthBeta + "," + clientBetaHeader
	}

	// 客户端没传，根据模型生成
	var modelID string
	var reqMap map[string]interface{}
	if json.Unmarshal(body, &reqMap) == nil {
		if m, ok := reqMap["model"].(string); ok {
			modelID = m
		}
	}

	// haiku模型不需要claude-code beta
	if strings.Contains(strings.ToLower(modelID), "haiku") {
		return "oauth-2025-04-20,interleaved-thinking-2025-05-14"
	}

	return "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,fine-grained-tool-streaming-2025-05-14"
}

func (s *GatewayService) forceRefreshToken(ctx context.Context, account *model.Account) (string, string, error) {
	tokenInfo, err := s.oauthService.RefreshAccountToken(ctx, account)
	if err != nil {
		return "", "", err
	}

	account.Credentials["access_token"] = tokenInfo.AccessToken
	account.Credentials["expires_at"] = strconv.FormatInt(tokenInfo.ExpiresAt, 10)
	if tokenInfo.RefreshToken != "" {
		account.Credentials["refresh_token"] = tokenInfo.RefreshToken
	}

	if err := s.repos.Account.Update(ctx, account); err != nil {
		log.Printf("Failed to update account credentials: %v", err)
	}

	return tokenInfo.AccessToken, "oauth", nil
}

func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *model.Account) (*ForwardResult, error) {
	body, _ := io.ReadAll(resp.Body)

	// apikey 类型账号：检查自定义错误码配置
	// 如果启用且错误码不在列表中，返回通用 500 错误（不做任何账号状态处理）
	if !account.ShouldHandleErrorCode(resp.StatusCode) {
		c.JSON(http.StatusInternalServerError, gin.H{
			"type": "error",
			"error": gin.H{
				"type":    "upstream_error",
				"message": "Upstream gateway error",
			},
		})
		return nil, fmt.Errorf("upstream error: %d (not in custom error codes)", resp.StatusCode)
	}

	// 处理上游错误，标记账号状态
	s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body)

	// 根据状态码返回适当的自定义错误响应（不透传上游详细信息）
	var errType, errMsg string
	var statusCode int

	switch resp.StatusCode {
	case 401:
		statusCode = http.StatusBadGateway
		errType = "upstream_error"
		errMsg = "Upstream authentication failed, please contact administrator"
	case 403:
		statusCode = http.StatusBadGateway
		errType = "upstream_error"
		errMsg = "Upstream access forbidden, please contact administrator"
	case 429:
		statusCode = http.StatusTooManyRequests
		errType = "rate_limit_error"
		errMsg = "Upstream rate limit exceeded, please retry later"
	case 529:
		statusCode = http.StatusServiceUnavailable
		errType = "overloaded_error"
		errMsg = "Upstream service overloaded, please retry later"
	case 500, 502, 503, 504:
		statusCode = http.StatusBadGateway
		errType = "upstream_error"
		errMsg = "Upstream service temporarily unavailable"
	default:
		statusCode = http.StatusBadGateway
		errType = "upstream_error"
		errMsg = "Upstream request failed"
	}

	// 返回自定义错误响应
	c.JSON(statusCode, gin.H{
		"type": "error",
		"error": gin.H{
			"type":    errType,
			"message": errMsg,
		},
	})

	return nil, fmt.Errorf("upstream error: %d", resp.StatusCode)
}

// streamingResult 流式响应结果
type streamingResult struct {
	usage        *ClaudeUsage
	firstTokenMs *int
}

func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *model.Account, startTime time.Time, originalModel, mappedModel string) (*streamingResult, error) {
	// 更新5h窗口状态
	s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header)

	// 设置SSE响应头
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("X-Accel-Buffering", "no")

	// 透传其他响应头
	if v := resp.Header.Get("x-request-id"); v != "" {
		c.Header("x-request-id", v)
	}

	w := c.Writer
	flusher, ok := w.(http.Flusher)
	if !ok {
		return nil, errors.New("streaming not supported")
	}

	usage := &ClaudeUsage{}
	var firstTokenMs *int
	scanner := bufio.NewScanner(resp.Body)
	// 设置更大的buffer以处理长行
	scanner.Buffer(make([]byte, 64*1024), 1024*1024)

	needModelReplace := originalModel != mappedModel

	for scanner.Scan() {
		line := scanner.Text()

		// 如果有模型映射，替换响应中的model字段
		if needModelReplace && strings.HasPrefix(line, "data: ") {
			line = s.replaceModelInSSELine(line, mappedModel, originalModel)
		}

		// 转发行
		fmt.Fprintf(w, "%s\n", line)
		flusher.Flush()

		// 解析usage数据
		if strings.HasPrefix(line, "data: ") {
			data := line[6:]
			// 记录首字时间：第一个有效的 content_block_delta 或 message_start
			if firstTokenMs == nil && data != "" && data != "[DONE]" {
				ms := int(time.Since(startTime).Milliseconds())
				firstTokenMs = &ms
			}
			s.parseSSEUsage(data, usage)
		}
	}

	if err := scanner.Err(); err != nil {
		return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream read error: %w", err)
	}

	return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, nil
}

// replaceModelInSSELine 替换SSE数据行中的model字段
func (s *GatewayService) replaceModelInSSELine(line, fromModel, toModel string) string {
	data := line[6:] // 去掉 "data: " 前缀
	if data == "" || data == "[DONE]" {
		return line
	}

	var event map[string]interface{}
	if err := json.Unmarshal([]byte(data), &event); err != nil {
		return line
	}

	// 只替换 message_start 事件中的 message.model
	if event["type"] != "message_start" {
		return line
	}

	msg, ok := event["message"].(map[string]interface{})
	if !ok {
		return line
	}

	model, ok := msg["model"].(string)
	if !ok || model != fromModel {
		return line
	}

	msg["model"] = toModel
	newData, err := json.Marshal(event)
	if err != nil {
		return line
	}

	return "data: " + string(newData)
}

func (s *GatewayService) parseSSEUsage(data string, usage *ClaudeUsage) {
	// 解析message_start获取input tokens
	var msgStart struct {
		Type    string `json:"type"`
		Message struct {
			Usage ClaudeUsage `json:"usage"`
		} `json:"message"`
	}
	if json.Unmarshal([]byte(data), &msgStart) == nil && msgStart.Type == "message_start" {
		usage.InputTokens = msgStart.Message.Usage.InputTokens
		usage.CacheCreationInputTokens = msgStart.Message.Usage.CacheCreationInputTokens
		usage.CacheReadInputTokens = msgStart.Message.Usage.CacheReadInputTokens
	}

	// 解析message_delta获取output tokens
	var msgDelta struct {
		Type  string `json:"type"`
		Usage struct {
			OutputTokens int `json:"output_tokens"`
		} `json:"usage"`
	}
	if json.Unmarshal([]byte(data), &msgDelta) == nil && msgDelta.Type == "message_delta" {
		usage.OutputTokens = msgDelta.Usage.OutputTokens
	}
}

func (s *GatewayService) handleNonStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *model.Account, originalModel, mappedModel string) (*ClaudeUsage, error) {
	// 更新5h窗口状态
	s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header)

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// 解析usage
	var response struct {
		Usage ClaudeUsage `json:"usage"`
	}
	if err := json.Unmarshal(body, &response); err != nil {
		return nil, fmt.Errorf("parse response: %w", err)
	}

	// 如果有模型映射，替换响应中的model字段
	if originalModel != mappedModel {
		body = s.replaceModelInResponseBody(body, mappedModel, originalModel)
	}

	// 透传响应头
	for key, values := range resp.Header {
		for _, value := range values {
			c.Header(key, value)
		}
	}

	// 写入响应
	c.Data(resp.StatusCode, "application/json", body)

	return &response.Usage, nil
}

// replaceModelInResponseBody 替换响应体中的model字段
func (s *GatewayService) replaceModelInResponseBody(body []byte, fromModel, toModel string) []byte {
	var resp map[string]interface{}
	if err := json.Unmarshal(body, &resp); err != nil {
		return body
	}

	model, ok := resp["model"].(string)
	if !ok || model != fromModel {
		return body
	}

	resp["model"] = toModel
	newBody, err := json.Marshal(resp)
	if err != nil {
		return body
	}

	return newBody
}

// RecordUsageInput 记录使用量的输入参数
type RecordUsageInput struct {
	Result       *ForwardResult
	ApiKey       *model.ApiKey
	User         *model.User
	Account      *model.Account
	Subscription *model.UserSubscription // 可选：订阅信息
}

// RecordUsage 记录使用量并扣费（或更新订阅用量）
func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInput) error {
	result := input.Result
	apiKey := input.ApiKey
	user := input.User
	account := input.Account
	subscription := input.Subscription

	// 计算费用
	tokens := UsageTokens{
		InputTokens:         result.Usage.InputTokens,
		OutputTokens:        result.Usage.OutputTokens,
		CacheCreationTokens: result.Usage.CacheCreationInputTokens,
		CacheReadTokens:     result.Usage.CacheReadInputTokens,
	}

	// 获取费率倍数
	multiplier := s.cfg.Default.RateMultiplier
	if apiKey.GroupID != nil && apiKey.Group != nil {
		multiplier = apiKey.Group.RateMultiplier
	}

	cost, err := s.billingService.CalculateCost(result.Model, tokens, multiplier)
	if err != nil {
		log.Printf("Calculate cost failed: %v", err)
		// 使用默认费用继续
		cost = &CostBreakdown{ActualCost: 0}
	}

	// 判断计费方式：订阅模式 vs 余额模式
	isSubscriptionBilling := subscription != nil && apiKey.Group != nil && apiKey.Group.IsSubscriptionType()
	billingType := model.BillingTypeBalance
	if isSubscriptionBilling {
		billingType = model.BillingTypeSubscription
	}

	// 创建使用日志
	durationMs := int(result.Duration.Milliseconds())
	usageLog := &model.UsageLog{
		UserID:              user.ID,
		ApiKeyID:            apiKey.ID,
		AccountID:           account.ID,
		RequestID:           result.RequestID,
		Model:               result.Model,
		InputTokens:         result.Usage.InputTokens,
		OutputTokens:        result.Usage.OutputTokens,
		CacheCreationTokens: result.Usage.CacheCreationInputTokens,
		CacheReadTokens:     result.Usage.CacheReadInputTokens,
		InputCost:           cost.InputCost,
		OutputCost:          cost.OutputCost,
		CacheCreationCost:   cost.CacheCreationCost,
		CacheReadCost:       cost.CacheReadCost,
		TotalCost:           cost.TotalCost,
		ActualCost:          cost.ActualCost,
		RateMultiplier:      multiplier,
		BillingType:         billingType,
		Stream:              result.Stream,
		DurationMs:          &durationMs,
		FirstTokenMs:        result.FirstTokenMs,
		CreatedAt:           time.Now(),
	}

	// 添加分组和订阅关联
	if apiKey.GroupID != nil {
		usageLog.GroupID = apiKey.GroupID
	}
	if subscription != nil {
		usageLog.SubscriptionID = &subscription.ID
	}

	if err := s.repos.UsageLog.Create(ctx, usageLog); err != nil {
		log.Printf("Create usage log failed: %v", err)
	}

	// 根据计费类型执行扣费
	if isSubscriptionBilling {
		// 订阅模式：更新订阅用量（使用 TotalCost 原始费用，不考虑倍率）
		if cost.TotalCost > 0 {
			if err := s.repos.UserSubscription.IncrementUsage(ctx, subscription.ID, cost.TotalCost); err != nil {
				log.Printf("Increment subscription usage failed: %v", err)
			}
			// 异步更新订阅缓存
			go func() {
				cacheCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
				defer cancel()
				if err := s.billingCacheService.UpdateSubscriptionUsage(cacheCtx, user.ID, *apiKey.GroupID, cost.TotalCost); err != nil {
					log.Printf("Update subscription cache failed: %v", err)
				}
			}()
		}
	} else {
		// 余额模式：扣除用户余额（使用 ActualCost 考虑倍率后的费用）
		if cost.ActualCost > 0 {
			if err := s.repos.User.DeductBalance(ctx, user.ID, cost.ActualCost); err != nil {
				log.Printf("Deduct balance failed: %v", err)
			}
			// 异步更新余额缓存
			go func() {
				cacheCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
				defer cancel()
				if err := s.billingCacheService.DeductBalanceCache(cacheCtx, user.ID, cost.ActualCost); err != nil {
					log.Printf("Update balance cache failed: %v", err)
				}
			}()
		}
	}

	// 更新账号最后使用时间
	if err := s.repos.Account.UpdateLastUsed(ctx, account.ID); err != nil {
		log.Printf("Update last used failed: %v", err)
	}

	return nil
}