#include <linux/sched.h>
#include <linux/slab.h>

#include <litmus/litmus.h>
#include <litmus/reservation.h>

void reservation_init(struct reservation *res)
{
	memset(res, sizeof(*res), 0);
	res->state = RESERVATION_INACTIVE;
	INIT_LIST_HEAD(&res->clients);
}

struct task_struct* default_dispatch_client(
	struct reservation *res,
	lt_t *for_at_most)
{
	struct reservation_client *client, *next;
	struct task_struct* tsk;

	BUG_ON(res->state != RESERVATION_ACTIVE);
	*for_at_most = 0;

	list_for_each_entry_safe(client, next, &res->clients, list) {
		tsk = client->dispatch(client);
		if (likely(tsk)) {
			return tsk;
		}
	}
	return NULL;
}

static struct task_struct * task_client_dispatch(struct reservation_client *client)
{
	struct task_client *tc = container_of(client, struct task_client, client);
	return tc->task;
}

void task_client_init(struct task_client *tc, struct task_struct *tsk,
	struct reservation *res)
{
	memset(&tc->client, sizeof(tc->client), 0);
	tc->client.dispatch = task_client_dispatch;
	tc->client.reservation = res;
	tc->task = tsk;
}

static void sup_scheduler_update_at(
	struct sup_reservation_environment* sup_env,
	lt_t when)
{
	//TRACE("SCHEDULER_UPDATE_AT update: %llu > when %llu\n", sup_env->next_scheduler_update, when);
	if (sup_env->next_scheduler_update > when)
		sup_env->next_scheduler_update = when;
}

void sup_scheduler_update_after(
	struct sup_reservation_environment* sup_env,
	lt_t timeout)
{
	sup_scheduler_update_at(sup_env, sup_env->env.current_time + timeout);
}

static int _sup_queue_depleted(
	struct sup_reservation_environment* sup_env,
	struct reservation *res)
{
	struct list_head *pos;
	struct reservation *queued;
	int passed_earlier = 0;

	list_for_each(pos, &sup_env->depleted_reservations) {
		queued = list_entry(pos, struct reservation, list);
		if (queued->next_replenishment > res->next_replenishment) {
			list_add(&res->list, pos->prev);
			return passed_earlier;
		} else
			passed_earlier = 1;
	}

	list_add_tail(&res->list, &sup_env->depleted_reservations);

	return passed_earlier;
}

static void sup_queue_depleted(
	struct sup_reservation_environment* sup_env,
	struct reservation *res)
{
	int passed_earlier = _sup_queue_depleted(sup_env, res);

	/* check for updated replenishment time */
	if (!passed_earlier)
		sup_scheduler_update_at(sup_env, res->next_replenishment);
}

static int _sup_queue_active(
	struct sup_reservation_environment* sup_env,
	struct reservation *res)
{
	struct list_head *pos;
	struct reservation *queued;
	int passed_active = 0;

	list_for_each(pos, &sup_env->active_reservations) {
		queued = list_entry(pos, struct reservation, list);
		if (queued->priority > res->priority) {
			list_add(&res->list, pos->prev);
			return passed_active;
		} else if (queued->state == RESERVATION_ACTIVE)
			passed_active = 1;
	}

	list_add_tail(&res->list, &sup_env->active_reservations);
	return passed_active;
}

static void sup_queue_active(
	struct sup_reservation_environment* sup_env,
	struct reservation *res)
{
	int passed_active = _sup_queue_active(sup_env, res);

	/* check for possible preemption */
	if (res->state == RESERVATION_ACTIVE && !passed_active)
		sup_env->next_scheduler_update = SUP_RESCHEDULE_NOW;
	else {
		/* Active means this reservation is draining budget => make sure
		 * the scheduler is called to notice when the reservation budget has been
		 * drained completely. */
		sup_scheduler_update_after(sup_env, res->cur_budget);
	}
}

static void sup_queue_reservation(
	struct sup_reservation_environment* sup_env,
	struct reservation *res)
{
	switch (res->state) {
		case RESERVATION_INACTIVE:
			list_add(&res->list, &sup_env->inactive_reservations);
			break;

		case RESERVATION_DEPLETED:
			sup_queue_depleted(sup_env, res);
			break;

		case RESERVATION_ACTIVE_IDLE:
		case RESERVATION_ACTIVE:
			sup_queue_active(sup_env, res);
			break;
	}
}

