#include <linux/hrtimer.h>
#include <linux/percpu.h>
#include <linux/sched.h>
#include <linux/uaccess.h>
#include <linux/ctype.h>

#include <litmus/bheap.h>
#include <litmus/litmus.h>
#include <litmus/litmus_proc.h>
#include <litmus/sched_trace.h>
#include <litmus/servers.h>

#define DEBUG_SERVERS

#define TIME(x)                                 \
        ({lt_t y = x;                           \
        do_div(y, NSEC_PER_MSEC);               \
        y;})
#ifdef  DEBUG_SERVERS
#define _TRACE_TIMER(fmt, args...)				       \
        sched_trace_log_message("%d P%d*[%s@%s:%d]: " fmt " at %d\n",  \
				TRACE_ARGS,  ## args, TIME(litmus_clock()))
#define TRACE_TIMER(s, fmt, args...)					\
	do {								\
	if (is_server_linked(s))					\
	       _TRACE_TIMER(TASK_FMT " " SERVER_FMT " " fmt,		\
			    TASK_ARGS(server_task(s)),			\
			    SERVER_ARGS(s), ##args);			\
	else								\
		_TRACE_TIMER("(NULL) " SERVER_FMT " " fmt,		\
			     SERVER_ARGS(s), ##args);			\
	} while(0)
#else
#define TRACE_TIMER(s, fmt, args...)
#define _TRACE_TIMER(fmt, args...)
#endif

/* Used to run a server on a remote CPU */
DEFINE_PER_CPU(struct hrtimer_start_on_info, server_cpu_infos);

/* Memory slabs for servers */
struct kmem_cache *server_release_cache;
struct kmem_cache *server_cache;

/*
 * Okay to call if the timer is not armed.
 */
static inline int timer_cancel(struct hrtimer *timer)
{
        if (hrtimer_active(timer))
                return hrtimer_try_to_cancel(timer);
	else
		return 0;
}

static int completion_timer_arm(server_domain_t* domain, int cpu)
{
	int err = 0, on_cpu;
	lt_t now = domain->start_times[cpu];
	server_t *server = domain->linked_servers[cpu];
	lt_t budget_exhausted = now + server->budget;
	completion_timer_t *timer = &domain->completion_timers[cpu];

	/* This happens when someone attempts to call server_run when
	 * the server completes. When this happens, we can ignore the request
	 * here because completion_timer_fire will re-arm the timer if
	 * the server is still running / was run again.
	 */
	if (hrtimer_active(&timer->timer)) {
		return 0;
	}
	if (timer->armed) {
		return 0;
	}

	if (lt_after(budget_exhausted, server->deadline))
		budget_exhausted = server->deadline;

#ifdef COMPLETION_ON_MASTER
	if (domain->release_master != NO_CPU)
		on_cpu = domain->release_master;
	else
#endif
		on_cpu = cpu;

	err = 1;
	if (cpu != smp_processor_id()) {
                err = hrtimer_start_on(on_cpu, &timer->info, &timer->timer,
                                       ns_to_ktime(budget_exhausted),
                                       HRTIMER_MODE_ABS_PINNED);
	} else if (atomic_read(&timer->info.state)== HRTIMER_START_ON_INACTIVE){
		err = __hrtimer_start_range_ns(&timer->timer,
					 ns_to_ktime(budget_exhausted),
					 0 /* delta */,
					 HRTIMER_MODE_ABS_PINNED,
					 0 /* no wakeup */);
	}

	timer->armed = (err) ? 0 : 1;

	return !err;
}

static enum hrtimer_restart completion_timer_fire(struct hrtimer *timer)
{
	int cpu;
	unsigned long flags;
	enum hrtimer_restart rv;
	struct task_struct *was_running;
	completion_timer_t *completion_timer;
	server_domain_t *domain;
	server_t *server;
	lt_t budget_exhausted;

	rv  = HRTIMER_NORESTART;

	completion_timer = container_of(timer, completion_timer_t, timer);
	domain = completion_timer->domain;
	cpu = completion_timer->cpu;

	raw_spin_lock_irqsave(domain->completion_lock, flags);

	_TRACE_TIMER("completion timer firing on P%d", cpu);

	/* We got the lock before someone tried to re-arm. Proceed. */
	if (completion_timer->armed) {
		server = domain->linked_servers[cpu];
		was_running = server_task(server);

		server->budget = 0;
		server->cpu = NO_CPU;
		domain->start_times[cpu] = 0;
		domain->linked_servers[cpu] = NULL;
		domain->linked_tasks[cpu] = NULL;

		domain->server_completed(server, was_running);
	}

	/* Someone either beat us to the lock or hooked up a new server
	 * when we called server_completed. Rearm the timer.
	 */
	if (domain->linked_servers[cpu] && !completion_timer->armed) {
		server = domain->linked_servers[cpu];
		budget_exhausted = domain->start_times[cpu] + server->budget;
		if (lt_after(budget_exhausted, server->deadline))
			budget_exhausted = server->deadline;
		hrtimer_set_expires(timer, ns_to_ktime(budget_exhausted));
		completion_timer->armed = 1;

		rv = HRTIMER_RESTART;
	} else {
		completion_timer->armed = 0;
	}

	raw_spin_unlock_irqrestore(domain->completion_lock, flags);

	return rv;
}

struct kmem_cache *server_release_cache; /* In litmus.c */
static enum hrtimer_restart release_servers_fire(struct hrtimer *timer);

/*
 * Initialize heap.
 */
static server_release_heap_t* release_heap_alloc(int gfp_flags)
{
	server_release_heap_t *rh;
	rh = kmem_cache_alloc(server_release_cache, gfp_flags);
	if (rh) {
		hrtimer_init(&rh->timer, CLOCK_MONOTONIC, HRTIMER_MODE_ABS);
		rh->timer.function = release_servers_fire;
	}
	return rh;
}

static void release_heap_free(server_release_heap_t* rh)
{
	kmem_cache_free(server_release_cache, rh);
}

void server_init(server_t *server, server_domain_t *domain,
		 int id, lt_t wcet, lt_t period, int grouped)
{
	server->id = id;
	server->wcet = wcet;
	server->period = period;

	server->deadline = 0;
	server->release = 0;
	server->budget = 0;
	server->job_no = 0;
	server->cpu = NO_CPU;

	server->domain = domain;

	server->data = NULL;

	server->hn = bheap_node_alloc(GFP_ATOMIC);
	bheap_node_init(&server->hn, server);
	INIT_LIST_HEAD(&server->list);

	server->release_heap = NULL;
	if (grouped) {
		server->release_heap = release_heap_alloc(GFP_ATOMIC);
		INIT_LIST_HEAD(&server->release_list);
	}
}

void server_destroy(server_t *server)
{
	bheap_node_free(server->hn);
	if (server->release_heap) {
		release_heap_free(server->release_heap);
	}
}

server_t* server_alloc(int gfp_flags)
{
	return kmem_cache_alloc(server_cache, gfp_flags);
}

void server_free(server_t *server)
{
	kmem_cache_free(server_cache, server);
}

/*
 * Handles subtraction of lt_t without underflows.
 */
static inline lt_t lt_subtract(lt_t a, lt_t b)
{
        long long sub = (long long)a - (long long)b;
        if (sub >= 0)
                return sub;
        else
                return 0;
}

void server_run(server_t *server, struct task_struct *task)
{
	int armed, cpu = task->rt_param.linked_on;
	server_domain_t *domain = server->domain;

	server->cpu = cpu;

	domain->linked_servers[cpu] = server;
	domain->linked_tasks[cpu] = task;
	domain->start_times[cpu] = litmus_clock();

	/* Arm completion timer */
	armed = completion_timer_arm(domain, cpu);
	domain->completion_timers[cpu].armed = armed;
}

void server_stop(server_t *server)
{
	int cpu;
	lt_t elapsed_time, now = litmus_clock();
	server_domain_t *domain = server->domain;

	if (!is_server_linked(server)) {
		return;
	}

	cpu = server->cpu;
	BUG_ON(cpu == NO_CPU);

	/* Calculate remaining budget */
	elapsed_time = lt_subtract(now, domain->start_times[cpu]);
	server->budget = lt_subtract(server->budget, elapsed_time);

	server->cpu = NO_CPU;

	/* Set domain state */
	domain->completion_timers[cpu].armed = 0;
	domain->linked_servers[cpu] = NULL;
	domain->linked_tasks[cpu] = NULL;
	timer_cancel(&domain->completion_timers[cpu].timer);
}

void server_release(server_t *server)
{
	server->budget = server->wcet;
	server->release = server->deadline;
	server->deadline += server->period;
	++server->job_no;

	/* Need to reset for budget calculations */
	if (is_server_linked(server))
		server->domain->start_times[server->cpu] = litmus_clock();
}

void server_release_at(server_t *server, lt_t time)
{
	server->deadline = time;
	server_release(server);
}

/******************************************************************************
 * Proc methods
 ******************************************************************************/

static int server_proc_read(char* page, char **start, off_t off,
			     int count,	int *eof, void *data)
{
	int length;
	server_proc_t *proc = (server_proc_t*)data;

	proc->page   = page;
	proc->length = 0;
	proc->list_servers(proc);

	length = proc->length;
	*eof = 1;

	proc->length = 0;
	proc->page = NULL;

	return length;
}

void list_server(server_t *server, int cpu, server_proc_t *proc)
{
	if (cpu == NO_CPU) {
		proc->length +=
			snprintf(proc->page + proc->length,
				 PAGE_SIZE - proc->length,
				 "%8llu    %8llu\n",
				 server->wcet, server->period);
	} else {
		proc->length +=
			snprintf(proc->page + proc->length,
				 PAGE_SIZE - proc->length,
				 "%8llu    %8llu    %3d\n",
				 server->wcet, server->period, cpu);
	}
}

/*
 * Validate server parameters.
 */
static inline int server_param_check(unsigned long long wcet,
				     unsigned long long period,
				     int cpu)
{
	int rv = 0;

	if (wcet <= 0) {
		printk(KERN_WARNING "Invalid WCET '%llu'\n", wcet);
		rv = -EINVAL;
		goto out;
	}

	if (period < wcet) {
		printk(KERN_WARNING "Invalid period '%llu'\n", period);
		rv = -EINVAL;
		goto out;
	}

	if (cpu != NO_CPU && (cpu < 0 || cpu >= nr_cpu_ids)) {
		printk(KERN_WARNING "Invalid CPU '%d'\n", cpu);
		rv = -EINVAL;
		goto out;
	}
 out:
	return rv;
}

/* Macro to see if we are in the buffer's range and not at the null byte */
#define buf_in_range(buf, pos, max) (buf <= pos && pos < (buf + max)  && *pos)

#define find_newline(buf, pos, max)					\
	do {								\
		while (buf_in_range(buf, pos, max) &&			\
				*pos != '\n')				\
			++pos;						\
	} while (0)

static int server_proc_write(struct file *file, const char __user *input,
			      unsigned long count, void *data)
{
	server_proc_t *proc = (server_proc_t*)data;
#define SERVER_PROC_BUF 512
	char buffer[SERVER_PROC_BUF];
	unsigned long long wcet, period;
	char *pos, *newline, *space_check;
	int nums_converted, chars_seen, ret, cpu;

	/* Allow plugin to stop any running servers */
	proc->stop_servers();

	if (count >= SERVER_PROC_BUF){
		printk(KERN_WARNING "proc buffer possibly too small in %s.\n",
				__func__);
		return -ENOSPC;
	}

	memset(buffer, 0, SERVER_PROC_BUF);

	/* Input is definitely < SERVER_PROC_BUF (see above check) */
	if (copy_from_user(buffer, input, count))
		return -EFAULT;

	buffer[SERVER_PROC_BUF-1] = '\0';
	pos = buffer;

	while (buf_in_range(buffer, pos, SERVER_PROC_BUF)) {
		newline = pos;
		find_newline(buffer, newline, SERVER_PROC_BUF);
		if (buf_in_range(buffer, newline, SERVER_PROC_BUF)) {
			/* If there was a newline character */
			*newline = '\0';
		}
		nums_converted = sscanf(pos, "%llu %llu %d%n", &wcet,
					&period, &cpu, &chars_seen);
		if (nums_converted == 2)
			cpu = NO_CPU;
		if (nums_converted != 2 && nums_converted != 3) {
			printk(KERN_WARNING "Didn't see 2-3 integers for "
					"server config: %s\n", pos);
			goto loop_end;
		}
		/* space_check = pos + chars_seen; */
		/* if (space_check != newline) { */
		/* 	/\* If the newline was not right after the numbers */
		/* 	 * converted, ensure extra characters are just space */
		/* 	 *\/ */
		/* 	for (; *space_check; space_check++) { */
		/* 		if (!isspace(*space_check)) { */
		/* 			printk(KERN_WARNING "Extra characters " */
		/* 					"in line: %s\n", pos); */
		/* 			goto loop_end; */
		/* 		} */
		/* 	} */
		/* } */

		ret = server_param_check(wcet, period, cpu);
		if (ret) goto loop_end;

		ret = proc->admit_server(wcet, period, cpu);
		if (ret) {
			printk(KERN_WARNING "Litmus plugin rejects server with "
			      "period: %llu, wcet: %llu, cpu: %d\n",
			       period, wcet, cpu);
			goto loop_end; /* Currently does nothing */
		}
loop_end:
		pos = newline + 1; /* Consider next line */
	}

	return count;
}

server_proc_t* server_proc_init(server_domain_t *domain,
				struct proc_dir_entry *proc_dir, char *file,
				admit_server_t admit_server,
				list_servers_t list_servers,
				stop_servers_t stop_servers)
{
	server_proc_t *server_proc = NULL;
	struct proc_dir_entry *entry;

	entry = create_proc_entry(file, 0644, proc_dir);
	if (!entry) {
		printk(KERN_ERR "Could not create proc entry: %s.\n", file);
		goto out;
 	}

	server_proc = kmalloc(sizeof(server_proc_t), GFP_ATOMIC);

	entry->data = server_proc;
	entry->read_proc  = server_proc_read;
	entry->write_proc = server_proc_write;

	server_proc->entry = entry;
	server_proc->admit_server = admit_server;
	server_proc->list_servers = list_servers;
	server_proc->stop_servers = stop_servers;
	server_proc->length = 0;
	server_proc->page = NULL;

	INIT_LIST_HEAD(&server_proc->list);
	list_add(&server_proc->list, &domain->server_procs);

 out:
	return server_proc;
}

void server_proc_exit(server_proc_t *proc)
{
	remove_proc_entry(proc->entry->name, proc->entry->parent);
	list_del(&proc->list);
	kfree(proc);
}

/******************************************************************************
 * Domain methods
 ******************************************************************************/

void server_domain_init(server_domain_t *domain,
			servers_released_t servers_released,
			server_completed_t server_completed,
			int release_master, raw_spinlock_t *completion_lock)
{
	int i;
	BUG_ON(!servers_released || !server_completed);

	INIT_LIST_HEAD(&domain->tobe_released);
	for (i = 0; i < SERVER_RELEASE_QUEUE_SLOTS; i++)
		INIT_LIST_HEAD(&domain->release_queue[i]);

	raw_spin_lock_init(&domain->release_lock);
	raw_spin_lock_init(&domain->tobe_lock);


	domain->release_master   = release_master;
	domain->completion_lock  = completion_lock;
	domain->server_completed = server_completed;
	domain->servers_released = servers_released;

	INIT_LIST_HEAD(&domain->server_procs);

	domain->completion_timers =
		kmalloc(NR_CPUS*sizeof(completion_timer_t), GFP_ATOMIC);
	domain->linked_servers =
		kmalloc(NR_CPUS*sizeof(server_t*), GFP_ATOMIC);
	domain->linked_tasks =
		kmalloc(NR_CPUS*sizeof(struct task_struct*), GFP_ATOMIC);
	domain->start_times =
		kmalloc(NR_CPUS*sizeof(lt_t), GFP_ATOMIC);

	for_each_online_cpu(i) {
		domain->linked_tasks[i] = NULL;
		domain->linked_servers[i] = NULL;
		domain->start_times[i] = 0;

		/* Initialize the completion timer info */
		domain->completion_timers[i].armed = 0;
		domain->completion_timers[i].cpu = i;
		hrtimer_init(&domain->completion_timers[i].timer,
			     CLOCK_MONOTONIC,
			     HRTIMER_MODE_ABS);
		domain->completion_timers[i].timer.function =
			completion_timer_fire;
		hrtimer_start_on_info_init(&domain->completion_timers[i].info);
		domain->completion_timers[i].domain = domain;
	}
}

void server_domain_destroy(server_domain_t *domain)
{
	struct list_head *pos, *safe;
	server_proc_t *proc;

	kfree(domain->completion_timers);
	kfree(domain->linked_tasks);
	kfree(domain->linked_servers);
	kfree(domain->start_times);

	list_for_each_safe(pos, safe, &domain->server_procs) {
		proc = list_entry(pos, server_proc_t, list);
		server_proc_exit(proc);
	}
}

static unsigned int time2slot(lt_t time)
{
	return (unsigned int) time2quanta(time, FLOOR) %
		SERVER_RELEASE_QUEUE_SLOTS;
}

/*
 * Send a list of servers to a client callback.
 */
static enum hrtimer_restart release_servers_fire(struct hrtimer *timer)
{
	unsigned long flags;
	server_release_heap_t *rh;

	rh = container_of(timer, server_release_heap_t, timer);

	raw_spin_lock_irqsave(&rh->domain->release_lock, flags);

	/* Remove from release queue */
	list_del(&rh->list);

	raw_spin_unlock_irqrestore(&rh->domain->release_lock, flags);

	/* Call release callback */
	rh->domain->servers_released(&rh->servers);

	return HRTIMER_NORESTART;
}

/*
 * Caller must hold release lock.
 * Will return heap for given time. If no such heap exists prior to
 * the invocation it will be created.
 */
static server_release_heap_t* get_release_heap(server_domain_t *rt,
					       server_t *server,
					       int use_server_heap)
{
	struct list_head *pos;
	server_release_heap_t *heap = NULL;
	server_release_heap_t *rh;
	lt_t release_time = server->release;
	unsigned int slot = time2slot(release_time);

	/* Initialize pos for the case that the list is empty */
	pos = rt->release_queue[slot].next;
	list_for_each(pos, &rt->release_queue[slot]) {
		rh = list_entry(pos, server_release_heap_t, list);
		if (release_time == rh->release_time) {
			/* Perfect match -- this happens on hyperperiod
			 * boundaries
			 */
			heap = rh;
			break;
		} else if (lt_before(release_time, rh->release_time)) {
			/* We need to insert a new node since rh is
			 * already in the future
			 */
			break;
		}
	}
	if (!heap && use_server_heap) {
		/* Use pre-allocated release heap */
		rh = server->release_heap;
		rh->domain = rt;
		rh->release_time = release_time;

		/* Add to release queue */
		list_add(&rh->list, pos->prev);
		heap = rh;
	}
	return heap;
}

/*
 * Prepare a server's release_heap for use.
 */
static int reinit_release_heap(server_t *server)
{
	int rv = 0;
	server_release_heap_t* rh;

	/* Use pre-allocated release heap */
	rh = server->release_heap;

	/* WARNING: If the CPU still holds the release_lock at this point,
	 *          deadlock may occur!
	 */
	rv = hrtimer_try_to_cancel(&rh->timer);

	/* The timer callback is running, it is useless to add
	 * to the release heap now.
	 */
	if (rv == -1) {
		rv = 0;
		goto out;
	}

	/* Under no cirumstances should the timer have been active
	 * but not running.
	 */
	rv = 1;

	/* initialize */
	INIT_LIST_HEAD(&rh->servers);
	atomic_set(&rh->info.state, HRTIMER_START_ON_INACTIVE);
 out:
	return rv;
}

/*
 * Arm the release timer for the next set of servers.
 */
static int arm_release_timer(server_domain_t *domain)
{
	int rv = 1;
	struct list_head list;
	struct list_head *pos, *safe;
	server_t *server;
	server_release_heap_t *rh;

	list_replace_init(&domain->tobe_released, &list);

	list_for_each_safe(pos, safe, &list) {
		/* Pick server from work list */
		server = list_entry(pos, server_t, release_list);
		list_del(pos);

		/* Put into release heap while holding release_lock */
		raw_spin_lock(&domain->release_lock);

		rh = get_release_heap(domain, server, 0);
		if (!rh) {
			/* Need to use our own, but drop lock first */
			raw_spin_unlock(&domain->release_lock);

			rv = reinit_release_heap(server);

			/* Bail! We missed the release time */
			if (!rv) {
				rv = 0;
				goto out;
			}

			raw_spin_lock(&domain->release_lock);

			rh = get_release_heap(domain, server, 1);
		}

		list_add(&server->release_list, &rh->servers);

		raw_spin_unlock(&domain->release_lock);

		/* To avoid arming the timer multiple times, we only let the
		 * owner do the arming (which is the "first" task to reference
		 * this release_heap anyway).
		 */
		if (rh == server->release_heap) {
			/* We cannot arm the timer using hrtimer_start()
			 * as it may deadlock on rq->lock
			 *
			 * PINNED mode is ok on both local and remote CPU
			 */
			if (domain->release_master == NO_CPU) {
				__hrtimer_start_range_ns(&rh->timer,
						ns_to_ktime(rh->release_time),
						0, HRTIMER_MODE_ABS_PINNED, 0);
			} else {
				hrtimer_start_on(domain->release_master,
					&rh->info, &rh->timer,
					ns_to_ktime(rh->release_time),
					HRTIMER_MODE_ABS_PINNED);
			}
		}
	}
 out:
	return rv;
}

int add_server_release(server_t *server, server_domain_t *domain)
{
	list_add(&server->release_list, &domain->tobe_released);
	return arm_release_timer(domain);
}

static int __init init_servers(void)
{
	server_cache = KMEM_CACHE(server, SLAB_PANIC);
	server_release_cache = KMEM_CACHE(server_release_heap, SLAB_PANIC);
	return 1;
}

static void exit_servers(void)
{
	kmem_cache_destroy(server_cache);
	kmem_cache_destroy(server_release_cache);
}


module_init(init_servers);
module_exit(exit_servers);