#include <litmus/sched_plugin.h>
#include <litmus/trace.h>
#include <litmus/litmus.h>

#ifdef CONFIG_REALTIME_AUX_TASKS
#include <litmus/rt_param.h>
#include <litmus/aux_tasks.h>

static int admit_aux_task(struct task_struct *t)
{
	int retval = 0;
	struct task_struct *leader = t->group_leader;


	/* budget enforcement increments job numbers.  job numbers are used in
	 * tie-breaking of aux_tasks.  method helps ensure:
	 * 1) aux threads with no inherited priority can starve another (they share
	 *    the CPUs equally.
	 * 2) aux threads that inherit the same priority cannot starve each other.
	 *
	 * Assuming aux threads are well-behavied (they do very little work and
	 * suspend), risk of starvation should not be an issue, but this is a
	 * fail-safe.
	 */
	struct rt_task tp = {
		.period = 1000000,   /* 1ms */
		.relative_deadline = 1000000,
		.exec_cost = 1000000, /* allow full utilization */
		.phase = 0,
		.cpu = task_cpu(leader),  /* take CPU of group leader */
		.budget_policy = QUANTUM_ENFORCEMENT,
		.budget_signal_policy = NO_SIGNALS,
		.cls = RT_CLASS_BEST_EFFORT
	};

	struct sched_param param = { .sched_priority = 0};

	tsk_rt(t)->task_params = tp;
	retval = sched_setscheduler_nocheck(t, SCHED_LITMUS, &param);

	return retval;
}

int exit_aux_task(struct task_struct *t)
{
	int retval = 0;
	struct task_struct *leader = t->group_leader;

	BUG_ON(!tsk_rt(t)->is_aux_task);

	TRACE_CUR("Aux task %s/%d is exiting from %s/%d.\n", t->comm, t->pid, leader->comm, leader->pid);

	list_del(&tsk_rt(t)->aux_task_node);

	tsk_rt(t)->is_aux_task = 0;

	if (tsk_rt(t)->inh_task) {
		litmus->decrease_prio(t, NULL);
	}

	return retval;
}

static int aux_tasks_increase_priority(struct task_struct *leader, struct task_struct *hp)
{
	int retval = 0;
	struct list_head *pos;

	TRACE_CUR("Increasing priority of aux tasks in group %s/%d.\n", leader->comm, leader->pid);

	list_for_each(pos, &tsk_aux(leader)->aux_tasks) {
		struct task_struct *aux =
			container_of(list_entry(pos, struct rt_param, aux_task_node),
						struct task_struct, rt_param);

		if (!is_realtime(aux)) {
			TRACE_CUR("skipping non-real-time aux task %s/%d\n", aux->comm, aux->pid);
		}

		// aux tasks don't touch rt locks, so no nested call needed.
		TRACE_CUR("increasing %s/%d.\n", aux->comm, aux->pid);
		retval = litmus->__increase_prio(aux, hp);
	}

	return retval;
}

static int aux_tasks_decrease_priority(struct task_struct *leader, struct task_struct *hp)
{
	int retval = 0;
	struct list_head *pos;

	TRACE_CUR("Decreasing priority of aux tasks in group %s/%d.\n", leader->comm, leader->pid);

	list_for_each(pos, &tsk_aux(leader)->aux_tasks) {
		struct task_struct *aux =
			container_of(list_entry(pos, struct rt_param, aux_task_node),
						 struct task_struct, rt_param);

		if (!is_realtime(aux)) {
			TRACE_CUR("skipping non-real-time aux task %s/%d\n", aux->comm, aux->pid);
		}
		else {
			TRACE_CUR("decreasing %s/%d.\n", aux->comm, aux->pid);
			retval = litmus->__decrease_prio(aux, hp);
		}
	}

	return retval;
}

int aux_task_owner_increase_priority(struct task_struct *t)
{
	int retval = 0;
	struct task_struct *leader;
	struct task_struct *hp = NULL;

	BUG_ON(!tsk_rt(t)->has_aux_tasks);
	BUG_ON(!is_realtime(t));
	BUG_ON(!binheap_is_in_heap(&tsk_rt(t)->aux_task_owner_node));

	leader = t->group_leader;

	TRACE_CUR("task %s/%d in group %s/%d increasing priority.\n", t->comm, t->pid, leader->comm, leader->pid);

	hp = container_of(binheap_top_entry(&tsk_aux(leader)->aux_task_owners, struct rt_param, aux_task_owner_node),
					  struct task_struct, rt_param);

	if (hp == t) {
		goto out; // already hp, nothing to do.
	}

	binheap_decrease(&tsk_rt(t)->aux_task_owner_node, &tsk_aux(leader)->aux_task_owners);

	hp = container_of(binheap_top_entry(&tsk_aux(leader)->aux_task_owners, struct rt_param, aux_task_owner_node),
					  struct task_struct, rt_param);

	if (hp == t) {
		TRACE_CUR("%s/%d is new hp in group %s/%d.\n", t->comm, t->pid, leader->comm, leader->pid);
		retval = aux_tasks_increase_priority(leader,
					   (tsk_rt(hp)->inh_task) ? tsk_rt(hp)->inh_task : hp);
	}

out:
	return retval;
}

