main.go

package main

import (
	"context"
	"log"
	"net/http"

	"api_qps_go/config"
	"api_qps_go/models"
	"time"

	"github.com/gin-gonic/gin"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

func main() {
	// 初始化MongoDB连接
	ctx := context.Background()
	clientOptions := options.Client().
		ApplyURI(config.MongoURI).
		SetConnectTimeout(10 * time.Second)

	mongoDB, err := mongo.Connect(ctx, clientOptions)
	if err != nil {
		log.Fatal("MongoDB连接失败:", err)
	}
	defer mongoDB.Disconnect(ctx)

	// 测试MongoDB连接
	if err := mongoDB.Ping(ctx, nil); err != nil {
		log.Fatal("MongoDB ping失败:", err)
	}
	log.Println("MongoDB连接成功")

	// 创建IP限流器
	limiter := models.NewIPLimiter()

	// 初始化Gin路由
	r := gin.Default()

	// 添加IP限流中间件
	r.Use(func(c *gin.Context) {
		ip := c.ClientIP()
		if err := limiter.CheckIPLimit(c.Request.Context(), ip); err != nil {
			c.JSON(http.StatusTooManyRequests, gin.H{"error": err.Error()})
			c.Abort()
			return
		}
		c.Next()
	})

	// 根路由 - API说明
	r.GET("/", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"apis": []gin.H{
				{
					"path":        "/random",
					"method":      "GET",
					"description": "获取随机数据",
					"params": gin.H{
						"type": gin.H{
							"required":    false,
							"description": "数据类型",
							"options":     config.TypeMap,
						},
					},
					"about": "https://home.lideshan.top",
				},
			},
		})
	})

	// 定义API路由
	r.GET("/random", func(c *gin.Context) {
		typeParam := c.Query("type")
		data, err := models.GetRandomData(c.Request.Context(), mongoDB, typeParam)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		c.JSON(http.StatusOK, data)
	})

	// 启动服务器
	if err := r.Run(":8080"); err != nil {
		log.Fatal(err)
	}
}