package app

import (
	"context"
	"crypto/sha256"
	"fmt"
	"net"
	"net/http"
	"strings"
	"time"
)

func (a *App) allowRequest(ctx context.Context, action, subject string, limit int, window time.Duration) bool {
	if limit < 1 || window < time.Second {
		return false
	}
	windowSeconds := int64(window / time.Second)
	slot := time.Now().Unix() / windowSeconds
	key := sha256.Sum256([]byte(fmt.Sprintf("%s|%s|%d", action, subject, slot)))
	expiresAt := time.Unix((slot+1)*windowSeconds, 0).Add(time.Minute)
	_, err := a.db.ExecContext(ctx, `INSERT INTO api_rate_limits(bucket_key,action_name,hits,expires_at) VALUES(?,?,1,?) ON DUPLICATE KEY UPDATE hits=hits+1,expires_at=VALUES(expires_at)`, key[:], action, expiresAt)
	if err != nil {
		return false
	}
	var hits int
	if err = a.db.QueryRowContext(ctx, `SELECT hits FROM api_rate_limits WHERE bucket_key=?`, key[:]).Scan(&hits); err != nil {
		return false
	}
	if key[0] == 0 {
		_, _ = a.db.ExecContext(ctx, `DELETE FROM api_rate_limits WHERE expires_at<NOW(3) LIMIT 1000`)
	}
	return hits <= limit
}

func (a *App) rateLimit(w http.ResponseWriter, r *http.Request, action, subject string, limit int, window time.Duration) bool {
	if a.allowRequest(r.Context(), action, subject, limit, window) {
		return true
	}
	w.Header().Set("Retry-After", fmt.Sprintf("%d", int(window/time.Second)))
	fail(w, http.StatusTooManyRequests, 20002, "请求过于频繁，请稍后再试")
	return false
}

func clientIP(r *http.Request) string {
	if forwarded := strings.TrimSpace(strings.Split(r.Header.Get("X-Forwarded-For"), ",")[0]); net.ParseIP(forwarded) != nil {
		return forwarded
	}
	if realIP := strings.TrimSpace(r.Header.Get("X-Real-IP")); net.ParseIP(realIP) != nil {
		return realIP
	}
	host, _, err := net.SplitHostPort(strings.TrimSpace(r.RemoteAddr))
	if err == nil && net.ParseIP(host) != nil {
		return host
	}
	return "unknown"
}