sub2api/backend/internal/model/account.go

package model

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"time"

	"gorm.io/gorm"
)

// JSONB 用于存储JSONB数据
type JSONB map[string]any

func (j JSONB) Value() (driver.Value, error) {
	if j == nil {
		return nil, nil
	}
	return json.Marshal(j)
}

func (j *JSONB) Scan(value any) error {
	if value == nil {
		*j = nil
		return nil
	}
	bytes, ok := value.([]byte)
	if !ok {
		return errors.New("type assertion to []byte failed")
	}
	return json.Unmarshal(bytes, j)
}

type Account struct {
	ID           int64          `gorm:"primaryKey" json:"id"`
	Name         string         `gorm:"size:100;not null" json:"name"`
	Platform     string         `gorm:"size:50;not null" json:"platform"`           // anthropic/openai/gemini
	Type         string         `gorm:"size:20;not null" json:"type"`               // oauth/apikey
	Credentials  JSONB          `gorm:"type:jsonb;default:'{}'" json:"credentials"` // 凭证(加密存储)
	Extra        JSONB          `gorm:"type:jsonb;default:'{}'" json:"extra"`       // 扩展信息
	ProxyID      *int64         `gorm:"index" json:"proxy_id"`
	Concurrency  int            `gorm:"default:3;not null" json:"concurrency"`
	Priority     int            `gorm:"default:50;not null" json:"priority"`           // 1-100，越小越高
	Status       string         `gorm:"size:20;default:active;not null" json:"status"` // active/disabled/error
	ErrorMessage string         `gorm:"type:text" json:"error_message"`
	LastUsedAt   *time.Time     `gorm:"index" json:"last_used_at"`
	CreatedAt    time.Time      `gorm:"not null" json:"created_at"`
	UpdatedAt    time.Time      `gorm:"not null" json:"updated_at"`
	DeletedAt    gorm.DeletedAt `gorm:"index" json:"-"`

	// 调度控制
	Schedulable bool `gorm:"default:true;not null" json:"schedulable"`

	// 限流状态 (429)
	RateLimitedAt    *time.Time `gorm:"index" json:"rate_limited_at"`
	RateLimitResetAt *time.Time `gorm:"index" json:"rate_limit_reset_at"`

	// 过载状态 (529)
	OverloadUntil *time.Time `gorm:"index" json:"overload_until"`

	// 5小时时间窗口
	SessionWindowStart  *time.Time `json:"session_window_start"`
	SessionWindowEnd    *time.Time `json:"session_window_end"`
	SessionWindowStatus string     `gorm:"size:20" json:"session_window_status"` // allowed/allowed_warning/rejected

	// 关联
	Proxy         *Proxy         `gorm:"foreignKey:ProxyID" json:"proxy,omitempty"`
	AccountGroups []AccountGroup `gorm:"foreignKey:AccountID" json:"account_groups,omitempty"`

	// 虚拟字段 (不存储到数据库)
	GroupIDs []int64  `gorm:"-" json:"group_ids,omitempty"`
	Groups   []*Group `gorm:"-" json:"groups,omitempty"`
}

func (Account) TableName() string {
	return "accounts"
}

// IsActive 检查是否激活
func (a *Account) IsActive() bool {
	return a.Status == "active"
}

// IsSchedulable 检查账号是否可调度
func (a *Account) IsSchedulable() bool {
	if !a.IsActive() || !a.Schedulable {
		return false
	}
	now := time.Now()
	if a.OverloadUntil != nil && now.Before(*a.OverloadUntil) {
		return false
	}
	if a.RateLimitResetAt != nil && now.Before(*a.RateLimitResetAt) {
		return false
	}
	return true
}

// IsRateLimited 检查是否处于限流状态
func (a *Account) IsRateLimited() bool {
	if a.RateLimitResetAt == nil {
		return false
	}
	return time.Now().Before(*a.RateLimitResetAt)
}

// IsOverloaded 检查是否处于过载状态
func (a *Account) IsOverloaded() bool {
	if a.OverloadUntil == nil {
		return false
	}
	return time.Now().Before(*a.OverloadUntil)
}

// IsOAuth 检查是否为OAuth类型账号（包括oauth和setup-token）
func (a *Account) IsOAuth() bool {
	return a.Type == AccountTypeOAuth || a.Type == AccountTypeSetupToken
}

// CanGetUsage 检查账号是否可以获取usage信息（只有oauth类型可以，setup-token没有profile权限）
func (a *Account) CanGetUsage() bool {
	return a.Type == AccountTypeOAuth
}

