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) }