int aux_task_owner_decrease_priority(struct task_struct *t)
{
	int retval = 0;
	struct task_struct *leader;
	struct task_struct *hp = NULL;
	struct task_struct *new_hp = NULL;

	BUG_ON(!tsk_rt(t)->has_aux_tasks);
	BUG_ON(!is_realtime(t));
	BUG_ON(!binheap_is_in_heap(&tsk_rt(t)->aux_task_owner_node));

	leader = t->group_leader;

	TRACE_CUR("task %s/%d in group %s/%d decresing priority.\n", t->comm, t->pid, leader->comm, leader->pid);

	hp = container_of(binheap_top_entry(&tsk_aux(leader)->aux_task_owners, struct rt_param, aux_task_owner_node),
					  struct task_struct, rt_param);
	binheap_delete(&tsk_rt(t)->aux_task_owner_node, &tsk_aux(leader)->aux_task_owners);
	binheap_add(&tsk_rt(t)->aux_task_owner_node, &tsk_aux(leader)->aux_task_owners,
				struct rt_param, aux_task_owner_node);
	new_hp = container_of(binheap_top_entry(&tsk_aux(leader)->aux_task_owners, struct rt_param, aux_task_owner_node),
						  struct task_struct, rt_param);

	if (hp == t && new_hp != t) {
		TRACE_CUR("%s/%d is no longer hp in group %s/%d.\n", t->comm, t->pid, leader->comm, leader->pid);
		retval = aux_tasks_decrease_priority(leader,
				   (tsk_rt(new_hp)->inh_task) ? tsk_rt(new_hp)->inh_task : new_hp);
	}

	return retval;
}



long enable_aux_task_owner(struct task_struct *t)
{
	long retval = 0;
	struct task_struct *leader = t->group_leader;
	struct task_struct *hp;

	if (!tsk_rt(t)->has_aux_tasks) {
		TRACE_CUR("task %s/%d is not an aux owner\n", t->comm, t->pid);
		return -1;
	}

	BUG_ON(!is_realtime(t));

	if (binheap_is_in_heap(&tsk_rt(t)->aux_task_owner_node)) {
		TRACE_CUR("task %s/%d is already active\n", t->comm, t->pid);
		goto out;
	}

	binheap_add(&tsk_rt(t)->aux_task_owner_node, &tsk_aux(leader)->aux_task_owners,
				struct rt_param, aux_task_owner_node);

	hp = container_of(binheap_top_entry(&tsk_aux(leader)->aux_task_owners, struct rt_param, aux_task_owner_node),
					  struct task_struct, rt_param);
	if (hp == t) {
		/* we're the new hp */
		TRACE_CUR("%s/%d is new hp in group %s/%d.\n", t->comm, t->pid, leader->comm, leader->pid);

		retval = aux_tasks_increase_priority(leader,
					   (tsk_rt(hp)->inh_task)? tsk_rt(hp)->inh_task : hp);
	}


out:
	return retval;
}

long disable_aux_task_owner(struct task_struct *t)
{
	long retval = 0;
	struct task_struct *leader = t->group_leader;
	struct task_struct *hp;
	struct task_struct *new_hp = NULL;

	if (!tsk_rt(t)->has_aux_tasks) {
		TRACE_CUR("task %s/%d is not an aux owner\n", t->comm, t->pid);
		return -1;
	}

	BUG_ON(!is_realtime(t));

	if (!binheap_is_in_heap(&tsk_rt(t)->aux_task_owner_node)) {
		TRACE_CUR("task %s/%d is already not active\n", t->comm, t->pid);
		goto out;
	}

	TRACE_CUR("task %s/%d exiting from group %s/%d.\n", t->comm, t->pid, leader->comm, leader->pid);

	hp = container_of(binheap_top_entry(&tsk_aux(leader)->aux_task_owners, struct rt_param, aux_task_owner_node),
					  struct task_struct, rt_param);
	binheap_delete(&tsk_rt(t)->aux_task_owner_node, &tsk_aux(leader)->aux_task_owners);

	if (!binheap_empty(&tsk_aux(leader)->aux_task_owners)) {
		new_hp = container_of(binheap_top_entry(&tsk_aux(leader)->aux_task_owners, struct rt_param, aux_task_owner_node),
							  struct task_struct, rt_param);
	}

	if (hp == t && new_hp != t) {
		struct task_struct *to_inh = NULL;

		TRACE_CUR("%s/%d is no longer hp in group %s/%d.\n", t->comm, t->pid, leader->comm, leader->pid);

		if (new_hp) {
			to_inh = (tsk_rt(new_hp)->inh_task) ? tsk_rt(new_hp)->inh_task : new_hp;
		}

		retval = aux_tasks_decrease_priority(leader, to_inh);
	}

out:
	return retval;
}


