package ratelimit

import (
	"context"
	"fmt"
	"time"
)

// RedisClient is the subset of Redis operations needed for rate limiting.
type RedisClient interface {
	Incr(ctx context.Context, key string) (int64, error)
	Expire(ctx context.Context, key string, expiration time.Duration) error
	Decr(ctx context.Context, key string) (int64, error)
}

// CheckMessageRate checks if a user has exceeded the message rate limit.
// Returns true if allowed, false if rate limited.
func CheckMessageRate(ctx context.Context, redis RedisClient, username string, max int, window time.Duration) (bool, error) {
	if redis == nil {
		return true, nil // No Redis = no rate limiting
	}
	key := fmt.Sprintf("rate:%s", username)
	count, err := redis.Incr(ctx, key)
	if err != nil {
		return true, err // Fail open on Redis error
	}
	if count == 1 {
		_ = redis.Expire(ctx, key, window)
	}
	return count <= int64(max), nil
}

// CheckLoginRate checks if an IP has exceeded the login attempt limit.
// Returns true if allowed, false if rate limited.
func CheckLoginRate(ctx context.Context, redis RedisClient, clientIP string) (bool, error) {
	if redis == nil {
		return true, nil
	}
	key := fmt.Sprintf("login_attempts:%s", clientIP)
	count, err := redis.Incr(ctx, key)
	if err != nil {
		return true, err
	}
	if count == 1 {
		_ = redis.Expire(ctx, key, 5*time.Minute)
	}
	return count <= 5, nil
}

// CheckWSConnRate checks if an IP has exceeded the WebSocket connection limit.
// Returns true if allowed, false if rate limited.
// Caller must call Decr on disconnect.
func CheckWSConnRate(ctx context.Context, redis RedisClient, clientIP string) (bool, error) {
	if redis == nil {
		return true, nil
	}
	key := fmt.Sprintf("ws_conn:%s", clientIP)
	count, err := redis.Incr(ctx, key)
	if err != nil {
		return true, err
	}
	if count == 1 {
		_ = redis.Expire(ctx, key, 60*time.Second)
	}
	return count <= 5, nil
}

// DecrWSConn decrements the WebSocket connection counter for an IP.
func DecrWSConn(ctx context.Context, redis RedisClient, clientIP string) {
	if redis == nil {
		return
	}
	key := fmt.Sprintf("ws_conn:%s", clientIP)
	_, _ = redis.Decr(ctx, key)
}