void sup_add_new_reservation(
	struct sup_reservation_environment* sup_env,
	struct reservation* new_res)
{
	new_res->env = &sup_env->env;
	sup_queue_reservation(sup_env, new_res);
}

struct reservation* sup_find_by_id(struct sup_reservation_environment* sup_env,
	unsigned int id)
{
	struct reservation *res;

	list_for_each_entry(res, &sup_env->active_reservations, list) {
		if (res->id == id)
			return res;
	}
	list_for_each_entry(res, &sup_env->inactive_reservations, list) {
		if (res->id == id)
			return res;
	}
	list_for_each_entry(res, &sup_env->depleted_reservations, list) {
		if (res->id == id)
			return res;
	}

	return NULL;
}

static void sup_charge_budget(
	struct sup_reservation_environment* sup_env,
	lt_t delta)
{
	struct list_head *pos, *next;
	struct reservation *res;

	int encountered_active = 0;

	list_for_each_safe(pos, next, &sup_env->active_reservations) {
		/* charge all ACTIVE_IDLE up to the first ACTIVE reservation */
		res = list_entry(pos, struct reservation, list);
		if (res->state == RESERVATION_ACTIVE) {
			TRACE("sup_charge_budget ACTIVE R%u drain %llu\n", res->id, delta);
			if (encountered_active == 0 && res->blocked_by_ghost == 0) {
				TRACE("DRAIN !!\n");
				res->ops->drain_budget(res, delta);
				encountered_active = 1;
			}			
		} else {
			BUG_ON(res->state != RESERVATION_ACTIVE_IDLE);
			TRACE("sup_charge_budget INACTIVE R%u drain %llu\n", res->id, delta);
			res->ops->drain_budget(res, delta);
		}
		if (res->state == RESERVATION_ACTIVE ||
			res->state == RESERVATION_ACTIVE_IDLE)
		{
			/* make sure scheduler is invoked when this reservation expires
			 * its remaining budget */
			 TRACE("requesting scheduler update for reservation %u in %llu nanoseconds\n",
				res->id, res->cur_budget);
			 sup_scheduler_update_after(sup_env, res->cur_budget);
		}
		//if (encountered_active == 2)
			/* stop at the first ACTIVE reservation */
		//	break;
	}
	//TRACE("finished charging budgets\n");
}

static void sup_replenish_budgets(struct sup_reservation_environment* sup_env)
{
	struct list_head *pos, *next;
	struct reservation *res;

	list_for_each_safe(pos, next, &sup_env->depleted_reservations) {
		res = list_entry(pos, struct reservation, list);
		if (res->next_replenishment <= sup_env->env.current_time) {
			res->ops->replenish(res);
		} else {
			/* list is ordered by increasing depletion times */
			break;
		}
	}
	//TRACE("finished replenishing budgets\n");

	/* request a scheduler update at the next replenishment instant */
	res = list_first_entry_or_null(&sup_env->depleted_reservations,
		struct reservation, list);
	if (res)
		sup_scheduler_update_at(sup_env, res->next_replenishment);
}

void sup_update_time(
	struct sup_reservation_environment* sup_env,
	lt_t now)
{
	lt_t delta;

	/* If the time didn't advance, there is nothing to do.
	 * This check makes it safe to call sup_advance_time() potentially
	 * multiple times (e.g., via different code paths. */
	//TRACE("(sup_update_time) now: %llu, current_time: %llu\n", now, sup_env->env.current_time);
	if (unlikely(now <= sup_env->env.current_time))
		return;

	delta = now - sup_env->env.current_time;
	sup_env->env.current_time = now;

	/* check if future updates are required */
	if (sup_env->next_scheduler_update <= sup_env->env.current_time)
		sup_env->next_scheduler_update = SUP_NO_SCHEDULER_UPDATE;

	/* deplete budgets by passage of time */
	//TRACE("CHARGE###\n");
	sup_charge_budget(sup_env, delta);

	/* check if any budgets where replenished */
	//TRACE("REPLENISH###\n");
	sup_replenish_budgets(sup_env);
}

struct task_struct* sup_dispatch(struct sup_reservation_environment* sup_env)
{
	struct reservation *res, *next;
	struct task_struct *tsk = NULL;
	lt_t time_slice;

