From e6899832f73cf94e55d2329fd6900c508ef1e593 Mon Sep 17 00:00:00 2001
From: Adrien <adrien.bouvais.pro@gmail.com>
Date: Sat, 1 Jun 2024 15:39:18 +0200
Subject: [PATCH] multiple plan

---
 Chat.go   | 12 ++++++++----
 Stripe.go | 40 +++++++++++++++++++++++++++++++++-------
 2 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/Chat.go b/Chat.go
index 272cb4d..94e4fc2 100644
--- a/Chat.go
+++ b/Chat.go
@@ -46,7 +46,8 @@ func LoadChatHandler(c *fiber.Ctx) error {
 	deleteLLMtoDelete(c)
 
 	if checkIfLogin(c) {
-		if IsCurrentUserLimiteReached(c) && !IsCurrentUserSubscribed(c) {
+		isPremium, isBasic := IsCurrentUserSubscribed(c)
+		if IsCurrentUserLimiteReached(c, isBasic) && !isPremium {
 			return c.SendString(generateLimitReachedChatHTML(c))
 		} else if !checkIfHaveKey(c) {
 			return c.SendString(generateEnterKeyChatHTML())
@@ -66,7 +67,8 @@ func LoadChatInputHandler(c *fiber.Ctx) error {
 		}
 		return c.SendString(out)
 	} else {
-		out, err := chatInputTmpl.Execute(pongo2.Context{"IsLogin": checkIfLogin(c), "HaveKey": checkIfHaveKey(c), "IsSubscribed": IsCurrentUserSubscribed(c), "IsLimiteReached": IsCurrentUserLimiteReached(c)})
+		isPremium, isBasic := IsCurrentUserSubscribed(c)
+		out, err := chatInputTmpl.Execute(pongo2.Context{"IsLogin": checkIfLogin(c), "HaveKey": checkIfHaveKey(c), "IsSubscribed": isPremium, "IsLimiteReached": IsCurrentUserLimiteReached(c, isBasic)})
 		if err != nil {
 			fmt.Println("Error executing chat input template")
 			panic(err)
@@ -697,6 +699,7 @@ func GenerateModelPopoverHTML(refresh bool, c *fiber.Ctx) string {
 	}
 
 	modelInfos := GetAvailableModels(c)
+	isPremium, isBasic := IsCurrentUserSubscribed(c)
 
 	out, err := modelPopoverTmpl.Execute(pongo2.Context{
 		"IsLogin":         checkIfLogin(c),
@@ -710,7 +713,7 @@ func GenerateModelPopoverHTML(refresh bool, c *fiber.Ctx) string {
 		"LLMs":            llms,
 		"ModelInfos":      modelInfos,
 		"DeleteUpdate":    refresh,
-		"IsSub":           IsCurrentUserSubscribed(c),
+		"IsSub":           isPremium || isBasic,
 	})
 	if err != nil {
 		fmt.Println("Error generating model popover")
@@ -798,6 +801,7 @@ func LoadSettingsHandler(c *fiber.Ctx) error {
 	stripeSubLink := "https://billing.stripe.com/p/login/test_eVa5kC1q7dogaaIcMM?prefilled_email=" + user.Email
 
 	openaiExists, anthropicExists, mistralExists, groqExists, gooseaiExists, googleExists := getExistingKeys(c)
+	isPremium, isBasic := IsCurrentUserSubscribed(c)
 
 	out, err := settingPopoverTmpl.Execute(pongo2.Context{
 		"IsLogin":         checkIfLogin(c),
@@ -808,7 +812,7 @@ func LoadSettingsHandler(c *fiber.Ctx) error {
 		"GooseaiExists":   gooseaiExists,
 		"GoogleExists":    googleExists,
 		"AnyExists":       openaiExists || anthropicExists || mistralExists || groqExists || gooseaiExists || googleExists,
-		"IsSub":           IsCurrentUserSubscribed(c),
+		"IsSub":           isPremium || isBasic,
 		"StripeSubLink":   stripeSubLink,
 	})
 	if err != nil {
diff --git a/Stripe.go b/Stripe.go
index 0832ebd..1b4718d 100644
--- a/Stripe.go
+++ b/Stripe.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"time"
 
 	"github.com/flosch/pongo2"
 	"github.com/gofiber/fiber/v2"
@@ -83,7 +84,8 @@ func CreateNewStripeCustomer(name string, email string) string {
 	return result.ID
 }
 
-func IsCurrentUserSubscribed(c *fiber.Ctx) bool {
+func IsCurrentUserSubscribed(c *fiber.Ctx) (bool, bool) {
+	// TODO Change to add both plan
 	var user User
 	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, "SELECT global currentUser { stripe_id } LIMIT 1;", &user)
 	if err != nil {
@@ -96,28 +98,52 @@ func IsCurrentUserSubscribed(c *fiber.Ctx) bool {
 		panic(err)
 	}
 	if result.Subscriptions.TotalCount == 0 {
-		return false
+		return false, false
 	}
 
-	return result.Subscriptions.Data[0].Plan.Product.ID == "prod_PnDjBBwvQN36cQ" && result.Subscriptions.Data[0].Plan.Active
+	// Check if an active product "prod_PnDjBBwvQN36cQ" is in the subscription. It is the premium plan
+	isPremium := false
+	for _, sub := range result.Subscriptions.Data {
+		if sub.Plan.Product.ID == "prod_Q9TpVos0t8QtX3" && sub.Plan.Active {
+			isPremium = true
+			break
+		}
+	}
+
+	isBasic := false
+	for _, sub := range result.Subscriptions.Data {
+		if sub.Plan.Product.ID == "prod_QDHakyHUeHLMET" && sub.Plan.Active {
+			isBasic = true
+			break
+		}
+	}
+
+	return isPremium, isBasic
 }
 
-func IsCurrentUserLimiteReached(c *fiber.Ctx) bool {
+func IsCurrentUserLimiteReached(c *fiber.Ctx, isBasic bool) bool {
+	// 30 days ago
+	date := time.Now().AddDate(0, 0, -30)
 	// Count the number of Usage for the current user
 	var count int64
 	err := edgeGlobalClient.WithGlobals(map[string]interface{}{"ext::auth::client_token": c.Cookies("jade-edgedb-auth-token")}).QuerySingle(edgeCtx, `
 	WITH
 		U := (
 			SELECT Usage
-			FILTER .user = global currentUser
+			FILTER .user = global currentUser AND .date >= <datetime>$0
 		)
 	SELECT count(U)
-	`, &count)
+	`, &count, date)
 	if err != nil {
 		fmt.Println("Error counting Usage")
 		panic(err)
 	}
-	return count >= 500
+	// 500 if isBasic is false, otherwise 10000
+	if isBasic {
+		return count >= 10000
+	} else {
+		return count >= 500
+	}
 }
 
 func CreateClientSecretSession(c *fiber.Ctx) string {