fanqie-go/internal/handlers/common.go

package handlers

import (
	"encoding/json"
	"fanqie/internal/config"
	"fanqie/internal/fqapi"
	"fmt"
	"net/http"
	"sync/atomic"
	"time"
)

var (
	startTime    = time.Now()
	deviceHeader string

	// 请求统计
	requestStats = newRequestStats()

	// 统计保存定时器
	statsSaveTicker *time.Ticker
)

// statsKinds 统计类别
var statsKinds = []string{"total", "search", "detail", "catalog", "content", "manga", "audio", "video", "category"}

// newRequestStats 初始化各类计数器
func newRequestStats() map[string]*atomic.Int64 {
	m := make(map[string]*atomic.Int64, len(statsKinds))
	for _, k := range statsKinds {
		m[k] = &atomic.Int64{}
	}
	return m
}

// incrementStat 原子递增请求计数
func incrementStat(kinds ...string) {
	for _, k := range kinds {
		requestStats[k].Add(1)
	}
}

// getStatsSnapshot 获取所有计数器的快照
func getStatsSnapshot() map[string]int64 {
	snapshot := make(map[string]int64, len(requestStats))
	for k, c := range requestStats {
		snapshot[k] = c.Load()
	}
	return snapshot
}

// SetDeviceHeader 设置设备header
func SetDeviceHeader(header string) {
	deviceHeader = header
}

// LoadStatistics 从数据库加载统计数据
func LoadStatistics() {
	stats, err := fqapi.LoadStatistics()
	if err != nil {
		fmt.Printf("加载统计数据失败: %v\n", err)
		return
	}
	
	for k, v := range stats {
		if c, ok := requestStats[k]; ok {
			c.Store(v)
		}
	}
}

// SaveStatistics 保存统计数据到数据库
func SaveStatistics() {
	if err := fqapi.SaveStatistics(getStatsSnapshot()); err != nil {
		fmt.Printf("保存统计数据失败: %v\n", err)
	}
}

// StartStatisticsSaver 启动统计数据定期保存
func StartStatisticsSaver() {
	// 每30秒保存一次统计数据
	statsSaveTicker = time.NewTicker(30 * time.Second)
	go func() {
		for range statsSaveTicker.C {
			SaveStatistics()
		}
	}()
}

// StopStatisticsSaver 停止统计数据定期保存
func StopStatisticsSaver() {
	if statsSaveTicker != nil {
		statsSaveTicker.Stop()
		// 最后保存一次
		SaveStatistics()
	}
}

// RequireAccessKey 验证访问密钥中间件
func RequireAccessKey(handler http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		cfg := config.GetConfig()
		
		// 如果未配置 access_key 或为空，则不启用验证
		if cfg == nil || cfg.Server.AccessKey == "" {
			handler(w, r)
			return
		}
		
		// 从查询参数或 Header 中获取 key
		key := r.URL.Query().Get("key")
		if key == "" {
			key = r.Header.Get("X-Access-Key")
		}
		
		// 验证密钥
		if key != cfg.Server.AccessKey {
			w.Header().Set("Content-Type", "application/json; charset=utf-8")
			w.WriteHeader(http.StatusUnauthorized)
			fmt.Fprintf(w, `{"code": -1, "msg": "访问被拒绝：无效的访问密钥"}`)
			return
		}
		
		// 验证通过，继续处理请求
		handler(w, r)
	}
}

// HandleMain 主页处理
func HandleMain(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	response := map[string]interface{}{
		"code":    0,
		"message": "success",
		"data":    nil,
	}
	json.NewEncoder(w).Encode(response)
}

// HandleSelfInfo 设备信息处理
func HandleSelfInfo(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	
	// 获取设备池统计信息
	pool := fqapi.GetDeviceKeyPool()
	poolStats := pool.GetStats()
	
	response := map[string]interface{}{
		"code":    0,
		"message": "SUCCESS",
		"data": map[string]interface{}{
			"device_pool_enabled": pool.IsEnabled(),
			"device_pool_stats":   poolStats,
		},
	}
	json.NewEncoder(w).Encode(response)
}