	list_for_each_entry_safe(res, next, &sup_env->active_reservations, list) {
		if (res->state == RESERVATION_ACTIVE) {
			tsk = res->ops->dispatch_client(res, &time_slice);
			if (likely(tsk)) {
				if (time_slice)
				    sup_scheduler_update_after(sup_env, time_slice);
				sup_scheduler_update_after(sup_env, res->cur_budget);
				return tsk;
			}
		}
	}

	return NULL;
}

static void sup_res_change_state(
	struct reservation_environment* env,
	struct reservation *res,
	reservation_state_t new_state)
{
	struct sup_reservation_environment* sup_env;

	sup_env = container_of(env, struct sup_reservation_environment, env);

	TRACE("reservation R%d state %d->%d at %llu\n",
		res->id, res->state, new_state, env->current_time);

	list_del(&res->list);
	/* check if we need to reschedule because we lost an active reservation */
	if (res->state == RESERVATION_ACTIVE && !sup_env->will_schedule)
		sup_env->next_scheduler_update = SUP_RESCHEDULE_NOW;
	res->state = new_state;
	sup_queue_reservation(sup_env, res);
}

void sup_init(struct sup_reservation_environment* sup_env)
{
	memset(sup_env, sizeof(*sup_env), 0);

	INIT_LIST_HEAD(&sup_env->active_reservations);
	INIT_LIST_HEAD(&sup_env->depleted_reservations);
	INIT_LIST_HEAD(&sup_env->inactive_reservations);

	sup_env->env.change_state = sup_res_change_state;

	sup_env->next_scheduler_update = SUP_NO_SCHEDULER_UPDATE;
}

struct reservation* gmp_find_by_id(struct gmp_reservation_environment* gmp_env,
	unsigned int id)
{
	struct reservation *res;

	list_for_each_entry(res, &gmp_env->active_reservations, list) {
		if (res->id == id)
			return res;
	}
	list_for_each_entry(res, &gmp_env->inactive_reservations, list) {
		if (res->id == id)
			return res;
	}
	list_for_each_entry(res, &gmp_env->depleted_reservations, list) {
		if (res->id == id)
			return res;
	}

	return NULL;
}

/*
struct next_timer_event* gmp_find_event_by_id(struct gmp_reservation_environment* gmp_env,
	unsigned int id)
{
	struct next_timer_event *event;

	list_for_each_entry(event, &gmp_env->next_events, list) {
		if (event->id == id)
			return event;
	}

	return NULL;
}
*/

struct next_timer_event* gmp_find_event_by_time(struct gmp_reservation_environment* gmp_env,
	lt_t when)
{
	struct next_timer_event *event;

	list_for_each_entry(event, &gmp_env->next_events, list) {
		if (event->next_update == when)
			return event;
	}

	return NULL;
}

/*
static void gmp_scheduler_update_at(
	struct gmp_reservation_environment* gmp_env, unsigned int id,
	event_type_t type, lt_t when)
{
	struct next_timer_event *nevent, *queued;
	struct list_head *pos;
	int found = 0;
	
	nevent = gmp_find_event_by_id(gmp_env, id);
	
	if (!nevent) {
		nevent = kzalloc(sizeof(*nevent), GFP_KERNEL);
		nevent->next_update = when;
		nevent->id = id;
		nevent->timer_armed_on = NO_CPU;
		nevent->type = type;
		
		list_for_each(pos, &gmp_env->next_events) {
			queued = list_entry(pos, struct next_timer_event, list);
			if (queued->next_update > nevent->next_update) {
				list_add(&nevent->list, pos->prev);
				found = 1;
				TRACE("NEXT_EVENT ADDED after %llu\n", queued->next_update);
				break;
			}
		}
		
		if (!found) {
			list_add_tail(&nevent->list, &gmp_env->next_events);
			TRACE("NEXT_EVENT ADDED at [0]\n");
		}
	} else {
		TRACE("EVENT FOUND at %llu T(%d), NEW EVENT %llu T(%d)\n", nevent->next_update, nevent->type, when, type);
	}
}
*/
#define TIMER_RESOLUTION 100000L

