junhong_cmp_fiber/internal/service/traffic/query_service.go

// Package traffic 提供流量数据查询服务
// 合并 Redis 缓冲（今日）和 DB（历史）两段数据源，对外提供统一的日粒度流量查询
package traffic

import (
	"context"
	"sort"
	"strconv"
	"time"

	cardflow "github.com/break/junhong_cmp_fiber/internal/domain/cardflow"
	"github.com/break/junhong_cmp_fiber/internal/store/postgres"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
	"github.com/redis/go-redis/v9"
)

// DailyUsageItem 日粒度流量记录
type DailyUsageItem struct {
	Date    string  `json:"date"`     // 日期，格式 YYYY-MM-DD
	UsageMB float64 `json:"usage_mb"` // 当天累计使用流量(MB)
}

// QueryService 流量查询服务
// 未切换日快照算法的卡合并 Redis（当天实时）和 DB（历史落盘）两段数据源；已切换的卡读日流量快照。
type QueryService struct {
	redis               *redis.Client
	cardDailyUsageStore *postgres.CardDailyUsageStore
	switched            func(cardID uint) bool
	flowDays            *postgres.CardFlowDayStore
}

// WithFlowSnapshots 注入日快照数据源：switched 判定卡是否已切换日快照算法，已切换的卡改读快照。
func (s *QueryService) WithFlowSnapshots(switched func(cardID uint) bool, flowDays *postgres.CardFlowDayStore) *QueryService {
	s.switched = switched
	s.flowDays = flowDays
	return s
}

// dailyUsageFromSnapshots 按快照计算每天用量：当天最大快照 − 本月此前最大快照，小于 0 按 0，跨月重新累计。
func (s *QueryService) dailyUsageFromSnapshots(ctx context.Context, cardID uint, startDate, endDate time.Time) ([]*DailyUsageItem, error) {
	start := cardflow.ShanghaiDate(startDate)
	end := cardflow.ShanghaiDate(endDate)
	rows, err := s.flowDays.ListBetween(ctx, cardID, cardflow.MonthStart(start), end)
	if err != nil {
		return nil, err
	}
	items := make([]*DailyUsageItem, 0, len(rows))
	var month string
	var priorMax float64
	for _, row := range rows {
		date := row.FlowDate.Format("2006-01-02")
		if rowMonth := date[:7]; rowMonth != month {
			month, priorMax = rowMonth, 0
		}
		usage := cardflow.DailyUsage(row.MonthUsedMB, priorMax)
		if row.MonthUsedMB > priorMax {
			priorMax = row.MonthUsedMB
		}
		if date >= start.Format("2006-01-02") {
			items = append(items, &DailyUsageItem{Date: date, UsageMB: usage})
		}
	}
	return items, nil
}

// NewQueryService 创建流量查询服务实例
func NewQueryService(redisClient *redis.Client, cardDailyUsageStore *postgres.CardDailyUsageStore) *QueryService {
	return &QueryService{
		redis:               redisClient,
		cardDailyUsageStore: cardDailyUsageStore,
	}
}

// GetDailyUsage 查询卡在指定日期范围内的日粒度流量
// 今天的数据从 Redis 缓冲读取（可能尚未落盘），历史数据从 DB 读取，合并后按日期 ASC 排序返回
func (s *QueryService) GetDailyUsage(ctx context.Context, cardID uint, startDate, endDate time.Time) ([]*DailyUsageItem, error) {
	if s.switched != nil && s.flowDays != nil && s.switched(cardID) {
		return s.dailyUsageFromSnapshots(ctx, cardID, startDate, endDate)
	}
	now := time.Now()
	todayStr := now.Format("2006-01-02")

	// 判断结束日期是否包含今天
	endDateStr := endDate.Format("2006-01-02")
	includeToday := endDateStr >= todayStr

	// 历史部分：从 DB 查询（排除今天，因为今天的数据可能还在 Redis 中）
	dbEndDate := endDate
	if includeToday {
		// DB 只查到昨天，今天从 Redis 取
		dbEndDate = now.AddDate(0, 0, -1)
	}

	var items []*DailyUsageItem

	// 只有 startDate <= dbEndDate 时才查 DB
	if !startDate.After(dbEndDate) {
		records, err := s.cardDailyUsageStore.ListByCardIDAndDateRange(ctx, cardID, startDate, dbEndDate)
		if err != nil {
			return nil, err
		}
		for _, r := range records {
			items = append(items, &DailyUsageItem{
				Date:    r.Date.Format("2006-01-02"),
				UsageMB: r.UsageMB,
			})
		}
	}

	// 今天部分：从 Redis 读取
	if includeToday && !startDate.After(now) {
		key := constants.RedisCardDailyTrafficKey(cardID, todayStr)
		val, err := s.redis.Get(ctx, key).Result()
		if err == nil && val != "" {
			if usageMB, parseErr := strconv.ParseFloat(val, 64); parseErr == nil && usageMB > 0 {
				items = append(items, &DailyUsageItem{
					Date:    todayStr,
					UsageMB: usageMB,
				})
			}
		}
		// Redis key 不存在（redis.Nil）时忽略，表示今天暂无流量
	}

	// 按日期 ASC 排序
	sort.Slice(items, func(i, j int) bool {
		return items[i].Date < items[j].Date
	})

	return items, nil
}