#define EXPORT_SYMTAB
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/init.h>
#include <linux/fs.h>
#include <linux/sched.h>
#include <linux/pid.h> /* For pid types */
#include <linux/tty.h> /* For the tty declarations */
#include <linux/version.h> /* For LINUX_VERSION_CODE */
#include <linux/slab.h>
#include <linux/spinlock.h>
#include <linux/delay.h>
#include <linux/random.h>

#include "./include/chardev.h"
#include "./include/utils.h"
#include "./include/queue.h"
#include "./include/message.h"

#define AUDIT if (1)

// ------------------------------
// Data Structures

/**
 * struct session_state
 *
 * Contains data regarding a single I/O session.
 */
typedef struct session_state {
	unsigned long send_timeout; // in nanoseconds
	unsigned long recv_timeout; // in nanoseconds

	atomic_t valid;

	int minor;

	struct mutex writes_synchronizer;
	atomic_t mtimer_size; // num of msg yet to be posted
	struct mtimer {
		// message to be posted
		message *m;

		// add this in order so that we don't have to iterate
		// the list everytime in order to get session_state
		// pointer.
		struct session_state *state;

		unsigned int valid;

		struct hrtimer hr_timer;
		struct list_head list;
	} m_timer;

} session_state;

/**
 * struct session_state
 *
 * Contains data regarding a single device instance
 */
typedef struct _object_state {
	// FIFO queue
	struct mutex queue_synchronizer;
	queue *q;

	// list of active I/O sessions
	struct mutex session_list_synchronizer;
	struct active_session {
		session_state *session;
		struct list_head list;
	} active_session;

	// waitqueue in which to place read I/O sessions in timeout
	wait_queue_head_t wq;

	// list of read I/O sessions in timeout. Used to set the
	// correct value for next_reader ptr.
	struct mutex readers_synchronizer;
	struct reader {
		int valid;
		int id;
		struct list_head list;
	} reader;

	// next I/O session ready to read in case anyone writes a
	// message to the queue.
	unsigned int next_reader_id;

	// indicates when the device is being flushed
	int flushing;

} object_state;

// ------------------------------
// Global data

int MAX_MESSAGE_SIZE = 4096;
int MAX_STORAGE_SIZE = 4096 * 20;
int size_so_far = 0;

module_param(MAX_MESSAGE_SIZE, int, 0660);
module_param(MAX_STORAGE_SIZE, int, 0660);

object_state objects[DEVICE_MINORS];

// ------------------------------
// Main functions signatures

static int dev_open(struct inode *, struct file *);
static int dev_release(struct inode *, struct file *);
static ssize_t dev_write(struct file *, const char *, size_t, loff_t *);
static ssize_t dev_read(struct file *filp, char *buff, size_t len, loff_t *off);
static long dev_ioctl(struct file *filp, unsigned int command,
		      unsigned long param);
static int dev_flush(struct file *filp, fl_owner_t id);

// ------------------------------
// Secondary functions signatures

static void wake_readers(object_state *the_object);
static enum hrtimer_restart hrtimer_callback(struct hrtimer *timer);
static void remove_session(session_state *state);

// ------------------------------
// Secondary functions implementation

/**
 * wake_readers()
 *
 * @the_object: ptr to object_state, representing a specific device
 * instance
 *
 * Checks if there are any readers waiting, and in case wakes them up.
 */
static void wake_readers(object_state *the_object)
{
	struct reader *reader;
	struct list_head *head, *pos;

	head = &the_object->reader.list;

	mutex_lock(&(the_object->readers_synchronizer));
	list_for_each (pos, head) {
		reader = list_entry(pos, struct reader, list);

		if (reader->valid) {
			the_object->next_reader_id = reader->id;
			// invalidates reader so next time it is not
			// taken even if it was not deleted in time.
			reader->valid = 0;
			wake_up(&(the_object->wq));
			break;
		}
	}
	mutex_unlock(&(the_object->readers_synchronizer));
}

/**
 * hrtimer_callback()
 *
 * @timer: ptr to high-resolution timer struct
 *
 * Execute the delayed write operation.
 */
static enum hrtimer_restart hrtimer_callback(struct hrtimer *timer)
{
	struct mtimer *m_timer;
	object_state *the_object;
	session_state *state;

	m_timer = (struct mtimer *)container_of(timer, struct mtimer, hr_timer);
	state = m_timer->state;
	the_object = &objects[state->minor];

	AUDIT
	printk("%s: hrtimer_callback() for state 0x%p with hrtimer 0x%p\n",
	       MODNAME, state, timer);