static void gmp_scheduler_update_at(
	struct gmp_reservation_environment* gmp_env,
	lt_t when)
{
	struct next_timer_event *nevent, *queued;
	struct list_head *pos;
	int found = 0;

	//when = div64_u64(when, TIMER_RESOLUTION);
	//when *= TIMER_RESOLUTION;
	
	nevent = gmp_find_event_by_time(gmp_env, when);
	
	if (!nevent) {
		nevent = kzalloc(sizeof(*nevent), GFP_KERNEL);
		nevent->next_update = when;
		nevent->timer_armed_on = NO_CPU;
		
		list_for_each(pos, &gmp_env->next_events) {
			queued = list_entry(pos, struct next_timer_event, list);
			if (queued->next_update > nevent->next_update) {
				list_add(&nevent->list, pos->prev);
				found = 1;
				TRACE("NEXT_EVENT at %llu ADDED before %llu\n", nevent->next_update, queued->next_update);
				break;
			}
		}
		
		if (!found) {
			list_add_tail(&nevent->list, &gmp_env->next_events);
			TRACE("NEXT_EVENT ADDED at %llu ADDED at HEAD\n", nevent->next_update);
		}
	} else {
		; //TRACE("EVENT FOUND at %llu, NEW EVENT %llu\n", nevent->next_update, when);
	}
}

void gmp_scheduler_update_after(
	struct gmp_reservation_environment* gmp_env, lt_t timeout)
{
	gmp_scheduler_update_at(gmp_env, gmp_env->env.current_time + timeout);
}

static void gmp_queue_depleted(
	struct gmp_reservation_environment* gmp_env,
	struct reservation *res)
{
	struct list_head *pos;
	struct reservation *queued;
	int found = 0;

	list_for_each(pos, &gmp_env->depleted_reservations) {
		queued = list_entry(pos, struct reservation, list);
		if (queued->next_replenishment > res->next_replenishment) {
			list_add(&res->list, pos->prev);
			found = 1;
		}
	}

	if (!found)
		list_add_tail(&res->list, &gmp_env->depleted_reservations);

	gmp_scheduler_update_at(gmp_env, res->next_replenishment);
}

static void gmp_queue_active(
	struct gmp_reservation_environment* gmp_env,
	struct reservation *res)
{
	struct list_head *pos;
	struct reservation *queued;
	int check_preempt = 1, found = 0;

	list_for_each(pos, &gmp_env->active_reservations) {
		queued = list_entry(pos, struct reservation, list);
		if (queued->priority > res->priority) {
			list_add(&res->list, pos->prev);
			found = 1;
			break;
		} else if (queued->scheduled_on == NO_CPU)
			check_preempt = 0;
	}

	if (!found)
		list_add_tail(&res->list, &gmp_env->active_reservations);

	/* check for possible preemption */
	if (res->state == RESERVATION_ACTIVE && !check_preempt)
		gmp_env->schedule_now = true;
	
	gmp_scheduler_update_after(gmp_env, res->cur_budget);
}

static void gmp_queue_reservation(
	struct gmp_reservation_environment* gmp_env,
	struct reservation *res)
{
	switch (res->state) {
		case RESERVATION_INACTIVE:
			list_add(&res->list, &gmp_env->inactive_reservations);
			break;

		case RESERVATION_DEPLETED:
			gmp_queue_depleted(gmp_env, res);
			break;

		case RESERVATION_ACTIVE_IDLE:
		case RESERVATION_ACTIVE:
			gmp_queue_active(gmp_env, res);
			break;
	}
}

void gmp_add_new_reservation(
	struct gmp_reservation_environment* gmp_env,
	struct reservation* new_res)
{
	new_res->env = &gmp_env->env;
	gmp_queue_reservation(gmp_env, new_res);
}

