sub2api/backend/internal/service/tiered_routing.go

package service

import (
	"context"
	"encoding/json"
	"strings"
	"sync"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
)

// SettingKeyTieredRoutingSettings 分级路由规则（JSON，见 TieredRoutingSettings）。
const SettingKeyTieredRoutingSettings = "tiered_routing_settings"

// 账号池与用户等级：premium 用户优先 premium 池账号，无可用时回退全池。
const (
	TierStandard = "standard"
	TierPremium  = "premium"
)

// AccountPoolExtraKey 账号 extra 中的账号池标记键。
const AccountPoolExtraKey = "pool"

// TieredRoutingSettings 分级路由规则。
type TieredRoutingSettings struct {
	Enabled              bool    `json:"enabled"`
	PremiumUserIDs       []int64 `json:"premium_user_ids"`
	MinBalanceForPremium float64 `json:"min_balance_for_premium"`
}

// DefaultTieredRoutingSettings 默认关闭，需管理员显式启用并配置规则。
func DefaultTieredRoutingSettings() TieredRoutingSettings {
	return TieredRoutingSettings{Enabled: false}
}

// TieredRoutingService 分级路由服务：规则快照缓存 + 用户等级判定。
type TieredRoutingService struct {
	settingRepo SettingRepository

	mu       sync.RWMutex
	snapshot *TieredRoutingSettings
	loadedAt time.Time
}

// tieredRoutingSnapshotTTL 规则快照 TTL。
const tieredRoutingSnapshotTTL = 60 * time.Second

// NewTieredRoutingService 创建分级路由服务。
func NewTieredRoutingService(settingRepo SettingRepository) *TieredRoutingService {
	return &TieredRoutingService{settingRepo: settingRepo}
}

func (s *TieredRoutingService) currentSettings() TieredRoutingSettings {
	def := DefaultTieredRoutingSettings()
	if s == nil || s.settingRepo == nil {
		return def
	}
	s.mu.RLock()
	if s.snapshot != nil && time.Since(s.loadedAt) < tieredRoutingSnapshotTTL {
		out := *s.snapshot
		s.mu.RUnlock()
		return out
	}
	s.mu.RUnlock()

	s.mu.Lock()
	defer s.mu.Unlock()
	if s.snapshot != nil && time.Since(s.loadedAt) < tieredRoutingSnapshotTTL {
		return *s.snapshot
	}
	cfg := def
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if raw, err := s.settingRepo.GetValue(ctx, SettingKeyTieredRoutingSettings); err == nil && strings.TrimSpace(raw) != "" {
		var parsed TieredRoutingSettings
		if json.Unmarshal([]byte(raw), &parsed) == nil {
			cfg = normalizeTieredRoutingSettings(parsed)
		}
	}
	s.snapshot = &cfg
	s.loadedAt = time.Now()
	return cfg
}

func normalizeTieredRoutingSettings(in TieredRoutingSettings) TieredRoutingSettings {
	ids := make([]int64, 0, len(in.PremiumUserIDs))
	seen := make(map[int64]bool, len(in.PremiumUserIDs))
	for _, id := range in.PremiumUserIDs {
		if id > 0 && !seen[id] {
			seen[id] = true
			ids = append(ids, id)
		}
	}
	in.PremiumUserIDs = ids
	if in.MinBalanceForPremium < 0 {
		in.MinBalanceForPremium = 0
	}
	return in
}

// Invalidate 清空规则快照（写操作后调用）。
func (s *TieredRoutingService) Invalidate() {
	if s == nil {
		return
	}
	s.mu.Lock()
	s.snapshot = nil
	s.loadedAt = time.Time{}
	s.mu.Unlock()
}

// GetSettings 读取当前规则。
func (s *TieredRoutingService) GetSettings(_ context.Context) TieredRoutingSettings {
	if s == nil {
		return DefaultTieredRoutingSettings()
	}
	return s.currentSettings()
}

// UpdateSettings 保存规则。
func (s *TieredRoutingService) UpdateSettings(ctx context.Context, in TieredRoutingSettings) (TieredRoutingSettings, error) {
	if s == nil || s.settingRepo == nil {
		return DefaultTieredRoutingSettings(), nil
	}
	norm := normalizeTieredRoutingSettings(in)
	raw, err := json.Marshal(norm)
	if err != nil {
		return norm, err
	}
	if err := s.settingRepo.Set(ctx, SettingKeyTieredRoutingSettings, string(raw)); err != nil {
		return norm, err
	}
	s.Invalidate()
	return norm, nil
}

// ResolveUserTier 判定用户等级：显式名单或余额达标即 premium，否则 standard。
// 未启用时恒返回 standard（关闭即无行为变化）。
func (s *TieredRoutingService) ResolveUserTier(user *User) string {
	if s == nil || user == nil {
		return TierStandard
	}
	cfg := s.currentSettings()
	if !cfg.Enabled {
		return TierStandard
	}
	return resolveTierByRules(user, cfg)
}

func resolveTierByRules(user *User, cfg TieredRoutingSettings) string {
	if user == nil {
		return TierStandard
	}
	for _, id := range cfg.PremiumUserIDs {
		if user.ID == id {
			return TierPremium
		}
	}
	if cfg.MinBalanceForPremium > 0 && user.Balance >= cfg.MinBalanceForPremium {
		return TierPremium
	}
	return TierStandard
}

// ResolvePool 供中间件调用：解析用户等级对应的目标账号池。
func (s *TieredRoutingService) ResolvePool(user *User) string {
	return s.ResolveUserTier(user)
}

// AccountPool 读取账号所属池：extra.pool 显式标记 premium 才进 premium 池。
func AccountPool(account *Account) string {
	if account == nil || len(account.Extra) == 0 {
		return TierStandard
	}
	if v, ok := account.Extra[AccountPoolExtraKey]; ok {
		if str, ok := v.(string); ok && strings.EqualFold(strings.TrimSpace(str), TierPremium) {
			return TierPremium
		}
	}
	return TierStandard
}

// WithTierPool 把目标账号池写入 ctx（认证中间件在鉴权成功后调用一次）。
func WithTierPool(ctx context.Context, pool string) context.Context {
	return context.WithValue(ctx, ctxkey.TierPool, pool)
}

// TierPoolFromContext 读取目标账号池，缺省 standard（不过滤）。
func TierPoolFromContext(ctx context.Context) string {
	if ctx == nil {
		return TierStandard
	}
	if pool, ok := ctx.Value(ctxkey.TierPool).(string); ok && pool == TierPremium {
		return TierPremium
	}
	return TierStandard
}

// FilterTierPoolAccounts 按目标池过滤候选账号：premium 用户优先 premium 池；
// 池内无可用时回退全量（永不 fail-closed）。
func FilterTierPoolAccounts(accounts []Account, wantPool string) []Account {
	if wantPool != TierPremium || len(accounts) == 0 {
		return accounts
	}
	matched := make([]Account, 0, len(accounts))
	for _, acc := range accounts {
		if AccountPool(&acc) == TierPremium {
			matched = append(matched, acc)
		}
	}
	if len(matched) == 0 {
		return accounts
	}
	return matched
}