import java.time.Duration;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.LongSupplier;

public class TimeoutBudgetDemo {
    static final class Budget {
        private final LongSupplier clock;
        private final long started;
        private final long totalNanos;
        Budget(Duration total, LongSupplier clock) {
            if (total.isNegative() || total.isZero()) throw new IllegalArgumentException("positive budget required");
            this.clock = clock;
            this.started = clock.getAsLong();
            this.totalNanos = total.toNanos();
        }
        long remainingNanos() {
            return Math.max(0, totalNanos - (clock.getAsLong() - started));
        }
        Duration stage(Duration cap, Duration reserve) throws TimeoutException {
            long available = remainingNanos() - reserve.toNanos();
            if (available <= 0) throw new TimeoutException("deadline/reserve reached");
            if (cap.isNegative() || cap.isZero() || reserve.isNegative()) throw new IllegalArgumentException();
            return Duration.ofNanos(Math.min(available, cap.toNanos()));
        }
        boolean retry(boolean transientFailure, boolean idempotent, boolean emitted,
                      Duration backoff, Duration minAttempt, Duration reserve) {
            if (backoff.isNegative() || minAttempt.isNegative() || reserve.isNegative()) throw new IllegalArgumentException();
            long needed = Math.addExact(Math.addExact(backoff.toNanos(), minAttempt.toNanos()), reserve.toNanos());
            return transientFailure && idempotent && !emitted && remainingNanos() > needed;
        }
    }
    public static void main(String[] args) throws Exception {
        AtomicLong fakeNanos = new AtomicLong(); // deterministic logical time, not measured latency
        Budget budget = new Budget(Duration.ofSeconds(10), fakeNanos::get);
        fakeNanos.addAndGet(TimeUnit.MILLISECONDS.toNanos(500));
        System.out.println("modelCapMs=" + budget.stage(Duration.ofSeconds(6), Duration.ofMillis(500)).toMillis());
        fakeNanos.addAndGet(TimeUnit.SECONDS.toNanos(6));
        System.out.println("remainingMs=" + TimeUnit.NANOSECONDS.toMillis(budget.remainingNanos()));
        System.out.println("retryWithBackoff=" + budget.retry(true, true, false,
            Duration.ofSeconds(1), Duration.ofSeconds(3), Duration.ofMillis(500)));
        System.out.println("shortRetry=" + budget.retry(true, true, false,
            Duration.ofMillis(100), Duration.ofSeconds(1), Duration.ofMillis(500)));
        System.out.println("writeRetry=" + budget.retry(true, false, false,
            Duration.ZERO, Duration.ofSeconds(1), Duration.ZERO));
        System.out.println("afterSseRetry=" + budget.retry(true, true, true,
            Duration.ZERO, Duration.ofSeconds(1), Duration.ZERO));
        fakeNanos.set(TimeUnit.SECONDS.toNanos(10));
        try {
            budget.stage(Duration.ofSeconds(1), Duration.ZERO);
            throw new AssertionError("expired budget accepted");
        } catch (TimeoutException expected) {
            System.out.println("atDeadline=TimeoutException");
        }
        // Production construction: new Budget(Duration.ofSeconds(10), System::nanoTime).
    }
}