static void gmp_charge_budget(
	struct gmp_reservation_environment* gmp_env,
	lt_t delta)
{
	struct list_head *pos, *next;
	struct reservation *res;

	list_for_each_safe(pos, next, &gmp_env->active_reservations) {
		int drained = 0;
		/* charge all ACTIVE_IDLE up to the first ACTIVE reservation */
		res = list_entry(pos, struct reservation, list);
		if (res->state == RESERVATION_ACTIVE) {
			TRACE("gmp_charge_budget ACTIVE R%u drain %llu\n", res->id, delta);
			if (res->scheduled_on != NO_CPU && res->blocked_by_ghost == 0) {
				TRACE("DRAIN !!\n");
				drained = 1;
				res->ops->drain_budget(res, delta);
			}			
		} else {
			//BUG_ON(res->state != RESERVATION_ACTIVE_IDLE);
			if (res->state != RESERVATION_ACTIVE_IDLE)
				TRACE("BUG!!!!!!!!!!!! gmp_charge_budget()\n");
			TRACE("gmp_charge_budget INACTIVE R%u drain %llu\n", res->id, delta);
			//if (res->is_ghost == 1) {
				TRACE("DRAIN !!\n");
				drained = 1;
				res->ops->drain_budget(res, delta);
			//}
		}
		if ((res->state == RESERVATION_ACTIVE ||
			res->state == RESERVATION_ACTIVE_IDLE) && (drained == 1))
		{
			/* make sure scheduler is invoked when this reservation expires
			 * its remaining budget */
			 TRACE("requesting gmp_scheduler update for reservation %u in %llu nanoseconds\n",
				res->id, res->cur_budget);
			 gmp_scheduler_update_after(gmp_env, res->cur_budget);
		}
		//if (encountered_active == 2)
			/* stop at the first ACTIVE reservation */
		//	break;
	}
	//TRACE("finished charging budgets\n");
}

static void gmp_replenish_budgets(struct gmp_reservation_environment* gmp_env)
{
	struct list_head *pos, *next;
	struct reservation *res;

	list_for_each_safe(pos, next, &gmp_env->depleted_reservations) {
		res = list_entry(pos, struct reservation, list);
		if (res->next_replenishment <= gmp_env->env.current_time) {
			res->ops->replenish(res);
		} else {
			/* list is ordered by increasing depletion times */
			break;
		}
	}
	//TRACE("finished replenishing budgets\n");

	/* request a scheduler update at the next replenishment instant */
	res = list_first_entry_or_null(&gmp_env->depleted_reservations,
		struct reservation, list);
	if (res)
		gmp_scheduler_update_at(gmp_env, res->next_replenishment);
}

/* return schedule_now */
bool gmp_update_time(
	struct gmp_reservation_environment* gmp_env,
	lt_t now)
{
	lt_t delta;

	if (!gmp_env) {
		TRACE("BUG****************************************\n");
		return false;
	}
	/* If the time didn't advance, there is nothing to do.
	 * This check makes it safe to call sup_advance_time() potentially
	 * multiple times (e.g., via different code paths. */
	//TRACE("(sup_update_time) now: %llu, current_time: %llu\n", now, sup_env->env.current_time);
	if (unlikely(now <= gmp_env->env.current_time))
		return gmp_env->schedule_now;

	delta = now - gmp_env->env.current_time;
	gmp_env->env.current_time = now;


	/* deplete budgets by passage of time */
	//TRACE("CHARGE###\n");
	gmp_charge_budget(gmp_env, delta);

	/* check if any budgets where replenished */
	//TRACE("REPLENISH###\n");
	gmp_replenish_budgets(gmp_env);
	
	return gmp_env->schedule_now;
}

static void gmp_res_change_state(
	struct reservation_environment* env,
	struct reservation *res,
	reservation_state_t new_state)
{
	struct gmp_reservation_environment* gmp_env;

	gmp_env = container_of(env, struct gmp_reservation_environment, env);

	TRACE("GMP reservation R%d state %d->%d at %llu\n",
		res->id, res->state, new_state, env->current_time);

	list_del(&res->list);
	/* check if we need to reschedule because we lost an active reservation */
	if (res->state == RESERVATION_ACTIVE && !gmp_env->will_schedule)
		gmp_env->schedule_now = true;
	res->state = new_state;
	gmp_queue_reservation(gmp_env, res);
}

void gmp_init(struct gmp_reservation_environment* gmp_env)
{
	memset(gmp_env, sizeof(*gmp_env), 0);

	INIT_LIST_HEAD(&gmp_env->active_reservations);
	INIT_LIST_HEAD(&gmp_env->depleted_reservations);
	INIT_LIST_HEAD(&gmp_env->inactive_reservations);
	INIT_LIST_HEAD(&gmp_env->next_events);

	gmp_env->env.change_state = gmp_res_change_state;

	gmp_env->schedule_now = false;
	gmp_env->will_schedule = false;
	
	raw_spin_lock_init(&gmp_env->lock);
}