From eb22de912a54873b3f80ed2225e596d87851ce48 Mon Sep 17 00:00:00 2001
From: Sakurasan <1173092237@qq.com>
Date: Mon, 17 Apr 2023 22:48:37 +0800
Subject: [PATCH] add usage

---
 opencat.go       |  2 ++
 router/router.go | 39 +++++++++++++++++++++++++++++++++++++++
 store/db.go      | 11 +++++++++++
 store/usage.go   | 38 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 90 insertions(+)
 create mode 100644 store/usage.go

diff --git a/opencat.go b/opencat.go
index 8dbba0b..1859eb9 100644
--- a/opencat.go
+++ b/opencat.go
@@ -38,6 +38,8 @@ func main() {
 		// 获取所有用户信息
 		group.GET("/users", router.HandleUsers)
 
+		group.GET("/usages", router.HandleUsage)
+
 		// 添加Key
 		group.POST("/keys", router.HandleAddKey)
 
diff --git a/router/router.go b/router/router.go
index a18ccd9..6510268 100644
--- a/router/router.go
+++ b/router/router.go
@@ -381,3 +381,42 @@ func HandleReverseProxy(c *gin.Context) {
 	proxy.ServeHTTP(c.Writer, req)
 
 }
+
+type Usage struct {
+	Cost      string `json:"cost"`
+	UserID    int    `json:"userId"`
+	TotalUnit int    `json:"totalUnit"`
+}
+
+func HandleUsage(c *gin.Context) {
+	fromStr := c.Query("from")
+	toStr := c.Query("to")
+
+	from, err := time.Parse("2006-01-02", fromStr)
+	if err != nil {
+		c.JSON(400, gin.H{"error": "Invalid from date format"})
+		return
+	}
+
+	to, err := time.Parse("2006-01-02", toStr)
+	if err != nil {
+		c.JSON(400, gin.H{"error": "Invalid to date format"})
+		return
+	}
+
+	err = store.QueryUsage(from, to)
+	if err != nil {
+		c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
+		return
+	}
+
+	// Mock data for testing
+	usage := Usage{
+		Cost:      "0.000076",
+		UserID:    1,
+		TotalUnit: 38,
+	}
+
+	c.JSON(200, []Usage{usage})
+
+}
diff --git a/store/db.go b/store/db.go
index 88b1294..65027bb 100644
--- a/store/db.go
+++ b/store/db.go
@@ -11,6 +11,8 @@ import (
 
 var db *gorm.DB
 
+var usage *gorm.DB
+
 func init() {
 	if _, err := os.Stat("db"); os.IsNotExist(err) {
 		errDir := os.MkdirAll("db", 0755)
@@ -31,4 +33,13 @@ func init() {
 	}
 	LoadKeysCache()
 	LoadAuthCache()
+
+	usage, err = gorm.Open(sqlite.Open("./db/usage.db"), &gorm.Config{})
+	if err != nil {
+		panic(err)
+	}
+	err = usage.AutoMigrate(&DailyUsage{}, &Usage{})
+	if err != nil {
+		panic(err)
+	}
 }
diff --git a/store/usage.go b/store/usage.go
new file mode 100644
index 0000000..23a755e
--- /dev/null
+++ b/store/usage.go
@@ -0,0 +1,38 @@
+package store
+
+import "time"
+
+type DailyUsage struct {
+	ID              int       `gorm:"column:id"`
+	UserID          int       `gorm:"column:user_id primarykey"`
+	Date            time.Time `gorm:"column:date"`
+	SKU             string    `gorm:"column:sku"`
+	PromptUnits     int       `gorm:"column:prompt_units"`
+	CompletionUnits int       `gorm:"column:completion_units"`
+	TotalUnit       int       `gorm:"column:total_unit"`
+	Cost            string    `gorm:"column:cost"`
+}
+
+func (DailyUsage) TableName() string {
+	return "daily_usages"
+}
+
+type Usage struct {
+	ID              int       `gorm:"column:id"`
+	PromptHash      string    `gorm:"column:prompt_hash"`
+	UserID          int       `gorm:"column:user_id"`
+	Date            time.Time `gorm:"column:date"`
+	SKU             string    `gorm:"column:sku"`
+	PromptUnits     int       `gorm:"column:prompt_units"`
+	CompletionUnits int       `gorm:"column:completion_units"`
+	TotalUnit       int       `gorm:"column:total_unit"`
+	Cost            string    `gorm:"column:cost"`
+}
+
+func (Usage) TableName() string {
+	return "usages"
+}
+
+func QueryUsage(from, to time.Time) error {
+	return nil
+}