From 45bf75709d52e4af73bc383a323369ccc2342f40 Mon Sep 17 00:00:00 2001
From: Zhang Minghan <zmh13054618081@dingtalk.com>
Date: Mon, 11 Sep 2023 14:44:22 +0800
Subject: [PATCH] update subscription

---
 api/chat.go                      | 16 ++++++---
 api/image.go                     | 25 ++++++-------
 api/stream.go                    | 12 +++++--
 app/src/conversation/addition.ts |  6 ++--
 app/src/routes/Subscription.tsx  | 28 +++++++--------
 auth/controller.go               | 60 +++++++++++++++++++++++++++++---
 auth/subscription.go             | 49 ++++++++++++++++++++++++++
 auth/usage.go                    |  9 +++++
 auth/user.go                     | 31 +++++++++++++++++
 connection/database.go           | 18 ++++++++++
 main.go                          |  2 ++
 middleware/throttle.go           |  2 ++
 utils/char.go                    |  4 +++
 13 files changed, 218 insertions(+), 44 deletions(-)
 create mode 100644 auth/subscription.go

diff --git a/api/chat.go b/api/chat.go
index 00266210..3c95365d 100644
--- a/api/chat.go
+++ b/api/chat.go
@@ -29,7 +29,7 @@ func SendSegmentMessage(conn *websocket.Conn, message types.ChatGPTSegmentRespon
 	_ = conn.WriteMessage(websocket.TextMessage, []byte(utils.ToJson(message)))
 }
 
