io.github.freya022.botcommands.internal.commands.ratelimit.RateLimitAnnotations.kt Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of BotCommands Show documentation

A Kotlin-first (and Java) framework that makes creating Discord bots a piece of cake, using the JDA library.

There is a newer version: 3.0.0-alpha.18

package io.github.freya022.botcommands.internal.commands.ratelimit

import io.github.freya022.botcommands.api.commands.annotations.*
import io.github.freya022.botcommands.api.commands.ratelimit.RateLimiter
import io.github.freya022.botcommands.api.commands.ratelimit.RateLimiterFactory
import io.github.freya022.botcommands.api.commands.ratelimit.bucket.BucketFactory
import io.github.freya022.botcommands.internal.utils.ReflectionUtils.declaringClass
import io.github.freya022.botcommands.internal.utils.annotationRef
import io.github.freya022.botcommands.internal.utils.requireUser
import java.time.Duration
import kotlin.reflect.KFunction
import kotlin.reflect.full.findAnnotation
import io.github.bucket4j.Bandwidth as BucketBandwidth
import io.github.bucket4j.Refill as BucketRefill

fun Refill.toRealRefill(): BucketRefill {
    val duration = Duration.of(period, periodUnit)
    return when (type) {
        RefillType.GREEDY -> BucketRefill.greedy(tokens, duration)
        RefillType.INTERVAL -> BucketRefill.intervally(tokens, duration)
    }
}

fun Bandwidth.toRealBandwidth(): BucketBandwidth {
    return BucketBandwidth.classic(capacity, refill.toRealRefill())
}

fun KFunction<*>.readRateLimit(): Pair? {
    val rateLimitAnnotation = findAnnotation() ?: this.declaringClass.findAnnotation()
    val cooldownAnnotation = findAnnotation() ?: this.declaringClass.findAnnotation()
    requireUser(cooldownAnnotation == null || rateLimitAnnotation == null, this) {
        "Cannot use both ${annotationRef()} and ${annotationRef()}"
    }

    return if (rateLimitAnnotation != null) {
        BucketFactory.custom(rateLimitAnnotation.bandwidths.map { it.toRealBandwidth() }) to RateLimiter.defaultFactory(rateLimitAnnotation.scope, rateLimitAnnotation.deleteOnRefill)
    } else if (cooldownAnnotation != null) {
        BucketFactory.ofCooldown(Duration.of(cooldownAnnotation.cooldown, cooldownAnnotation.unit)) to RateLimiter.defaultFactory(cooldownAnnotation.rateLimitScope, cooldownAnnotation.deleteOnRefill)
    } else {
        null
    }
}