#include <assert.h>
#include <limits.h>
#include "ucore/rate_limit.h"

void uc_ratelimit_init(struct RateLimit *r, long tickets, long period)
{
    r->tickets = tickets;

    r->period = period;
    r->funds = 0; //we fill it on first _allow call
    r->last_ts = -period; //-period ensures the first tick can start 
                          //at >= 0 and still refill the bicket

    assert(tickets > 0);
    assert(r->period > 0);
}

void uc_ratelimit_reset(struct RateLimit *r, long current_ts)
{
    r->funds = r->period * r->tickets;

    if (current_ts) {
        r->last_ts = current_ts;
    }
}

int uc_ratelimit_allow(struct RateLimit *r, long current_ts)
{
    int allowed;
    long diff_period;

    if (r->funds >= r->period) {
        //we have enough
        r->funds -= r->period;
        allowed = 1;
    } else { //refill
        //Find elapsed time
        diff_period = current_ts - r->last_ts;

        if (diff_period > r->period + 1) {
            //help prevent overflow when calculating available_tickets below
            diff_period = r->period + 1;
        } else if (diff_period < 0) {
            //time went backwards, skip this round
            diff_period = 0;  
        }

        //Calculate the cost of tickets that became available 
        r->funds += diff_period * r->tickets;


        //throttle handing out tickets
        if (r->funds > r->period * r->tickets) {
            r->funds = r->period * r->tickets;
        }

        //If we have enough to buy atleast one ticket, we can allow 
        if (r->funds >=  r->period) {
            r->funds -= r->period;
            allowed = 1;
        } else {
            //not enough to buy a ticket
            allowed = 0;
        }

        //save this period. 
        r->last_ts = current_ts;
    }


    return allowed;
}