	// remove timer from list of timers
	mutex_lock(&(state->writes_synchronizer));
	list_del(&m_timer->list);
	mutex_unlock(&(state->writes_synchronizer));

	// if message still valid, post message in the queue
	if (m_timer->valid) {
		mutex_lock(&(the_object->queue_synchronizer));
		enqueue(the_object->q, m_timer->m);
		wake_readers(the_object);
		mutex_unlock(&(the_object->queue_synchronizer));
	} else {
		AUDIT
		printk("%s: hrtimer_callback() for state 0x%p with hrtimer 0x%p cancalled\n",
		       MODNAME, state, timer);

		// otherwise update device's size and de-allocate mem
		atomic_set((atomic_t *)&size_so_far,
			   size_so_far - m_timer->m->length);

		destroy_message(m_timer->m);
	}

	// de-allocate memory
	kfree(m_timer);

	// signal we are done
	atomic_dec(&state->mtimer_size);

	if ((atomic_read(&state->valid) == 0) &&
	    (atomic_read(&state->mtimer_size) == 0)) {
		// we are the last hr-timer for current I/O session,
		// thus we have to free memory and remove our session
		// from the list of active sessions.
		remove_session(state);

		// TODO: fix double-free in potential race condition
		// maybe with a barrier before
		kfree(state);
	}

	return HRTIMER_NORESTART;
}

/**
 * remove_session()
 *
 * @filp: pointer to I/O session.
 *
 * Removes the I/O session passed as an argument from the object_data
 * struct related to the specific driver instance used by the session.
 */
static void remove_session(session_state *state)
{
	object_state *the_object;
	struct active_session *session, *active_session;
	struct list_head *head, *pos;

	the_object = &objects[state->minor];

	head = &the_object->active_session.list;

	mutex_lock(&the_object->session_list_synchronizer);
	list_for_each (pos, head) {
		session = list_entry(pos, struct active_session, list);

		if (session->session == state) {
			active_session = session;
			break;
		}
	}

	// remove elem from list
	list_del(&active_session->list);
	mutex_unlock(&the_object->session_list_synchronizer);

	kfree(active_session);
}

// ------------------------------
// Main functions implementation

/**
 * dev_open()
 *
 * @inode: ptr to inode
 * @filp: ptr to I/O session.
 * 
 * allocates memory for a new I/O session. It registers the newly
 * created session in the list of the particular device file instance
 * used.
 *
 */
static int dev_open(struct inode *inode, struct file *filp)
{
	session_state *state;
	struct active_session *session;
	object_state *the_object;

	the_object = &objects[get_minor(filp)];

	state = kmalloc(sizeof(session_state), GFP_KERNEL);
	session = kmalloc(sizeof(struct active_session), GFP_KERNEL);

	// initialize session data
	state->recv_timeout = 0;
	state->send_timeout = 0;
	state->m_timer.state = state;
	state->minor = get_minor(filp);
	atomic_set(&state->mtimer_size, 0);
	atomic_set(&state->valid, 1);
	INIT_LIST_HEAD(&state->m_timer.list);
	mutex_init(&state->writes_synchronizer);

	// access session state quickly from filp
	filp->private_data = (void *)state;

	// add session to list of active sessions
	session->session = state;
	mutex_lock(&the_object->session_list_synchronizer);
	list_add_tail(&session->list, &the_object->active_session.list);
	mutex_unlock(&the_object->session_list_synchronizer);

	AUDIT
	printk("%s: dev_open() for 0x%p with [major, minor] = [%d, %d] and session data at 0x%p\n",
	       MODNAME, filp, get_major(filp), get_minor(filp), state);
	return 0;
}

/**
 * dev_release()
 *
 * @inode: ptr to inode
 * @filp: ptr to I/O session.
 * 
 * closes I/O session. If there are no write in timeout that have yet
 * to be executed, it frees all the memory associated with the I/O
 * session being closed. Otherwise it signals that the session is no
 * longer valid, and the memory will be freed by the last delayed
 * write in hrtimer_callback().
 *
 */
static int dev_release(struct inode *inode, struct file *filp)
{
	session_state *state;

	state = (session_state *)filp->private_data;

	if (atomic_read(&state->mtimer_size) == 0) {
		// no delayed write to execute
		remove_session(state);
		kfree((session_state *)filp->private_data);

		AUDIT
		printk("%s: dev_release() for 0x%p with no delayed writes\n",
		       MODNAME, filp);

	} else {
		atomic_set(&state->valid, 0);

		AUDIT
		printk("%s: dev_release() for 0x%p with delayed writes\n",
		       MODNAME, filp);
	}

	return 0;
}

