ente/server/pkg/controller/storagebonus/referral.go

package storagebonus

import (
	"database/sql"
	"errors"
	"fmt"

	"github.com/ente-io/museum/ente"
	entity "github.com/ente-io/museum/ente/storagebonus"
	"github.com/ente-io/museum/pkg/controller/email"
	"github.com/ente-io/museum/pkg/controller/lock"

	"github.com/ente-io/museum/pkg/repo"
	"github.com/ente-io/museum/pkg/repo/storagebonus"
	"github.com/ente-io/museum/pkg/utils/auth"
	enteTime "github.com/ente-io/museum/pkg/utils/time"
	"github.com/ente-io/stacktrace"
	"github.com/gin-gonic/gin"
)

const (
	codeLength                 = 6
	referralAmountInGb         = 10
	maxClaimableReferralAmount = 2000
	numOfDaysToClaimReferral   = 32
	defaultPlanType            = entity.TenGbOnUpgrade
)

// Controller exposes functions to interact with family module
type Controller struct {
	UserRepo                    *repo.UserRepository
	StorageBonus                *storagebonus.Repository
	LockController              *lock.LockController
	CronRunning                 bool
	EmailNotificationController *email.EmailNotificationController
}

func (c *Controller) GetUserReferralView(ctx *gin.Context) (*entity.GetUserReferralView, error) {
	// Get the user id from the context
	userID := auth.GetUserID(ctx.Request.Header)

	// Use goroutines to fetch UserRepo.Get, HasAppliedReferral
	user, err := c.UserRepo.Get(userID)
	if err != nil {
		return nil, stacktrace.Propagate(err, "failed to get user")
	}
	appliedReferral, err := c.StorageBonus.HasAppliedReferral(ctx, userID)
	if err != nil {
		return nil, stacktrace.Propagate(err, "")
	}
	isFamilyMember := user.FamilyAdminID != nil && *user.FamilyAdminID != userID
	enableApplyCode := !appliedReferral && user.CreationTime > enteTime.MicrosecondBeforeDays(numOfDaysToClaimReferral) && !isFamilyMember
	// Get the referral code for the user or family admin
	codeUser := userID
	if isFamilyMember {
		codeUser = *user.FamilyAdminID
	}
	referralCode, err2 := c.GetOrCreateReferralCode(ctx, codeUser)
	if err2 != nil {
		return nil, stacktrace.Propagate(err2, "failed to get or create referral code")
	}
	storageClaimed, err2 := c.GetActiveReferralBonusValue(ctx, codeUser)
	if err2 != nil {
		return nil, stacktrace.Propagate(err2, "failed to get storage claimed")
	}

	return &entity.GetUserReferralView{
		PlanInfo: entity.PlanInfo{
			IsEnabled:               true,
			PlanType:                defaultPlanType,
			StorageInGB:             referralAmountInGb,
			MaxClaimableStorageInGB: maxClaimableReferralAmount,
		},
		Code:            referralCode,
		EnableApplyCode: enableApplyCode,
		IsFamilyMember:  isFamilyMember,
		HasAppliedCode:  appliedReferral,
		ClaimedStorage:  *storageClaimed,
	}, nil
}

func (c *Controller) ApplyReferralCode(ctx *gin.Context, code string) error {
	// Get user id from the context
	userID := auth.GetUserID(ctx.Request.Header)
	user, err := c.UserRepo.Get(userID)
	if err != nil {
		return stacktrace.Propagate(err, "failed to get user")
	}

	codeOwnerID, err := c.StorageBonus.GetUserIDByCode(ctx, code)
	if err != nil {
		return stacktrace.Propagate(err, "failed to get user id by code")
	}
	// Verify that the codeOwnerID is not deleted yet
	_, err = c.UserRepo.Get(*codeOwnerID)
	if err != nil {
		if errors.Is(err, ente.ErrUserDeleted) {
			return stacktrace.Propagate(entity.InvalidCodeErr, "code belongs to deleted user")
		}
		return stacktrace.Propagate(err, "failed to get user")
	}

	if user.CreationTime < enteTime.MicrosecondBeforeDays(numOfDaysToClaimReferral) {
		return stacktrace.Propagate(entity.CanNotApplyCodeErr, "account is too old to apply code")
	} else if user.FamilyAdminID != nil && userID != *user.FamilyAdminID {
		return stacktrace.Propagate(entity.CanNotApplyCodeErr, "user is member of a family plan")
	}

	err = c.StorageBonus.TrackReferralAndInviteeBonus(ctx, userID, *codeOwnerID, defaultPlanType)
	if err != nil {
		return stacktrace.Propagate(err, "failed to apply code")
	}
	return nil
}

func (c *Controller) GetOrCreateReferralCode(ctx *gin.Context, userID int64) (*string, error) {
	referralCode, err := c.StorageBonus.GetCode(ctx, userID)
	if err != nil {
		if !errors.Is(err, sql.ErrNoRows) {
			return nil, stacktrace.Propagate(err, "failed to get storagebonus code")
		}
		code, err := generateAlphaNumString(codeLength)
		if err != nil {
			return nil, stacktrace.Propagate(err, "")
		}
		err = c.StorageBonus.InsertCode(ctx, userID, code)
		if err != nil {
			return nil, stacktrace.Propagate(err, "failed to insert storagebonus code")
		}
		referralCode = &code
	}
	return referralCode, nil
}

// generateAlphaNumString returns AlphaNumeric code of given length
// which exclude number 0 and letter O. The code always starts with an
// alphabet
func generateAlphaNumString(length int) (string, error) {
	// Define the alphabet and numbers to be used in the string.
	alphabet := "ABCDEFGHIJKLMNPQRSTUVWXYZ"
	// Define the alphabet and numbers to be used in the string.
	alphaNum := fmt.Sprintf("%s123456789", alphabet)
	// Allocate a byte slice with the desired length.
	result := make([]byte, length)
	// Generate the first letter as an alphabet.
	r0, err := auth.GenerateRandomInt(int64(len(alphabet)))
	if err != nil {
		return "", stacktrace.Propagate(err, "")
	}
	result[0] = alphabet[r0]
	// Generate the remaining characters as alphanumeric.
	for i := 1; i < length; i++ {
		ri, err := auth.GenerateRandomInt(int64(len(alphaNum)))
		if err != nil {
			return "", stacktrace.Propagate(err, "")
		}
		result[i] = alphaNum[ri]
	}
	return string(result), nil
}