-func TextChat(db *sql.DB, user *auth.User, conn *websocket.Conn, instance *conversation.Conversation) string {
+func TextChat(db *sql.DB, cache *redis.Client, user *auth.User, conn *websocket.Conn, instance *conversation.Conversation) string {
 	var keyword string
 	var segment []types.ChatGPTMessage
 
@@ -41,7 +41,8 @@ func TextChat(db *sql.DB, user *auth.User, conn *websocket.Conn, instance *conve
 
 	SendSegmentMessage(conn, types.ChatGPTSegmentResponse{Keyword: keyword, End: false})
 
-	if instance.IsEnableGPT4() && !auth.CanEnableGPT4(db, user) {
+	isProPlan := auth.CanEnableSubscription(db, cache, user)
+	if instance.IsEnableGPT4() && (!isProPlan) && (!auth.CanEnableGPT4(db, user)) {
 		SendSegmentMessage(conn, types.ChatGPTSegmentResponse{
 			Message: defaultQuotaMessage,
 			Quota:   0,
@@ -51,7 +52,7 @@ func TextChat(db *sql.DB, user *auth.User, conn *websocket.Conn, instance *conve
 	}
 
 	buffer := NewBuffer(instance.IsEnableGPT4(), segment)
-	StreamRequest(instance.IsEnableGPT4(), segment, 2000, func(resp string) {
+	StreamRequest(instance.IsEnableGPT4(), isProPlan, segment, 2000, func(resp string) {
 		SendSegmentMessage(conn, types.ChatGPTSegmentResponse{
 			Message: buffer.Write(resp),
 			Quota:   buffer.GetQuota(),
@@ -59,6 +60,9 @@ func TextChat(db *sql.DB, user *auth.User, conn *websocket.Conn, instance *conve
 		})
 	})
 	if buffer.IsEmpty() {
+		if isProPlan {
+			auth.DecreaseSubscriptionUsage(cache, user)
+		}
 		SendSegmentMessage(conn, types.ChatGPTSegmentResponse{
 			Message: defaultErrorMessage,
 			Quota:   -0xe, // special value for error
@@ -68,7 +72,9 @@ func TextChat(db *sql.DB, user *auth.User, conn *websocket.Conn, instance *conve
 	}
 
 	// collect quota
-	user.UseQuota(db, buffer.GetQuota())
+	if !isProPlan {
+		user.UseQuota(db, buffer.GetQuota())
+	}
 	SendSegmentMessage(conn, types.ChatGPTSegmentResponse{End: true, Quota: buffer.GetQuota()})
 
 	return buffer.ReadWithDefault(defaultErrorMessage)
@@ -111,7 +117,7 @@ func ChatHandler(conn *websocket.Conn, instance *conversation.Conversation, user
 	if strings.HasPrefix(instance.GetLatestMessage(), "/image") {
 		return ImageChat(conn, instance, user, db, cache)
 	} else {
-		return TextChat(db, user, conn, instance)
+		return TextChat(db, cache, user, conn, instance)
 	}
 }
 
diff --git a/api/image.go b/api/image.go
index 5bfaa61e..21c1d4e0 100644
--- a/api/image.go
+++ b/api/image.go
@@ -48,26 +48,21 @@ func GetImageWithCache(ctx context.Context, prompt string, cache *redis.Client)
 }
 
 func GetLimitFormat(id int64) string {
-	t := time.Now().Format("2006-01-02")
-	return fmt.Sprintf(":imagelimit:%s:%d", t, id)
+	today := time.Now().Format("2006-01-02")
+	return fmt.Sprintf(":imagelimit:%s:%d", today, id)
 }
 
 func GetImageWithUserLimit(user *auth.User, prompt string, db *sql.DB, cache *redis.Client) (string, error) {
-	// 5 images one day per user (count by cache)
-	res, err := cache.Get(context.Background(), GetLimitFormat(user.GetID(db))).Result()
-	if err != nil || len(res) == 0 || res == "" {
-		cache.Set(context.Background(), GetLimitFormat(user.GetID(db)), "1", time.Hour*24)
-		return GetImageWithCache(context.Background(), prompt, cache)
-	}
+	// free plan: 5 images per day
+	// pro plan: 50 images per day
 
-	if res == "5" {
-		if auth.ReduceDalle(db, user) {
-			return GetImageWithCache(context.Background(), prompt, cache)
-		}
-		return "", fmt.Errorf("you have reached your limit of 5 free images per day, please buy more dalle usage or wait until tomorrow")
-	} else {
-		cache.Set(context.Background(), GetLimitFormat(user.GetID(db)), fmt.Sprintf("%d", utils.ToInt(res)+1), time.Hour*24)
+	key := GetLimitFormat(user.GetID(db))
+	usage := auth.GetDalleUsageLimit(db, user)
+
+	if utils.IncrWithLimit(cache, key, 1, int64(usage), 60*60*24) || auth.ReduceDalle(db, user) {
 		return GetImageWithCache(context.Background(), prompt, cache)
+	} else {
+		return "", fmt.Errorf("you have reached your limit of %d free images per day, please buy more quota or wait until tomorrow", usage)
 	}
 }
 
diff --git a/api/stream.go b/api/stream.go
index bab6ea05..8c93897b 100644
--- a/api/stream.go
+++ b/api/stream.go
@@ -56,7 +56,7 @@ func NativeStreamRequest(model string, endpoint string, apikeys string, messages
 
 	client := &http.Client{}
 	req, err := http.NewRequest("POST", endpoint+"/chat/completions", utils.ConvertBody(types.ChatGPTRequest{
-		Model:    model,
+		Model:    strings.Replace(model, "reverse", "free", -1),
 		Messages: messages,
 		MaxToken: token,
 		Stream:   true,
@@ -93,9 +93,15 @@ func NativeStreamRequest(model string, endpoint string, apikeys string, messages
 	}
 }
 
-func StreamRequest(enableGPT4 bool, messages []types.ChatGPTMessage, token int, callback func(string)) {
+func StreamRequest(enableGPT4 bool, isProPlan bool, messages []types.ChatGPTMessage, token int, callback func(string)) {
+	var model string
+	if isProPlan {
+		model = "gpt-4-reverse" // using reverse engine
+	} else {
+		model = "gpt-4"
+	}
 	if enableGPT4 {
-		NativeStreamRequest("gpt-4", viper.GetString("openai.gpt4_endpoint"), viper.GetString("openai.gpt4"), messages, token, callback)
+		NativeStreamRequest(model, viper.GetString("openai.gpt4_endpoint"), viper.GetString("openai.gpt4"), messages, token, callback)
 	} else {
 		NativeStreamRequest("gpt-3.5-turbo-16k-0613", viper.GetString("openai.user_endpoint"), viper.GetString("openai.user"), messages, token, callback)
 	}
diff --git a/app/src/conversation/addition.ts b/app/src/conversation/addition.ts
index 5bda6933..971cd271 100644
--- a/app/src/conversation/addition.ts
+++ b/app/src/conversation/addition.ts
@@ -57,8 +57,8 @@ export async function getSubscription(): Promise<SubscriptionResponse> {
     }
     return {
       status: resp.data.status,
-      is_subscribed: resp.data.data.is_subscribed,
-      expired: resp.data.data.expired,
+      is_subscribed: resp.data.is_subscribed,
+      expired: resp.data.expired,
     };
   } catch (e) {
     console.debug(e);
@@ -70,7 +70,7 @@ export async function buySubscription(
   month: number,
 ): Promise<BuySubscriptionResponse> {
   try {
-    const resp = await axios.post(`/subscription`, { month });
+    const resp = await axios.post(`/subscribe`, { month });
     return resp.data as BuySubscriptionResponse;
   } catch (e) {
     console.debug(e);
diff --git a/app/src/routes/Subscription.tsx b/app/src/routes/Subscription.tsx
index 4acc1dfe..0813c090 100644
--- a/app/src/routes/Subscription.tsx
+++ b/app/src/routes/Subscription.tsx
@@ -76,16 +76,16 @@ function Upgrade({ children }: UpgradeProps) {
         </DialogHeader>
         <div className="upgrade-wrapper">
           <Select onValueChange={
-            (value: number) => setMonth(value)
+            (value: string) => setMonth(parseInt(value))
           }>
             <SelectTrigger className="w-[200px]">
               <SelectValue placeholder={t(`sub.time.${month}`)} />
             </SelectTrigger>
             <SelectContent>
-              <SelectItem value={1}>{t(`sub.time.1`)}</SelectItem>
-              <SelectItem value={3}>{t(`sub.time.3`)}</SelectItem>
-              <SelectItem value={6}>{t(`sub.time.6`)}</SelectItem>
-              <SelectItem value={12}>
+              <SelectItem value={"1"}>{t(`sub.time.1`)}</SelectItem>
+              <SelectItem value={"3"}>{t(`sub.time.3`)}</SelectItem>
+              <SelectItem value={"6"}>{t(`sub.time.6`)}</SelectItem>
+              <SelectItem value={"12"}>
                 {t(`sub.time.12`)}
                 <Badge className={`ml-2 cent`}>{t(`percent`, { cent: 9 })}</Badge>
               </SelectItem>
@@ -97,16 +97,16 @@ function Upgrade({ children }: UpgradeProps) {
           <Button variant={`outline`} onClick={
             () => setOpen(false)
           }>{ t('cancel') }</Button>
-          <Button>
-            <Plus className={`h-4 w-4 mr-1`} onClick={
-              async () => {
-                const res = await callBuyAction(t, toast, month);
-                if (res) {
-                  setOpen(false);
-                  await refreshSubscription(dispatch);
-                }
+          <Button onClick={
+            async () => {
+              const res = await callBuyAction(t, toast, month);
+              if (res) {
+                setOpen(false);
+                await refreshSubscription(dispatch);
               }
-            } />
+            }
+          }>
+            <Plus className={`h-4 w-4 mr-1`} />
             { t('confirm') }
           </Button>
         </DialogFooter>
diff --git a/auth/controller.go b/auth/controller.go
index 44d9db72..db6aef24 100644
--- a/auth/controller.go
+++ b/auth/controller.go
@@ -9,6 +9,10 @@ type BuyForm struct {
 	Quota int `json:"quota" binding:"required"`
 }
 
+type SubscribeForm struct {
+	Month int `json:"month" binding:"required"`
+}
+
 func GetUserByCtx(c *gin.Context) *User {
 	user := c.MustGet("user").(string)
 	if len(user) == 0 {
@@ -50,6 +54,57 @@ func QuotaAPI(c *gin.Context) {
 	})
 }
 
+func SubscriptionAPI(c *gin.Context) {
+	user := GetUserByCtx(c)
+	if user == nil {
+		return
+	}
+
+	db := utils.GetDBFromContext(c)
+	c.JSON(200, gin.H{
+		"status":        true,
+		"is_subscribed": user.IsSubscribe(db),
+		"expired":       user.GetSubscriptionExpiredDay(db),
+	})
+}
+
+func SubscribeAPI(c *gin.Context) {
+	user := GetUserByCtx(c)
+	if user == nil {
+		return
+	}
+
+	db := utils.GetDBFromContext(c)
+	var form SubscribeForm
+	if err := c.ShouldBindJSON(&form); err != nil {
+		c.JSON(200, gin.H{
+			"status": false,
+			"error":  err.Error(),
+		})
+		return
+	}
+
+	if form.Month <= 0 || form.Month > 999 {
+		c.JSON(200, gin.H{
+			"status": false,
+			"error":  "invalid month range (1 ~ 999)",
+		})
+		return
+	}
+
+	if BuySubscription(db, user, form.Month) {
+		c.JSON(200, gin.H{
+			"status": true,
+			"error":  "success",
+		})
+	} else {
+		c.JSON(200, gin.H{
+			"status": false,
+			"error":  "not enough money",
+		})
+	}
+}
+
 func BuyAPI(c *gin.Context) {
 	user := GetUserByCtx(c)
 	if user == nil {
@@ -74,10 +129,7 @@ func BuyAPI(c *gin.Context) {
 		return
 	}
 
-	money := float32(form.Quota) * 0.1
-	if Pay(user.Username, money) {
-		user.IncreaseQuota(db, float32(form.Quota))
-
+	if BuyQuota(db, user, form.Quota) {
 		c.JSON(200, gin.H{
 			"status": true,
 			"error":  "success",
diff --git a/auth/subscription.go b/auth/subscription.go
new file mode 100644
index 00000000..4bc3db63
--- /dev/null
+++ b/auth/subscription.go
@@ -0,0 +1,49 @@
+package auth
+
+import (
+	"chat/utils"
+	"database/sql"
+	"fmt"
+	"github.com/go-redis/redis/v8"
+	"time"
+)
+
+func CountSubscriptionPrize(month int) float32 {
+	if month >= 12 {
+		return 8 * float32(month) * 0.9
+	}
+	return 8 * float32(month)
+}
+
+func BuySubscription(db *sql.DB, user *User, month int) bool {
+	if month < 1 || month > 999 {
+		return false
+	}
+	money := CountSubscriptionPrize(month)
+	if Pay(user.Username, money) {
+		user.AddSubscription(db, month)
+		return true
+	}
+	return false
+}
+
+func IncreaseSubscriptionUsage(cache *redis.Client, user *User) bool {
+	today := time.Now().Format("2006-01-02")
+	return utils.IncrWithLimit(cache, fmt.Sprintf(":subscription-usage:%s:%d", today, user.ID), 1, 999, 60*60*24) // 1 day
+}
+
+func DecreaseSubscriptionUsage(cache *redis.Client, user *User) bool {
+	today := time.Now().Format("2006-01-02")
+	return utils.DecrInt(cache, fmt.Sprintf(":subscription-usage:%s:%d", today, user.ID), 1)
+}
+
+func CanEnableSubscription(db *sql.DB, cache *redis.Client, user *User) bool {
+	return user.IsSubscribe(db) && IncreaseSubscriptionUsage(cache, user)
+}
+
+func GetDalleUsageLimit(db *sql.DB, user *User) int {
+	if user.IsSubscribe(db) {
+		return 50
+	}
+	return 5
+}
diff --git a/auth/usage.go b/auth/usage.go
index 0a8c78e5..3f66f6bb 100644
--- a/auth/usage.go
+++ b/auth/usage.go
@@ -37,3 +37,12 @@ func ReduceDalle(db *sql.DB, user *User) bool {
 func CanEnableGPT4(db *sql.DB, user *User) bool {
 	return user.GetQuota(db) >= 5
 }
+
+func BuyQuota(db *sql.DB, user *User, quota int) bool {
+	money := float32(quota) * 0.1
+	if Pay(user.Username, money) {
+		user.IncreaseQuota(db, float32(quota))
+		return true
+	}
+	return false
+}
diff --git a/auth/user.go b/auth/user.go
index 0cd65c8e..0ba7a1f2 100644
--- a/auth/user.go
+++ b/auth/user.go
@@ -8,6 +8,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/go-redis/redis/v8"
 	"github.com/spf13/viper"
+	"math"
 	"net/http"
 	"time"
 )
@@ -128,6 +129,36 @@ func (u *User) UseQuota(db *sql.DB, quota float32) bool {
 	return u.IncreaseUsedQuota(db, quota)
 }
 
+func (u *User) GetSubscription(db *sql.DB) time.Time {
+	var expiredAt []uint8
+	if err := db.QueryRow("SELECT expired_at FROM subscription WHERE user_id = ?", u.GetID(db)).Scan(&expiredAt); err != nil {
+		return time.Unix(0, 0)
+	}
+	return *utils.ConvertTime(expiredAt)
+}
+
+func (u *User) IsSubscribe(db *sql.DB) bool {
+	return u.GetSubscription(db).Unix() > time.Now().Unix()
+}
+
+func (u *User) GetSubscriptionExpiredDay(db *sql.DB) int {
+	stamp := u.GetSubscription(db).Sub(time.Now())
+	return int(math.Round(stamp.Hours() / 24))
+}
+
+func (u *User) AddSubscription(db *sql.DB, month int) bool {
+	current := u.GetSubscription(db)
+	if current.Unix() < time.Now().Unix() {
+		current = time.Now()
+	}
+	expiredAt := current.AddDate(0, month, 0)
+	_, err := db.Exec(`
+		INSERT INTO subscription (user_id, expired_at, total_month) VALUES (?, ?, ?) 
+		ON DUPLICATE KEY UPDATE expired_at = ?, total_month = total_month + ?
+	`, u.GetID(db), utils.ConvertSqlTime(expiredAt), month, utils.ConvertSqlTime(expiredAt), month)
+	return err == nil
+}
+
 func IsUserExist(db *sql.DB, username string) bool {
 	var count int
 	if err := db.QueryRow("SELECT COUNT(*) FROM auth WHERE username = ?", username).Scan(&count); err != nil {
diff --git a/connection/database.go b/connection/database.go
index 8f2d561f..370d2da5 100644
--- a/connection/database.go
+++ b/connection/database.go
@@ -30,6 +30,7 @@ func ConnectMySQL() *sql.DB {
 	CreateConversationTable(db)
 	CreatePackageTable(db)
 	CreateQuotaTable(db)
+	CreateSubscriptionTable(db)
 	return db
 }
 
@@ -97,3 +98,20 @@ func CreateConversationTable(db *sql.DB) {
 		log.Fatal(err)
 	}
 }
+
+func CreateSubscriptionTable(db *sql.DB) {
+	_, err := db.Exec(`
+		CREATE TABLE IF NOT EXISTS subscription (
+		  id INT PRIMARY KEY AUTO_INCREMENT,
+		  user_id INT UNIQUE,
+		  expired_at DATETIME,
+		  created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+		  updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+		  total_month INT DEFAULT 0,
+		  FOREIGN KEY (user_id) REFERENCES auth(id)
+		);
+	`)
+	if err != nil {
+		log.Fatal(err)
+	}
+}
diff --git a/main.go b/main.go
index 4d388a14..9aabca81 100644
--- a/main.go
+++ b/main.go
@@ -31,6 +31,8 @@ func main() {
 		app.GET("/package", auth.PackageAPI)
 		app.GET("/quota", auth.QuotaAPI)
 		app.POST("/buy", auth.BuyAPI)
+		app.GET("/subscription", auth.SubscriptionAPI)
+		app.POST("/subscribe", auth.SubscribeAPI)
 		app.GET("/conversation/list", conversation.ListAPI)
 		app.GET("/conversation/load", conversation.LoadAPI)
 		app.GET("/conversation/delete", conversation.DeleteAPI)
diff --git a/middleware/throttle.go b/middleware/throttle.go
index 753739e9..a1439ef1 100644
--- a/middleware/throttle.go
+++ b/middleware/throttle.go
@@ -33,6 +33,8 @@ var limits = map[string]Limiter{
 	"/package":      {Duration: 1, Count: 2},
 	"/quota":        {Duration: 1, Count: 2},
 	"/buy":          {Duration: 1, Count: 2},
+	"/subscribe":    {Duration: 1, Count: 2},
+	"/subscription": {Duration: 1, Count: 2},
 	"/chat":         {Duration: 1, Count: 5},
 	"/conversation": {Duration: 1, Count: 5},
 }
diff --git a/utils/char.go b/utils/char.go
index f5ba3dc4..9d1d8b72 100644
--- a/utils/char.go
+++ b/utils/char.go
@@ -52,3 +52,7 @@ func ToInt(value string) int {
 		return 0
 	}
 }
+
+func ConvertSqlTime(t time.Time) string {
+	return t.Format("2006-01-02 15:04:05")
+}