/**
 * dev_write()
 *
 * @filp: ptr to I/O session.
 * @buff: user space buff containting data to put in the queue
 * @len: size to write
 * @off: offset used
 * 
 * Write a single message to the device file instance's queue. Can
 * be controlled with dev_ioctl(). In particular if the option
 * IOCTL_SET_SEND_TIMEOUT was used in the current session it waits a
 * given timeout before actually writing the message.
 *
 */
static ssize_t dev_write(struct file *filp, const char *buff, size_t len,
			 loff_t *off)
{
	message *m;
	object_state *the_object;
	queue *q;
	session_state *state;
	struct mtimer *m_timer;
	ktime_t ktime_interval;

	the_object = &objects[get_minor(filp)];
	q = the_object->q;
	state = (session_state *)filp->private_data;

	AUDIT
	printk("%s: starting dev_write() for 0x%p\n", MODNAME, filp);

	if (len > MAX_MESSAGE_SIZE) {
		AUDIT
		printk("%s - [ERROR] dev_write() for 0x%p len %lu is too much\n",
		       MODNAME, filp, len);

		// file to large
		return -EFBIG;
	}

	if (atomic_read((atomic_t *)&size_so_far) + len > MAX_STORAGE_SIZE) {
		AUDIT
		printk("%s - [ERROR] dev_write() for 0x%p not enough space left on device\n",
		       MODNAME, filp);

		// no space on device
		return -ENOSPC;
	}

	atomic_set((atomic_t *)&size_so_far, size_so_far + len);

	// -- create new message from user buffer
	m = create_message(buff, len, 1);

	if (state->send_timeout == 0) {
		// store message in the queue immediately

		mutex_lock(&(the_object->queue_synchronizer));
		enqueue(q, m);
		mutex_unlock(&(the_object->queue_synchronizer));

		AUDIT
		printk("%s: dev_write() for 0x%p, write done, total size of driver: %d\n",
		       MODNAME, filp, size_so_far);

		// if there are readers waiting, wake them up
		wake_readers(the_object);

	} else {
		// store message after timeout

		// -- start high-res timer
		m_timer = kmalloc(sizeof(struct mtimer), GFP_KERNEL);
		m_timer->m = m;
		m_timer->state = state;
		m_timer->valid = 1;

		ktime_interval = ktime_set(0, state->send_timeout);

		hrtimer_init(&(m_timer->hr_timer), CLOCK_MONOTONIC,
			     HRTIMER_MODE_REL);
		m_timer->hr_timer.function = &hrtimer_callback;

		// add timer to list of timers
		mutex_lock(&(state->writes_synchronizer));
		list_add_tail(&(m_timer->list), &(state->m_timer.list));
		atomic_inc(&state->mtimer_size);
		mutex_unlock(&(state->writes_synchronizer));

		AUDIT
		printk("%s: dev_write() for 0x%p, starting hres timer with %lu timeout\n",
		       MODNAME, filp, state->send_timeout);

		// start timer
		hrtimer_start(&(m_timer->hr_timer), ktime_interval,
			      HRTIMER_MODE_REL);
	}

	// Always return as if we wrote all the bytes.
	return len;
}

/**
 * dev_read()
 *
 * @filp: ptr to I/O session.
 * @buff: user space buff in which to put the data read
 * @len: size to read
 * @off: offset used
 * 
 * Reads a single message from the device file instance's queue. Can
 * be controlled with dev_ioctl(). In particular if the option
 * IOCTL_SET_RECV_TIMEOUT was used in the current session it waits a
 * given timeout upon a read on an empty queue. During the timeout it
 * can be woken up by a thread that wrote on the queue.
 *
 */
