#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_refill_ts = -period; //-period ensures the first tick can start 
                          //at >= 0 and still refill the bicket

    assert(tickets > 0);
    assert(r->period > 0);
    assert(r->period !=  LONG_MAX);
}

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

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

static void uc_ratelimit_refill(struct RateLimit *r, long current_ts)
{

    long diff_period;
    //Find elapsed time
    diff_period = current_ts - r->last_refill_ts;

    if (diff_period > r->period) {
        //help prevent overflow when calculating available tickets below
        diff_period = r->period;
    } 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;
    }

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

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

    if (r->funds < r->period) { //not enough funds

        uc_ratelimit_refill(r, current_ts);

        //as we keep track of the last timestamp of refilling
        //We don't need to refill when there is enough funds
    }
  
    //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;
    }



    return allowed;
}