static int aux_task_owner_max_priority_order(struct binheap_node *a,
									   struct binheap_node *b)
{
	struct task_struct *d_a = container_of(binheap_entry(a, struct rt_param, aux_task_owner_node),
										   struct task_struct, rt_param);
	struct task_struct *d_b = container_of(binheap_entry(b, struct rt_param, aux_task_owner_node),
										   struct task_struct, rt_param);

	BUG_ON(!d_a);
	BUG_ON(!d_b);

	return litmus->compare(d_a, d_b);
}


static long __do_enable_aux_tasks(void)
{
	long retval = 0;
	struct task_struct *leader;
	struct task_struct *t;

	leader = current->group_leader;

	if (!tsk_aux(leader)->initialized) {
		INIT_LIST_HEAD(&tsk_aux(leader)->aux_tasks);
		INIT_BINHEAP_HANDLE(&tsk_aux(leader)->aux_task_owners, aux_task_owner_max_priority_order);
		tsk_aux(leader)->initialized = 1;
	}

	t = leader;
	do {
		/* doesn't hurt to initialize them both */
		INIT_LIST_HEAD(&tsk_rt(t)->aux_task_node);
		INIT_BINHEAP_NODE(&tsk_rt(t)->aux_task_owner_node);

		TRACE_CUR("Checking task in %s/%d: %s/%d = (p = %llu):\n",
				  leader->comm, leader->pid, t->comm, t->pid,
				  tsk_rt(t)->task_params.period);

		/* inspect heap_node to see if it is an rt task */
		if (tsk_rt(t)->task_params.period == 0) { //||
			//			tsk_rt(t)->task_params.period == MAGIC_AUX_TASK_PERIOD) {
			if (!tsk_rt(t)->is_aux_task) {
				TRACE_CUR("AUX task in %s/%d: %s/%d:\n", leader->comm, leader->pid, t->comm, t->pid);
				/* hasn't been aux_tasks_increase_priorityted into rt. make it a aux. */
				tsk_rt(t)->is_aux_task = 1;
				list_add_tail(&tsk_rt(t)->aux_task_node, &tsk_aux(leader)->aux_tasks);

				(void)admit_aux_task(t);
			}
			else {
				TRACE_CUR("AUX task in %s/%d is already set up: %s/%d\n", leader->comm, leader->pid, t->comm, t->pid);
			}
		}
		else {
			if (!tsk_rt(t)->has_aux_tasks) {
				TRACE_CUR("task in %s/%d: %s/%d:\n", leader->comm, leader->pid, t->comm, t->pid);
				tsk_rt(t)->has_aux_tasks = 1;
			}
			else {
				TRACE_CUR("task in %s/%d is already set up: %s/%d\n", leader->comm, leader->pid, t->comm, t->pid);
			}
		}

		t = next_thread(t);
	} while(t != leader);


	if (!binheap_empty(&tsk_aux(leader)->aux_task_owners)) {
		struct task_struct *hp = container_of(binheap_top_entry(&tsk_aux(leader)->aux_task_owners, struct rt_param, aux_task_owner_node),
											  struct task_struct, rt_param);
		TRACE_CUR("found hp in group: %s/%d\n", hp->comm, hp->pid);
		retval = aux_tasks_increase_priority(leader,
											 (tsk_rt(hp)->inh_task)? tsk_rt(hp)->inh_task : hp);
	}

	return retval;
}

static long __do_disable_aux_tasks(void)
{
	long retval = 0;
	struct task_struct *leader;
	struct task_struct *t;

	leader = current->group_leader;

	t = leader;
	do {
		if (tsk_rt(t)->is_aux_task) {

			TRACE_CUR("%s/%d is an aux task.\n", t->comm, t->pid);

			if (is_realtime(t)) {
				long temp_retval;
				struct sched_param param = { .sched_priority = 0};

				TRACE_CUR("%s/%d is real-time. Changing policy to SCHED_NORMAL.\n", t->comm, t->pid);

				temp_retval = sched_setscheduler_nocheck(t, SCHED_NORMAL, &param);

				if (temp_retval != 0) {
					TRACE_CUR("error changing policy of %s/%d to SCHED_NORMAL\n", t->comm, t->pid);
					if (retval == 0) {
						retval = temp_retval;
					}
					else {
						TRACE_CUR("prior error (%d) masks new error (%d)\n", retval, temp_retval);
					}
				}
			}

			tsk_rt(t)->is_aux_task = 0;
		}
		t = next_thread(t);
	} while(t != leader);

	return retval;
}

asmlinkage long sys_set_aux_tasks(int enable)
{
	long retval;

	read_lock_irq(&tasklist_lock);

	if (enable) {
		retval = __do_enable_aux_tasks();
	}
	else {
		retval = __do_disable_aux_tasks();
	}

	read_unlock_irq(&tasklist_lock);

	return retval;
}

#else

asmlinkage long sys_set_aux_tasks(int enable)
{
	printk("Unsupported. Recompile with CONFIG_REALTIME_AUX_TASKS.\n");
	return -EINVAL;
}

#endif