package middleware

import (
	"context"
	"fmt"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	goredis "github.com/redis/go-redis/v9"

	"security-center-backend/internal/pkg/response"
)

// quotaKeyTTL 是配额计数 key 的过期时间, 略大于一分钟窗口, 防止残留。
const quotaKeyTTL = 65 * time.Second

// quotaStore 抽象 INCR/EXPIRE, 便于测试注入 fake(go.mod 无 miniredis/redismock)。
type quotaStore interface {
	Incr(ctx context.Context, key string) (int64, error)
	Expire(ctx context.Context, key string, ttl time.Duration) error
}

// redisQuotaStore 是 quotaStore 的 go-redis 实现。
type redisQuotaStore struct {
	rdb *goredis.Client
}

func (s redisQuotaStore) Incr(ctx context.Context, key string) (int64, error) {
	return s.rdb.Incr(ctx, key).Result()
}

func (s redisQuotaStore) Expire(ctx context.Context, key string, ttl time.Duration) error {
	return s.rdb.Expire(ctx, key, ttl).Err()
}

// Quota 按 api_key_id 计每分钟调用数, 超额返回 429(项目信封: HTTP 200 + body code 429)。
// 无 api_key_id(人类 JWT 调用)直接跳过。
func Quota(rdb *goredis.Client, enabled bool, defaultRPM int) gin.HandlerFunc {
	return quotaWithStore(redisQuotaStore{rdb: rdb}, enabled, defaultRPM)
}

func quotaWithStore(store quotaStore, enabled bool, defaultRPM int) gin.HandlerFunc {
	var warnOnce sync.Once
	return func(c *gin.Context) {
		if !enabled {
			c.Next()
			return
		}
		if defaultRPM <= 0 {
			// 防御: enabled=true 但 rpm 未配, 若计数必然全量 429 自我断网, 直通并告警一次。
			warnOnce.Do(func() {
				slog.Warn("quota enabled but default_rpm <= 0, pass through", "default_rpm", defaultRPM)
			})
			c.Next()
			return
		}
		keyID, ok := c.Get("api_key_id")
		if !ok {
			// 人类 JWT 调用, 不计 agent 配额。
			c.Next()
			return
		}
		window := now().Truncate(time.Minute).Unix()
		cacheKey := fmt.Sprintf("quota:agent:%v:%d", keyID, window)
		n, err := store.Incr(c.Request.Context(), cacheKey)
		if err != nil {
			// Redis 故障 fail-open, 不因配额组件拖垮业务链路。
			slog.Warn("quota incr failed, pass through", "key", cacheKey, "error", err)
			c.Next()
			return
		}
		if n == 1 {
			if err := store.Expire(c.Request.Context(), cacheKey, quotaKeyTTL); err != nil {
				slog.Warn("quota expire failed", "key", cacheKey, "error", err)
			}
		}
		if n > int64(defaultRPM) {
			c.Header("Retry-After", "60")
			response.Error(c, http.StatusTooManyRequests, "rate limit exceeded")
			c.Abort()
			return
		}
		c.Next()
	}
}