package middleware

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"time"

	"git.echol.cn/loser/ai_proxy/server/global"
	"git.echol.cn/loser/ai_proxy/server/model/common/response"
	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
)

type LimitConfig struct {
	GenerationDuration time.Duration
	Limit              int
}

func (l LimitConfig) LimitKey(c *gin.Context) string {
	return "GVA_Limit" + c.ClientIP()
}

func (l LimitConfig) GetLimit(c *gin.Context) int {
	return l.Limit
}

func (l LimitConfig) Reached(c *gin.Context) response.Response {
	return response.Response{Code: response.ERROR, Data: nil, Msg: "操作过于频繁，请稍后再试"}
}

// IPLimit IP 限流中间件
func IPLimit() gin.HandlerFunc {
	return func(c *gin.Context) {
		if global.GVA_CONFIG.System.UseRedis {
			key := "GVA_Limit" + c.ClientIP()
			limit := global.GVA_CONFIG.System.IpLimitCount
			limitTime := global.GVA_CONFIG.System.IpLimitTime

			ctx := context.Background()
			count, err := global.GVA_REDIS.Get(ctx, key).Int()
			if err != nil && !errors.Is(err, context.DeadlineExceeded) {
				global.GVA_LOG.Error("get redis key error", zap.Error(err))
			}

			if count >= limit {
				c.JSON(http.StatusOK, gin.H{
					"code": response.ERROR,
					"msg":  fmt.Sprintf("操作过于频繁，请在 %d 秒后再试", limitTime),
				})
				c.Abort()
				return
			}

			pipe := global.GVA_REDIS.Pipeline()
			pipe.Incr(ctx, key)
			pipe.Expire(ctx, key, time.Second*time.Duration(limitTime))
			_, err = pipe.Exec(ctx)
			if err != nil {
				global.GVA_LOG.Error("redis pipeline error", zap.Error(err))
			}
		}
		c.Next()
	}
}