// GetCredential 获取凭证字段
func (a *Account) GetCredential(key string) string {
	if a.Credentials == nil {
		return ""
	}
	if v, ok := a.Credentials[key]; ok {
		if s, ok := v.(string); ok {
			return s
		}
	}
	return ""
}

// GetModelMapping 获取模型映射配置
// 返回格式: map[请求模型名]实际模型名
func (a *Account) GetModelMapping() map[string]string {
	if a.Credentials == nil {
		return nil
	}
	raw, ok := a.Credentials["model_mapping"]
	if !ok || raw == nil {
		return nil
	}
	// 处理map[string]interface{}类型
	if m, ok := raw.(map[string]any); ok {
		result := make(map[string]string)
		for k, v := range m {
			if s, ok := v.(string); ok {
				result[k] = s
			}
		}
		if len(result) > 0 {
			return result
		}
	}
	return nil
}

// IsModelSupported 检查请求的模型是否被该账号支持
// 如果没有设置模型映射，则支持所有模型
func (a *Account) IsModelSupported(requestedModel string) bool {
	mapping := a.GetModelMapping()
	if len(mapping) == 0 {
		return true // 没有映射配置，支持所有模型
	}
	_, exists := mapping[requestedModel]
	return exists
}

// GetMappedModel 获取映射后的实际模型名
// 如果没有映射，返回原始模型名
func (a *Account) GetMappedModel(requestedModel string) string {
	mapping := a.GetModelMapping()
	if len(mapping) == 0 {
		return requestedModel
	}
	if mappedModel, exists := mapping[requestedModel]; exists {
		return mappedModel
	}
	return requestedModel
}

// GetBaseURL 获取API基础URL（用于apikey类型账号）
func (a *Account) GetBaseURL() string {
	if a.Type != AccountTypeApiKey {
		return ""
	}
	baseURL := a.GetCredential("base_url")
	if baseURL == "" {
		return "https://api.anthropic.com" // 默认URL
	}
	return baseURL
}

// GetExtraString 从Extra字段获取字符串值
func (a *Account) GetExtraString(key string) string {
	if a.Extra == nil {
		return ""
	}
	if v, ok := a.Extra[key]; ok {
		if s, ok := v.(string); ok {
			return s
		}
	}
	return ""
}

// IsCustomErrorCodesEnabled 检查是否启用自定义错误码功能（仅适用于 apikey 类型）
func (a *Account) IsCustomErrorCodesEnabled() bool {
	if a.Type != AccountTypeApiKey || a.Credentials == nil {
		return false
	}
	if v, ok := a.Credentials["custom_error_codes_enabled"]; ok {
		if enabled, ok := v.(bool); ok {
			return enabled
		}
	}
	return false
}

// GetCustomErrorCodes 获取自定义错误码列表
func (a *Account) GetCustomErrorCodes() []int {
	if a.Credentials == nil {
		return nil
	}
	raw, ok := a.Credentials["custom_error_codes"]
	if !ok || raw == nil {
		return nil
	}
	// 处理 []interface{} 类型（JSON反序列化后的格式）
	if arr, ok := raw.([]any); ok {
		result := make([]int, 0, len(arr))
		for _, v := range arr {
			// JSON 数字默认解析为 float64
			if f, ok := v.(float64); ok {
				result = append(result, int(f))
			}
		}
		return result
	}
	return nil
}

// ShouldHandleErrorCode 检查指定错误码是否应该被处理（停止调度/标记限流等）
// 如果未启用自定义错误码或列表为空，返回 true（使用默认策略）
// 如果启用且列表非空，只有在列表中的错误码才返回 true
func (a *Account) ShouldHandleErrorCode(statusCode int) bool {
	if !a.IsCustomErrorCodesEnabled() {
		return true // 未启用，使用默认策略
	}
	codes := a.GetCustomErrorCodes()
	if len(codes) == 0 {
		return true // 启用但列表为空，fallback到默认策略
	}
	// 检查是否在自定义列表中
	for _, code := range codes {
		if code == statusCode {
			return true
		}
	}
	return false
}

// IsInterceptWarmupEnabled 检查是否启用预热请求拦截
// 启用后，标题生成、Warmup等预热请求将返回mock响应，不消耗上游token
func (a *Account) IsInterceptWarmupEnabled() bool {
	if a.Credentials == nil {
		return false
	}
	if v, ok := a.Credentials["intercept_warmup_requests"]; ok {
		if enabled, ok := v.(bool); ok {
			return enabled
		}
	}
	return false
}

// =============== OpenAI 相关方法 ===============