static ssize_t dev_read(struct file *filp, char *buff, size_t len, loff_t *off)
{
	int ret;
	message *m;
	object_state *the_object;
	queue *q;
	session_state *state;
	ktime_t ktime_interval;
	struct reader *reader;

	the_object = &objects[get_minor(filp)];
	q = the_object->q;
	state = (session_state *)filp->private_data;

	AUDIT
	printk("%s: starting dev_read() for 0x%p\n", MODNAME, filp);

	// -- lock the queue
	mutex_lock(&(the_object->queue_synchronizer));

	if (q->len == 0 && state->recv_timeout == 0) {
		// if no message and no rcv_timeout, return immediately
		mutex_unlock(&(the_object)->queue_synchronizer);

		AUDIT
		printk("%s: dev_read() for 0x%p found empty queue", MODNAME,
		       filp);

		return 0;

	} else if (q->len == 0 && state->recv_timeout > 0) {
		// go to sleep in a wait queue
		mutex_unlock(&(the_object->queue_synchronizer));

		// add session to list of read I/O session on timeout
		reader = kmalloc(sizeof(struct reader), GFP_KERNEL);
		reader->valid = 1;
		reader->id = get_random_int();

		mutex_lock(&(the_object->readers_synchronizer));
		list_add_tail(&(reader->list), &(the_object->reader.list));
		mutex_unlock(&(the_object->readers_synchronizer));

		ktime_interval = ktime_set(0, state->recv_timeout);

		AUDIT
		printk("%s: dev_read() for 0x%p starting sleep timeout of %lu\n",
		       MODNAME, filp, state->recv_timeout);

		// wait until we're next reader, device is being
		// flushed, or timeout is over.
		wait_event_hrtimeout(the_object->wq,
				     the_object->next_reader_id == reader->id ||
					     the_object->flushing == 1,
				     ktime_interval);

		AUDIT
		printk("%s: dev_read() for 0x%p (0x%p) ended sleep with next-reader %ud\n",
		       MODNAME, filp, reader, the_object->next_reader_id);

		// remove from list and free mem
		mutex_lock(&(the_object->readers_synchronizer));
		list_del(&(reader->list));
		mutex_unlock(&(the_object->readers_synchronizer));

		kfree(reader);

		if (atomic_read((atomic_t *)&(the_object->flushing)) == 1) {
			// immediately return if device is being flushed

			AUDIT
			printk("%s: dev_read() for 0x%p, device being flushed\n",
			       MODNAME, filp);

			return 0;
		}

		mutex_lock(&(the_object->queue_synchronizer));
		if (q->len == 0) {
			mutex_unlock(&(the_object->queue_synchronizer));

			AUDIT
			printk("%s: dev_read() for 0x%p found empty queue",
			       MODNAME, filp);

			return 0;
		}
	}

	// extract the message and return it to the user
	m = dequeue(q);
	mutex_unlock(&(the_object->queue_synchronizer));

	len = len > m->length ? m->length : len;
	ret = copy_to_user(buff, m->data, len);

	// free message mem
	destroy_message(m);

	// update device's size
	atomic_set((atomic_t *)&size_so_far, size_so_far - len);

	AUDIT
	printk("%s: dev_read() for 0x%p found message of size %lu\n", MODNAME,
	       filp, len);

	return len;
}

/**
 * dev_ioctl()
 *
 * @filp: ptr to I/O session.
 * @ioctl_num: specifies which command to execute.
 * @param: extra
 * 
 * Allows to control the behavior of the device for the current I/O
 * session. The various commands supported are defined in
 * './include/chardev.h'.
 *
 */
static long dev_ioctl(struct file *filp, unsigned int ioctl_num,
		      unsigned long param)
{
	session_state *state;
	struct list_head *head, *pos;
	struct mtimer *m_timer;

	// if filp is NULL then we're being called from
	// cleanup_module()
	state = filp ? (session_state *)filp->private_data :
			     (session_state *)param;

	switch (ioctl_num) {
	case IOCTL_SET_SEND_TIMEOUT:
		// do not directly store messages upon write() but wait a timeout
		state->send_timeout = param;

		AUDIT
		printk("%s: dev_ioctl() for 0x%p with IOCTL_SET_SEND_TIMEOUT of %lu\n",
		       MODNAME, filp, param);

		break;

	case IOCTL_SET_RECV_TIMEOUT:
		// do not directly return if read on empty queue but wait a timeout
		state->recv_timeout = param;

		AUDIT
		printk("%s: dev_ioctl() for 0x%p with IOCTL_SET_RECV_TIMEOUT of %lu\n",
		       MODNAME, filp, param);

		break;

	case IOCTL_REVOKE_DELAYED_MESSAGES:
		// revoke all write() currently in timeout

		AUDIT
		printk("%s: dev_ioctl() for 0x%p with IOCTL_REVOKE_DELAYED_MESSAGES\n",
		       MODNAME, filp);

		head = &(state->m_timer.list);

		// iterate list and invalid all messages
		mutex_lock(&(state->writes_synchronizer));
		list_for_each (pos, head) {
			m_timer = list_entry(pos, struct mtimer, list);
			m_timer->valid = 0;
		}
		mutex_unlock(&(state->writes_synchronizer));

		break;

	default:
		AUDIT
		printk("%s: dev_ioctl() for 0x%p with unknown command\n",
		       MODNAME, filp);

		break;
	}

	return 0;
}

