openteam/server/internal/billing/service.go

package billing

import (
	"errors"
	"time"

	"github.com/shopspring/decimal"
	"go.uber.org/zap"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	"openteam/server/internal/store"
)

var (
	ErrInsufficientBalance = errors.New("insufficient balance")
	ErrUserNotFound        = errors.New("user not found")
)

// PriceSnapshot mirrors a model's prices at billing time.
type PriceSnapshot struct {
	InputPrice     decimal.Decimal
	OutputPrice    decimal.Decimal
	CacheReadPrice decimal.Decimal
}

// CostFromPrices computes cost for token counts using per-1M-token prices.
func CostFromPrices(in, out, cacheRead int64, prices PriceSnapshot) decimal.Decimal {
	perM := decimal.NewFromInt(1_000_000)
	cost := prices.InputPrice.Mul(decimal.NewFromInt(in)).Div(perM).
		Add(prices.OutputPrice.Mul(decimal.NewFromInt(out)).Div(perM)).
		Add(prices.CacheReadPrice.Mul(decimal.NewFromInt(cacheRead)).Div(perM))
	return cost.Round(8)
}

// EstimateCost approximates cost from token estimates.
func (s *Service) EstimateCost(modelID int64, in, out, cacheRead int64) (decimal.Decimal, error) {
	var m store.Model
	if err := s.db.First(&m, modelID).Error; err != nil {
		return decimal.Zero, err
	}
	return CostFromPrices(in, out, cacheRead, PriceSnapshot{
		InputPrice: m.InputPrice, OutputPrice: m.OutputPrice, CacheReadPrice: m.CacheReadPrice,
	}), nil
}

type Service struct {
	db  *gorm.DB
	log *zap.Logger
}

func NewService(db *gorm.DB, log *zap.Logger) *Service {
	return &Service{db: db, log: log}
}

// CheckBalance returns whether the user can afford the estimated cost.
func (s *Service) CheckBalance(userID int64, estimated decimal.Decimal) error {
	var u store.User
	if err := s.db.First(&u, userID).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return ErrUserNotFound
		}
		return err
	}
	if u.Balance.LessThan(estimated) {
		return ErrInsufficientBalance
	}
	return nil
}

// Deduct atomically debits the user balance and appends a ledger entry.
func (s *Service) Deduct(userID int64, change decimal.Decimal, logType, refID string) (decimal.Decimal, error) {
	return s.mutateBalance(userID, change.Neg(), logType, refID)
}

// Credit adds to the user balance (recharge / refund / admin adjust).
func (s *Service) Credit(userID int64, change decimal.Decimal, logType, refID string) (decimal.Decimal, error) {
	return s.mutateBalance(userID, change, logType, refID)
}

// mutateBalance performs the balance update in a transaction so ledger and
// balance always agree.
func (s *Service) mutateBalance(userID int64, delta decimal.Decimal, logType, refID string) (decimal.Decimal, error) {
	var after decimal.Decimal
	err := s.db.Transaction(func(tx *gorm.DB) error {
		var u store.User
		if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&u, userID).Error; err != nil {
			return err
		}
		newBal := u.Balance.Add(delta)
		if newBal.LessThan(decimal.Zero) {
			return ErrInsufficientBalance
		}
		if err := tx.Model(&u).Update("balance", newBal).Error; err != nil {
			return err
		}
		after = newBal
		return tx.Create(&store.BalanceLog{
			UserID:       userID,
			Change:       delta,
			BalanceAfter: newBal,
			Type:         logType,
			RefID:        refID,
			CreatedAt:    time.Now(),
		}).Error
	})
	if err != nil {
		return decimal.Zero, err
	}
	return after, nil
}

// AdminAdjust changes a user's balance with an optional reason.
func (s *Service) AdminAdjust(userID int64, amount decimal.Decimal, reason string) (decimal.Decimal, error) {
	return s.mutateBalance(userID, amount, "admin_adjust", reason)
}