// IsOpenAI 检查是否为 OpenAI 平台账号
func (a *Account) IsOpenAI() bool {
	return a.Platform == PlatformOpenAI
}

// IsAnthropic 检查是否为 Anthropic 平台账号
func (a *Account) IsAnthropic() bool {
	return a.Platform == PlatformAnthropic
}

// IsOpenAIOAuth 检查是否为 OpenAI OAuth 类型账号
func (a *Account) IsOpenAIOAuth() bool {
	return a.IsOpenAI() && a.Type == AccountTypeOAuth
}

// IsOpenAIApiKey 检查是否为 OpenAI API Key 类型账号（Response 账号）
func (a *Account) IsOpenAIApiKey() bool {
	return a.IsOpenAI() && a.Type == AccountTypeApiKey
}

// GetOpenAIBaseURL 获取 OpenAI API 基础 URL
// 对于 API Key 类型账号，从 credentials 中获取 base_url
// 对于 OAuth 类型账号，返回默认的 OpenAI API URL
func (a *Account) GetOpenAIBaseURL() string {
	if !a.IsOpenAI() {
		return ""
	}
	if a.Type == AccountTypeApiKey {
		baseURL := a.GetCredential("base_url")
		if baseURL != "" {
			return baseURL
		}
	}
	return "https://api.openai.com" // OpenAI 默认 API URL
}

// GetOpenAIAccessToken 获取 OpenAI 访问令牌
func (a *Account) GetOpenAIAccessToken() string {
	if !a.IsOpenAI() {
		return ""
	}
	return a.GetCredential("access_token")
}

// GetOpenAIRefreshToken 获取 OpenAI 刷新令牌
func (a *Account) GetOpenAIRefreshToken() string {
	if !a.IsOpenAIOAuth() {
		return ""
	}
	return a.GetCredential("refresh_token")
}

// GetOpenAIIDToken 获取 OpenAI ID Token（JWT，包含用户信息）
func (a *Account) GetOpenAIIDToken() string {
	if !a.IsOpenAIOAuth() {
		return ""
	}
	return a.GetCredential("id_token")
}

// GetOpenAIApiKey 获取 OpenAI API Key（用于 Response 账号）
func (a *Account) GetOpenAIApiKey() string {
	if !a.IsOpenAIApiKey() {
		return ""
	}
	return a.GetCredential("api_key")
}

// GetOpenAIUserAgent 获取 OpenAI 自定义 User-Agent
// 返回空字符串表示透传原始 User-Agent
func (a *Account) GetOpenAIUserAgent() string {
	if !a.IsOpenAI() {
		return ""
	}
	return a.GetCredential("user_agent")
}

// GetChatGPTAccountID 获取 ChatGPT 账号 ID（从 ID Token 解析）
func (a *Account) GetChatGPTAccountID() string {
	if !a.IsOpenAIOAuth() {
		return ""
	}
	return a.GetCredential("chatgpt_account_id")
}

// GetChatGPTUserID 获取 ChatGPT 用户 ID（从 ID Token 解析）
func (a *Account) GetChatGPTUserID() string {
	if !a.IsOpenAIOAuth() {
		return ""
	}
	return a.GetCredential("chatgpt_user_id")
}

// GetOpenAIOrganizationID 获取 OpenAI 组织 ID
func (a *Account) GetOpenAIOrganizationID() string {
	if !a.IsOpenAIOAuth() {
		return ""
	}
	return a.GetCredential("organization_id")
}

// GetOpenAITokenExpiresAt 获取 OpenAI Token 过期时间
func (a *Account) GetOpenAITokenExpiresAt() *time.Time {
	if !a.IsOpenAIOAuth() {
		return nil
	}
	expiresAtStr := a.GetCredential("expires_at")
	if expiresAtStr == "" {
		return nil
	}
	// 尝试解析时间
	t, err := time.Parse(time.RFC3339, expiresAtStr)
	if err != nil {
		// 尝试解析为 Unix 时间戳
		if v, ok := a.Credentials["expires_at"].(float64); ok {
			t = time.Unix(int64(v), 0)
			return &t
		}
		return nil
	}
	return &t
}

// IsOpenAITokenExpired 检查 OpenAI Token 是否过期
func (a *Account) IsOpenAITokenExpired() bool {
	expiresAt := a.GetOpenAITokenExpiresAt()
	if expiresAt == nil {
		return false // 没有过期时间信息，假设未过期
	}
	// 提前 60 秒认为过期，便于刷新
	return time.Now().Add(60 * time.Second).After(*expiresAt)
}