From fe18d6c82857327a2c85281a0b38518695c8ba09 Mon Sep 17 00:00:00 2001
From: wood chen <wood@czl.net>
Date: Sat, 26 Oct 2024 12:39:00 +0800
Subject: [PATCH] feat(api, stats): add stats manager and endpoint call
 tracking

---
 main.go        | 22 ++++++++++++
 stats/stats.go | 97 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 119 insertions(+)
 create mode 100644 stats/stats.go

diff --git a/main.go b/main.go
index 47ac09f..6ba70c0 100644
--- a/main.go
+++ b/main.go
@@ -11,6 +11,7 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"random-api-go/stats"
 	"strings"
 	"sync"
 	"time"
@@ -31,6 +32,8 @@ var (
 	rng           *rand.Rand
 )
 
+var statsManager *stats.StatsManager
+
 type URLSelector struct {
 	URLs         []string
 	CurrentIndex int
@@ -80,11 +83,19 @@ func (us *URLSelector) GetRandomURL() string {
 	return us.URLs[0]
 }
 
+func init() {
+	// 确保数据目录存在
+	if err := os.MkdirAll("data", 0755); err != nil {
+		log.Fatal("Failed to create data directory:", err)
+	}
+}
+
 func main() {
 	source := rand.NewSource(time.Now().UnixNano())
 	rng = rand.New(source)
 
 	setupLogging()
+	statsManager = stats.NewStatsManager("data/stats.json")
 
 	if err := loadCSVPaths(); err != nil {
 		log.Fatal("Failed to load CSV paths:", err)
@@ -97,6 +108,8 @@ func main() {
 	// 设置 API 路由
 	http.HandleFunc("/pic/", handleAPIRequest)
 	http.HandleFunc("/video/", handleAPIRequest)
+	// 添加统计API路由
+	http.HandleFunc("/stats", handleStats)
 
 	log.Printf("Listening on %s...\n", port)
 	if err := http.ListenAndServe(port, nil); err != nil {
@@ -252,9 +265,18 @@ func handleAPIRequest(w http.ResponseWriter, r *http.Request) {
 
 	randomURL := selector.GetRandomURL()
 
+	// 记录统计
+	endpoint := fmt.Sprintf("%s/%s", prefix, suffix)
+	statsManager.IncrementCalls(endpoint)
+
 	duration := time.Since(start)
 	log.Printf("Request: %s %s from %s - Source: %s - Duration: %v - Redirecting to: %s",
 		r.Method, r.URL.Path, realIP, sourceDomain, duration, randomURL)
 
 	http.Redirect(w, r, randomURL, http.StatusFound)
 }
+
+func handleStats(w http.ResponseWriter, r *http.Request) {
+	stats := statsManager.GetStats()
+	json.NewEncoder(w).Encode(stats)
+}
diff --git a/stats/stats.go b/stats/stats.go
new file mode 100644
index 0000000..78e59ea
--- /dev/null
+++ b/stats/stats.go
@@ -0,0 +1,97 @@
+package stats
+
+import (
+	"encoding/json"
+	"os"
+	"sync"
+	"time"
+)
+
+type EndpointStats struct {
+	TotalCalls    int64  `json:"total_calls"`
+	TodayCalls    int64  `json:"today_calls"`
+	LastResetDate string `json:"last_reset_date"`
+}
+
+type StatsManager struct {
+	Stats    map[string]*EndpointStats `json:"stats"`
+	mu       sync.RWMutex
+	filepath string
+}
+
+func NewStatsManager(filepath string) *StatsManager {
+	sm := &StatsManager{
+		Stats:    make(map[string]*EndpointStats),
+		filepath: filepath,
+	}
+	sm.LoadStats()
+	go sm.startDailyReset()
+	return sm
+}
+
+func (sm *StatsManager) IncrementCalls(endpoint string) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	if _, exists := sm.Stats[endpoint]; !exists {
+		sm.Stats[endpoint] = &EndpointStats{
+			LastResetDate: time.Now().Format("2006-01-02"),
+		}
+	}
+
+	sm.Stats[endpoint].TotalCalls++
+	sm.Stats[endpoint].TodayCalls++
+
+	// 异步保存统计数据
+	go sm.SaveStats()
+}
+
+func (sm *StatsManager) GetStats() map[string]*EndpointStats {
+	sm.mu.RLock()
+	defer sm.mu.RUnlock()
+
+	return sm.Stats
+}
+
+func (sm *StatsManager) SaveStats() error {
+	sm.mu.RLock()
+	defer sm.mu.RUnlock()
+
+	data, err := json.MarshalIndent(sm, "", "  ")
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(sm.filepath, data, 0644)
+}
+
+func (sm *StatsManager) LoadStats() error {
+	data, err := os.ReadFile(sm.filepath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return err
+	}
+
+	return json.Unmarshal(data, sm)
+}
+
+func (sm *StatsManager) startDailyReset() {
+	for {
+		now := time.Now()
+		next := now.Add(24 * time.Hour)
+		next = time.Date(next.Year(), next.Month(), next.Day(), 0, 0, 0, 0, next.Location())
+		duration := next.Sub(now)
+
+		time.Sleep(duration)
+
+		sm.mu.Lock()
+		for _, stats := range sm.Stats {
+			stats.TodayCalls = 0
+			stats.LastResetDate = time.Now().Format("2006-01-02")
+		}
+		sm.mu.Unlock()
+
+		sm.SaveStats()
+	}
+}