thread-master/apps/backend/internal/module/usage/repository/mongo.go

package repository

import (
	"context"
	"fmt"

	libmongo "apps/backend/internal/lib/mongo"
	"apps/backend/internal/module/usage/domain"

	"github.com/zeromicro/go-zero/core/stores/mon"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type MonStore struct {
	events    *mon.Model
	counters  *mon.Model
	prefs     *mon.Model
	purchases *mon.Model
}

func NewMonStore(uri, database string) *MonStore {
	uri = libmongo.MustMongoURI(uri)
	return &MonStore{
		events:    mon.MustNewModel(uri, database, "usage_events"),
		counters:  mon.MustNewModel(uri, database, "usage_monthly_counters"),
		prefs:     mon.MustNewModel(uri, database, "usage_prefs"),
		purchases: mon.MustNewModel(uri, database, "usage_purchases"),
	}
}

func monthlyCounterID(uid int64, monthKey string) string {
	return fmt.Sprintf("%d:%s", uid, monthKey)
}

func meterCounterField(meter string) (string, error) {
	switch meter {
	case domain.MeterAICopy, domain.MeterAIResearch, domain.MeterWebSearch, domain.MeterAIImage:
		return "by_meter." + meter, nil
	default:
		return "", fmt.Errorf("invalid usage meter %q", meter)
	}
}

func (s *MonStore) ensureMonthlyCounter(ctx context.Context, uid int64, monthKey string) error {
	id := monthlyCounterID(uid, monthKey)
	var existing domain.MonthlyCounter
	if err := s.counters.FindOne(ctx, &existing, bson.M{"_id": id}); err == nil {
		return nil
	} else if err != mon.ErrNotFound {
		return err
	}
	events, err := s.ListEvents(ctx, uid, monthKey, domain.KeyModePlatform, 0)
	if err != nil {
		return err
	}
	counter := &domain.MonthlyCounter{
		ID: id, UID: uid, MonthKey: monthKey, ByMeter: map[string]int{}, UpdatedAt: domain.NowNano(),
	}
	for _, event := range events {
		credits := event.Credits
		if minimum := domain.MeterCost(event.Meter); credits < minimum {
			credits = minimum
		}
		counter.TotalCredits += credits
		counter.ByMeter[event.Meter] += credits
	}
	if _, err := s.counters.InsertOne(ctx, counter); err != nil && !mongo.IsDuplicateKeyError(err) {
		return err
	}
	return nil
}

func (s *MonStore) ReservePlatform(ctx context.Context, uid int64, monthKey, meter string, cost, totalLimit, meterLimit int) error {
	field, err := meterCounterField(meter)
	if err != nil {
		return err
	}
	if totalLimit >= 0 && cost > totalLimit {
		return domain.ErrQuotaExceeded
	}
	if meterLimit >= 0 && cost > meterLimit {
		return domain.ErrMeterCapExceeded
	}
	if err := s.ensureMonthlyCounter(ctx, uid, monthKey); err != nil {
		return err
	}
	filter := bson.M{"_id": monthlyCounterID(uid, monthKey)}
	if totalLimit >= 0 {
		filter["total_credits"] = bson.M{"$lte": totalLimit - cost}
	}
	if meterLimit >= 0 {
		filter["$or"] = bson.A{
			bson.M{field: bson.M{"$exists": false}},
			bson.M{field: bson.M{"$lte": meterLimit - cost}},
		}
	}
	res, err := s.counters.UpdateOne(ctx, filter, bson.M{
		"$inc": bson.M{"total_credits": cost, field: cost},
		"$set": bson.M{"updated_at": domain.NowNano()},
	})
	if err != nil {
		return err
	}
	if res.MatchedCount != 0 {
		return nil
	}
	counter, err := s.GetMonthlyCounter(ctx, uid, monthKey)
	if err != nil {
		return err
	}
	if totalLimit >= 0 && counter.TotalCredits+cost > totalLimit {
		return domain.ErrQuotaExceeded
	}
	if meterLimit >= 0 && counter.ByMeter[meter]+cost > meterLimit {
		return domain.ErrMeterCapExceeded
	}
	return fmt.Errorf("platform quota reservation did not match counter")
}

func (s *MonStore) GetMonthlyCounter(ctx context.Context, uid int64, monthKey string) (*domain.MonthlyCounter, error) {
	var counter domain.MonthlyCounter
	err := s.counters.FindOne(ctx, &counter, bson.M{"_id": monthlyCounterID(uid, monthKey)})
	if err == mon.ErrNotFound {
		return nil, domain.ErrNotFound
	}
	if err != nil {
		return nil, err
	}
	if counter.ByMeter == nil {
		counter.ByMeter = map[string]int{}
	}
	return &counter, nil
}

func (s *MonStore) InsertEvent(ctx context.Context, e *domain.Event) error {
	_, err := s.events.InsertOne(ctx, e)
	return err
}

func (s *MonStore) ListEvents(ctx context.Context, uid int64, monthKey, keyMode string, limit int) ([]*domain.Event, error) {
	filter := bson.M{"uid": uid}
	if monthKey != "" {
		filter["month_key"] = monthKey
	}
	if keyMode != "" && keyMode != "all" {
		filter["key_mode"] = keyMode
	}
	opts := options.Find().SetSort(bson.D{{Key: "created_at", Value: -1}})
	if limit > 0 {
		opts.SetLimit(int64(limit))
	}
	var list []*domain.Event
	err := s.events.Find(ctx, &list, filter, opts)
	return list, err
}

func (s *MonStore) ListEventsAll(ctx context.Context, monthKey string, limit int) ([]*domain.Event, error) {
	filter := bson.M{}
	if monthKey != "" {
		filter["month_key"] = monthKey
	}
	opts := options.Find().SetSort(bson.D{{Key: "created_at", Value: -1}})
	if limit > 0 {
		opts.SetLimit(int64(limit))
	}
	var list []*domain.Event
	err := s.events.Find(ctx, &list, filter, opts)
	return list, err
}

func (s *MonStore) ListEventsRange(ctx context.Context, fromNano, toExclusiveNano int64) ([]*domain.Event, error) {
	var list []*domain.Event
	err := s.events.Find(ctx, &list, bson.M{
		"created_at": bson.M{"$gte": fromNano, "$lt": toExclusiveNano},
	}, options.Find().SetSort(bson.D{{Key: "created_at", Value: 1}}))
	return list, err
}

func (s *MonStore) GetPrefs(ctx context.Context, uid int64) (*domain.MemberPrefs, error) {
	var p domain.MemberPrefs
	err := s.prefs.FindOne(ctx, &p, bson.M{"uid": uid})
	if err != nil {
		if err == mon.ErrNotFound {
			return &domain.MemberPrefs{UID: uid, PlanID: domain.PlanFree}, nil
		}
		return nil, err
	}
	return &p, nil
}

func (s *MonStore) SavePrefs(ctx context.Context, p *domain.MemberPrefs) error {
	_, err := s.prefs.ReplaceOne(ctx, bson.M{"uid": p.UID}, p, options.Replace().SetUpsert(true))
	return err
}

func (s *MonStore) InsertPurchase(ctx context.Context, p *domain.Purchase) error {
	_, err := s.purchases.InsertOne(ctx, p)
	return err
}

func (s *MonStore) ListPurchases(ctx context.Context, uid int64, limit int) ([]*domain.Purchase, error) {
	opts := options.Find().SetSort(bson.D{{Key: "created_at", Value: -1}})
	if limit > 0 {
		opts.SetLimit(int64(limit))
	}
	var list []*domain.Purchase
	err := s.purchases.Find(ctx, &list, bson.M{"uid": uid}, opts)
	return list, err
}

var _ domain.Repository = (*MonStore)(nil)