import Foundation

/// A token bucket matching the relay's per-connection budgets so frames are
/// paced instead of dropped as `rate-limited`.
public struct TokenBucket: Sendable {
    public let capacity: Double
    public let refillPerSecond: Double
    private var tokens: Double
    private var last: Duration

    public init(capacity: Double, refillPerSecond: Double, now: Duration) {
        self.capacity = capacity
        self.refillPerSecond = refillPerSecond
        tokens = capacity
        last = now
    }

    private mutating func refill(now: Duration) {
        let elapsed = max(0, (now - last).totalSeconds)
        tokens = min(capacity, tokens + elapsed * refillPerSecond)
        last = now
    }

    /// Takes `cost` tokens, returning how long the caller must wait first (zero when
    /// the budget allows it now). The bucket goes into debt, so concurrent callers
    /// queue behind each other: each waits for its own share of the refill.
    public mutating func take(_ cost: Double, now: Duration) -> Duration {
        refill(now: now)
        tokens -= cost
        if tokens >= 0 { return .zero }
        return .milliseconds(Int((-tokens / refillPerSecond * 1000).rounded(.up)))
    }

    public func available(now: Duration) -> Double {
        var copy = self
        copy.refill(now: now)
        return max(0, copy.tokens)
    }
}

/// The two relay budgets: messages per second (with burst) and bytes per second (4× burst).
public struct RelayPacer: Sendable {
    public var messages: TokenBucket
    public var bytes: TokenBucket
    public var enabled = true

    public init(limits: AppDescriptor.Limits, now: Duration) {
        messages = TokenBucket(capacity: Double(limits.messageBurst), refillPerSecond: Double(limits.messagesPerSec), now: now)
        bytes = TokenBucket(capacity: Double(max(limits.maxBytesPerSec * 4, limits.maxFrameBytes)), refillPerSecond: Double(limits.maxBytesPerSec), now: now)
    }

    /// The wait before a frame of `byteCount` bytes may be sent.
    public mutating func take(byteCount: Int, now: Duration) -> Duration {
        guard enabled else { return .zero }
        let a = messages.take(1, now: now)
        let b = bytes.take(Double(byteCount), now: now)
        return max(a, b)
    }
}