/**
 * dev_flush()
 *
 * @filp: ptr to I/O session.
 * @id: extra
 *
 * If called with filp != NULL flushes the state of a single instance
 * of the device file. Otherwise flushes the state of all instances of
 * the device file.
 *
 */
static int dev_flush(struct file *filp, fl_owner_t id)
{
	int ret, total_size;
	object_state *the_object;
	struct active_session *session;
	struct list_head *head, *pos;

	// check if dev_flush() is being called by cleanup_module()
	the_object = filp ? &objects[get_minor(filp)] : &objects[(int)id];

	// make sure dev_flush() can only be executed by a single
	// thread at any given time for any given device file instance.
	ret = __sync_bool_compare_and_swap(&(the_object->flushing), 0, 1);
	if (!ret) {
		AUDIT
		printk("%s: dev_flush() for 0x%p, another flush() executing for minor %d\n",
		       MODNAME, filp, filp ? get_minor(filp) : (int)id);

		return 0;
	}

	AUDIT
	printk("%s: dev_flush() for 0x%p, starting execution for minor %d\n",
	       MODNAME, filp, filp ? get_minor(filp) : (int)id);

	// lock all resources
	mutex_lock(&the_object->session_list_synchronizer);
	mutex_lock(&the_object->readers_synchronizer);
	mutex_lock(&the_object->queue_synchronizer);

	// revoke all delayed messages for all I/O sessions
	head = &the_object->active_session.list;
	list_for_each (pos, head) {
		session = list_entry(pos, struct active_session, list);

		// NOTE: instead of using filp we pass directly the
		// session as the third argument
		dev_ioctl(NULL, IOCTL_REVOKE_DELAYED_MESSAGES,
			  (unsigned long)session->session);
	}

	// revoke all messages
	total_size = the_object->q->total_size;
	flush_queue(the_object->q);

	// update device's size
	atomic_set((atomic_t *)&size_so_far, size_so_far - total_size);

	// wake up all readers
	wake_up(&(the_object->wq));

	// unlock all resources
	mutex_unlock(&(the_object->session_list_synchronizer));
	mutex_unlock(&(the_object->readers_synchronizer));
	mutex_unlock(&(the_object->queue_synchronizer));

	atomic_set((atomic_t *)&the_object->flushing, 0);

	return 0;
}

// ------------------------------
// Device driver

static struct file_operations fops = {
	.owner = THIS_MODULE,
	.write = dev_write,
	.read = dev_read,
	.open = dev_open,
	.release = dev_release,
	.unlocked_ioctl = dev_ioctl,
	// .flush = dev_flush,
};

// ------------------------------
// Module functions

int init_module(void)
{
	int i, ret;

	ret = __register_chrdev(DEVICE_MAJOR, 0, 256, DEVICE_NAME, &fops);

	if (ret < 0) {
		printk("%s - [ERROR]: registering device failed\n", MODNAME);
		return ret;
	}

	AUDIT
	printk("%s: Device registered with major number %d\n", MODNAME,
	       DEVICE_MAJOR);

	// init driver state
	size_so_far = 0;
	for (i = 0; i < DEVICE_MINORS; i++) {
		mutex_init(&(objects[i].queue_synchronizer));
		objects[i].q = init_queue();

		if (!objects[i].q) {
			printk("%s - [ERROR]: queue n. %d not allocated. \n",
			       MODNAME, i);
			return -1;
		}

		AUDIT
		printk("%s: Queue n. %d allocated at address 0x%p\n", MODNAME,
		       i, objects[i].q);

		mutex_init(&(objects[i].session_list_synchronizer));
		INIT_LIST_HEAD(&(objects[i].active_session.list));

		init_waitqueue_head(&(objects[i].wq));

		mutex_init(&(objects[i].readers_synchronizer));
		INIT_LIST_HEAD(&(objects[i].reader.list));

		objects[i].flushing = 0;
	}

	return 0;
}

void cleanup_module(void)
{
	int i;

	// free driver memory
	for (i = 0; i < DEVICE_MINORS; i++) {
		dev_flush(NULL, (fl_owner_t)i);
		kfree(objects[i].q);
	}

	unregister_chrdev(DEVICE_MAJOR, DEVICE_NAME);

	AUDIT
	printk("%s: Module unloaded!\n", MODNAME);
	return;
}