#include <sys/types.h>
#include <sys/module.h>
#include <sys/systm.h>  /* uprintf */
#include <sys/param.h>  /* defines used in kernel.h */
#include <sys/kernel.h> /* types used in module initialization */
#include <sys/conf.h>   /* cdevsw struct */
#include <sys/uio.h>    /* uio struct */
#include <sys/malloc.h>
#include <sys/kthread.h>
#include <sys/lockdoc.h>
#include <sys/mutex.h>
#include <sys/proc.h>
#include <sys/lockmgr.h>
#include "lockdoc_internal.h"
#include "opt_lockdoc.h"

#ifndef LOCKDOC_TEST_ITERATIONS
#define LOCKDOC_TEST_ITERATIONS	200
#endif

#define DEV_FILE_CONTROL_NAME	"control"
#define DEV_FILE_ITER_NAME		"iterations"
#define DEFAULT_ITERATIONS	LOCKDOC_TEST_ITERATIONS

#define MS_TO_HZ(x)		((x) * hz / 1000)
/*
 * For some reasons, our ring buffer (aka BSB ring buffer)
 * can only hold size - 1 elements.
 * If we want to store DEFAULT_ITERATIONS elements, as desired,
 * the buffer must be one element larger.
 * Hence, RING_BUFFER_SIZE_REAL is used for allocating the actual buffer
 * and used for the size member.
 * In contrast, RING_BUFFER_SIZE_VIRT is used when asigning a new value
 * for iterations in procfile_iter_write.
 */
#define RING_BUFFER_SIZE_REAL	(DEFAULT_ITERATIONS + 1)
#define RING_BUFFER_SIZE_VIRT	(RING_BUFFER_SIZE_REAL - 1)
#define RING_BUFFER_STORAGE_TYPE int
#define MK_STRING(x)    #x
#define START_AND_WAIT_THREAD(x)	start_and_wait_thread(MK_STRING(x), x)

// Shamelessly stolen from sys/buf.h:310
#define LOCKDOC_LOCKINIT(x, y)						\
	lockinit(&(x), PRIBIO + 4, (y), 0, LK_NEW)

#define	LOCKDOC_LOCK(x, locktype, interlock)				\
	_lockmgr_args_rw(&(x), (locktype), (interlock),	\
	    LK_WMESG_DEFAULT, LK_PRIO_DEFAULT, LK_TIMO_DEFAULT,		\
	    LOCK_FILE, LOCK_LINE)

#define	LOCKDOC_UNLOCK(x) do {						\
	(void)_lockmgr_args(&(x), LK_RELEASE, NULL,		\
	    LK_WMESG_DEFAULT, LK_PRIO_DEFAULT, LK_TIMO_DEFAULT,		\
	    LOCK_FILE, LOCK_LINE);					\
} while (0)

static d_write_t     lockdoc_ctl_write;

static d_read_t      lockdoc_iter_read;
static d_write_t     lockdoc_iter_write;

static struct cdevsw ctl_cdevsw = {
	.d_version = D_VERSION,
	.d_write = lockdoc_ctl_write,
	.d_name = "lockdoc_control",
};
static struct cdevsw iter_cdevsw = {
	.d_version = D_VERSION,
	.d_read = lockdoc_iter_read,
	.d_write = lockdoc_iter_write,
	.d_name = "lockdoc_iterations",
};

static struct cdev *ctl_dev;
static struct cdev *iter_dev;


static int iterations = DEFAULT_ITERATIONS;
static struct proc *control_thread = NULL;

static struct mtx rb_lock;
static struct mtx consumer_lock;
static struct mtx producer_lock;
static struct mtx counter_mtx;
static struct lock counter_lock;

static MALLOC_DEFINE(M_LOCKDOC, "lockdoc_ring_buffer", "ring buffer for LockDoc test");
struct lockdoc_ring_buffer {
	int next_in;
	int next_out;
	int size;
	RING_BUFFER_STORAGE_TYPE data[RING_BUFFER_SIZE_REAL];
	int counter;
};
static struct lockdoc_ring_buffer *ring_buffer = NULL;

static __noinline int is_full(volatile struct lockdoc_ring_buffer *buffer) { return (buffer->next_in + 1) % buffer->size == buffer->next_out; }
static __noinline int is_empty(volatile struct lockdoc_ring_buffer *buffer) { return buffer->next_out == buffer->next_in; }

static __noinline int produce(volatile struct lockdoc_ring_buffer *buffer, RING_BUFFER_STORAGE_TYPE data) {
	if (is_full(buffer)) {
		return -1;
	}

	buffer->data[buffer->next_in] = data;
	buffer->next_in = (buffer->next_in + 1) % buffer->size;

	return 0;
}

static __noinline RING_BUFFER_STORAGE_TYPE consume(volatile struct lockdoc_ring_buffer *buffer) {
	RING_BUFFER_STORAGE_TYPE result;

	if (is_empty(buffer)) {
		return -1;
	}
	result = buffer->data[buffer->next_out];
	buffer->next_out = (buffer->next_out + 1) % buffer->size;

	return result;
}

static void producer_thread_work(void *data) {
	int i, ret;

	/*
	 * Produce 'iterations' elements.
	 * This fills every element in the ring buffer.
	 * The 'iterations'+1 call to produce() would fail due to a full buffer.
	 * The consumer thread will completely empty the buffer.
	 */
	for (i = 0; i < iterations; i++) {
		mtx_lock(&rb_lock);
		ret = is_full(ring_buffer);
		mtx_unlock(&rb_lock);
		if (ret) {
			printf("%s: Ring buffer is full\n", __func__);
		}

		mtx_lock(&producer_lock);
		mtx_lock(&rb_lock);
		ret = produce(ring_buffer, i + 30);
		mtx_unlock(&producer_lock);
		mtx_unlock(&rb_lock);
		printf("%s-%03d: Produced(%d): %03d\n", __func__, i, ret, i + 30);

		pause("W", MS_TO_HZ(100));
	}

	kproc_exit(0);
}

static void consumer_thread_work(void *data) {
	int i, ret;

	for (i = 0; i < iterations; i++) {
		mtx_lock(&rb_lock);
		ret = is_empty(ring_buffer);
		mtx_unlock(&rb_lock);
		if (ret) {
			printf("%s: Ring buffer is empty\n", __func__);
		}

		mtx_lock(&consumer_lock);
		mtx_lock(&rb_lock);
		ret = consume(ring_buffer);
		mtx_unlock(&consumer_lock);
		mtx_unlock(&rb_lock);
		printf("%s-%03d: Consumed: %03d\n", __func__, i, ret);

		pause("W", MS_TO_HZ(100));
	}

	kproc_exit(0);
}

static void dirty_nolocks_thread_work(void *data) {
	int i = 0, ret;

	/*
	 * Wait a bit. Otherwise the call to tsleep() in control_thread_work() will wait for ever,
	 * because this thread has terminated and the caller will not be notified.
	 */
	pause("W", MS_TO_HZ(500));

	ret = is_full(ring_buffer);
	if (ret) {
		printf("%s: Ring buffer is full\n", __func__);
	}

	ret = produce(ring_buffer, i - 1);
	printf("%s-%03d: Produced(%d): %03d\n", __func__, i, ret, i - 1);
	
	ret = is_empty(ring_buffer);
	if (ret) {
		printf("%s: Ring buffer is empty\n",__func__);
	}

	ret = consume(ring_buffer);
	printf("%s-%03d: Consumed: %03d\n", __func__, i, ret);

	kproc_exit(0);
}

static void dirty_fewlocks_thread_work(void *data) {
	int i = 0, ret;

	/*
	 * Wait a bit. Otherwise the call to tsleep() in control_thread_work() will wait for ever,
	 * because this thread has terminated and the caller will not be notified.
	 */
	pause("W", MS_TO_HZ(500));

	mtx_lock(&rb_lock);
	ret = is_full(ring_buffer);
	mtx_unlock(&rb_lock);
	if (ret) {
		printf("%s: Ring buffer is full\n", __func__);
	}

	mtx_lock(&rb_lock);
	ret = produce(ring_buffer, i - 1);
	mtx_unlock(&rb_lock);
	printf("%s-%03d: Produced(%d): %03d\n", __func__, i, ret, i - 1);

	mtx_lock(&rb_lock);
	ret = is_empty(ring_buffer);
	mtx_unlock(&rb_lock);
	if (ret) {
		printf("%s: Ring buffer is empty\n", __func__);
	}

	mtx_lock(&rb_lock);
	ret = consume(ring_buffer);
	mtx_unlock(&rb_lock);
	printf("%s-%03d: Consumed: %03d\n", __func__, i, ret);

	kproc_exit(0);
}

static void dirty_alllocks_thread_work(void *data) {
	int i = 0, ret;

	/*
	 * Wait a bit. Otherwise the call to tsleep() in control_thread_work() will wait for ever,
	 * because this thread has terminated and the caller will not be notified.
	 */
	pause("W", MS_TO_HZ(500));

	mtx_lock(&producer_lock);
	mtx_lock(&consumer_lock);
	mtx_lock(&rb_lock);
	ret = is_full(ring_buffer);
	mtx_unlock(&rb_lock);
	mtx_unlock(&consumer_lock);
	mtx_unlock(&producer_lock);
	if (ret) {
		printf("%s: Ring buffer is full\n", __func__);
	}

	mtx_lock(&producer_lock);
	mtx_lock(&consumer_lock);
	mtx_lock(&rb_lock);
	ret = produce(ring_buffer, i - 1);
	mtx_unlock(&rb_lock);
	mtx_unlock(&consumer_lock);
	mtx_unlock(&producer_lock);
	printf("%s-%03d: Produced(%d): %03d\n", __func__, i, ret, i - 1);

	mtx_lock(&producer_lock);
	mtx_lock(&consumer_lock);
	mtx_lock(&rb_lock);
	ret = is_empty(ring_buffer);
	mtx_unlock(&rb_lock);
	mtx_unlock(&consumer_lock);
	mtx_unlock(&producer_lock);
	if (ret) {
		printf("%s: Ring buffer is empty\n", __func__);
	}

	mtx_lock(&producer_lock);
	mtx_lock(&consumer_lock);
	mtx_lock(&rb_lock);
	ret = consume(ring_buffer);
	mtx_unlock(&rb_lock);
	mtx_unlock(&consumer_lock);
	mtx_unlock(&producer_lock);
	printf("%s-%03d: Consumed: %03d\n", __func__, i, ret);

	kproc_exit(0);
}

static void dirty_order_thread_work(void *data) {
	int i = 0, ret;

	/*
	 * Wait a bit. Otherwise the call to tsleep() in control_thread_work() will wait for ever,
	 * because this thread has terminated and the caller will not be notified.
	 */
	pause("W", MS_TO_HZ(500));

	mtx_lock(&rb_lock);
	mtx_lock(&producer_lock);
	ret = produce(ring_buffer, i - 1);
	mtx_unlock(&producer_lock);
	mtx_unlock(&rb_lock);
	printf("%s-%03d: Produced(%d): %03d\n", __func__, i, ret, i - 1);

	mtx_lock(&rb_lock);
	mtx_lock(&consumer_lock);
	ret = consume(ring_buffer);
	mtx_unlock(&consumer_lock);
	mtx_unlock(&rb_lock);
	printf("%s-%03d: Consumed: %03d\n", __func__, i, ret);

	kproc_exit(0);
}

static void fake_txn_start_thread_work(void *data) {
	/*
	 * Wait a bit. Otherwise the call to tsleep() in control_thread_work() will wait for ever,
	 * because this thread has terminated and the caller will not be notified.
	 */
	pause("W", MS_TO_HZ(500));


	LOCKDOC_LOCK(counter_lock, LK_EXCLUSIVE | LK_NOWAIT, NULL);
	ring_buffer->counter++;	
	printf("%s-%04d-%04d: Incremented counter\n", __func__, curthread->td_tid, curthread->td_proc->p_pid);
	_lockmgr_disown(&counter_lock, LOCK_FILE, LOCK_LINE);

	kproc_exit(0);
}

static void fake_txn_end_thread_work_1(void *data) {
	/*
	 * Wait a bit. Otherwise the call to tsleep() in control_thread_work() will wait for ever,
	 * because this thread has terminated and the caller will not be notified.
	 */
	pause("W", MS_TO_HZ(500));

	ring_buffer->counter++;	
	printf("%s-%04d-%04d: Incremented counter\n", __func__, curthread->td_tid, curthread->td_proc->p_pid);
	LOCKDOC_UNLOCK(counter_lock);

	kproc_exit(0);
}

static void fake_txn_end_thread_work_2(void *data) {
	/*
	 * Wait a bit. Otherwise the call to tsleep() in control_thread_work() will wait for ever,
	 * because this thread has terminated and the caller will not be notified.
	 */
	pause("W", MS_TO_HZ(500));

	mtx_lock(&counter_mtx);
	ring_buffer->counter++;	
	mtx_unlock(&counter_mtx);
	printf("%s-%04d-%04d: Incremented counter\n", __func__, curthread->td_tid, curthread->td_proc->p_pid);
	LOCKDOC_UNLOCK(counter_lock);

	kproc_exit(0);
}

static void start_and_wait_thread(const char *fn_name, void (*work_fn)(void*)) {
	struct proc *temp = NULL;
	int error;
	

	printf("%s: Starting %s thread...\n", __func__, fn_name);	
	error = kproc_create(work_fn, NULL, &temp,0, 0, "lockdoc-control");
	if (error) {
		return;
	}
	printf("%s: Waiting for %s thread (%d) to terminate...\n", __func__, fn_name, temp->p_pid);
	error = tsleep(temp, 0, "ldctl", 0);
	if (error) {
		printf("%s: Error waiting for %s thread\n", __func__, fn_name);
	} else {
		printf("%s: %s thread terminated successfully\n", __func__, fn_name);
	}
}

static void control_thread_work(void *data) {

	mtx_init(&rb_lock, "LockDoc test rb lock", NULL, MTX_DEF);
	mtx_init(&consumer_lock, "LockDoc test consumer lock", NULL, MTX_DEF);
	mtx_init(&producer_lock, "LockDoc test producer lock", NULL, MTX_DEF);
	mtx_init(&counter_mtx, "LockDoc test counter mtx", NULL, MTX_DEF);
	LOCKDOC_LOCKINIT(counter_lock, "LockDoc test counter lock");

	ring_buffer = malloc(sizeof(*ring_buffer), M_LOCKDOC, M_WAITOK | M_ZERO);
	if (!ring_buffer) {
		printf("Cannot allocate %u bytes for ring buffer\n", sizeof(*ring_buffer));
		kproc_exit(1);
	}
	ring_buffer->size = RING_BUFFER_SIZE_REAL;
	log_memory(1, "lockdoc_ring_buffer", ring_buffer, sizeof(*ring_buffer));

	START_AND_WAIT_THREAD(producer_thread_work);
	START_AND_WAIT_THREAD(consumer_thread_work);
	START_AND_WAIT_THREAD(dirty_nolocks_thread_work);
	START_AND_WAIT_THREAD(dirty_fewlocks_thread_work);
	START_AND_WAIT_THREAD(dirty_alllocks_thread_work);
	START_AND_WAIT_THREAD(dirty_order_thread_work);
	START_AND_WAIT_THREAD(fake_txn_start_thread_work);
	START_AND_WAIT_THREAD(fake_txn_end_thread_work_1);
	START_AND_WAIT_THREAD(fake_txn_start_thread_work);
	START_AND_WAIT_THREAD(fake_txn_end_thread_work_2);

	log_memory(0, "lockdoc_ring_buffer", ring_buffer, sizeof(*ring_buffer));
	free(ring_buffer, M_LOCKDOC);

	mtx_destroy(&rb_lock);
	mtx_destroy(&producer_lock);
	mtx_destroy(&consumer_lock);

	kproc_exit(0);
}

#define BUFSIZE	5
static int lockdoc_ctl_write(struct cdev *dev __unused, struct uio *uio, int ioflag __unused) {

	unsigned long value = 0;
	size_t amt;
	int error;
	char buffer[BUFSIZE];

	if (uio->uio_offset != 0 && (uio->uio_offset != BUFSIZE)) {
		return (EINVAL);
	}

	amt = MIN(uio->uio_resid, BUFSIZE);
	error = uiomove(buffer, amt, uio);
	if (error) {
		return (error);
	}

	/* parse input */
	value = strtoul(buffer, NULL, 10);
	if (value == 1) {
		error = kproc_create(control_thread_work, NULL, &control_thread,0, 0, "lockdoc-control");
		if (error) {
			return (error);
		}

		pause("W", MS_TO_HZ(200));
		uprintf("%s: Waiting for control_thread to terminate...\n", __func__);
		// This will block the caller until all threads terminated
		error = tsleep(control_thread, 0, "ldctl", 0);
		if (error) {
			uprintf("%s: Wait for control thread timed out\n", __func__);
			return (error);
		}
		uprintf("%s: Control thread terminated successfully!\n", __func__);
	} else {
		return (EINVAL);
	}

	return 0;
}
#undef BUFSIZE

#define BUFSIZE 20
static int lockdoc_iter_read(struct cdev *dev __unused, struct uio *uio, int ioflag __unused)
{
	size_t amt;
	int error;
	char buffer[BUFSIZE];
	int ret = snprintf(buffer, BUFSIZE, "%u\n", iterations);

	if (ret >= BUFSIZE) {
		ret = BUFSIZE;
	}

	amt = MIN(uio->uio_resid, uio->uio_offset >= ret ? 0 :
	    ret - uio->uio_offset);

	error = uiomove(buffer, amt, uio);
	if (error != 0) {
		uprintf("uiomove failed!\n");
	}

	return (error);
}
#undef BUFSIZE

#define BUFSIZE 20
static int lockdoc_iter_write(struct cdev *dev __unused, struct uio *uio, int ioflag __unused) {
	unsigned long value = 0;
	size_t amt;
	int error;
	char buffer[BUFSIZE];

	if (uio->uio_offset != 0 && (uio->uio_offset != BUFSIZE)) {
		return (EINVAL);
	}

	amt = MIN(uio->uio_resid, BUFSIZE);
	error = uiomove(buffer, amt, uio);
	if (error) {
		return (error);
	}

	/* parse input */
	value = strtoul(buffer, NULL, 10);
	/*
	 * Iterations cannot be larger than the buffer size.
	 * If 'iterations' is larger than the buffer size, 
	 * the consumer/producer thread will execute different code paths (iterations - RING_BUFFER_SIZE_VIRT) times.
	 * We want to the consumer and producer thread to execute the same code 'iterations' times.
	 */
	if (value > RING_BUFFER_SIZE_VIRT) {
		uprintf("%s: Desired iterations (%lu) is larger than buffer size (%d)\n", __func__, value, RING_BUFFER_SIZE_VIRT); 
		return (EINVAL);
	}
	iterations = value;
	uprintf("Setting iterations to %d\n", iterations);

	return 0;
}
#undef BUFSIZE

static int lockdoc_test_loader(struct module *m __unused, int what, void *arg __unused)
{
	int error = 0;

	switch (what) {
	case MOD_LOAD:
		error = make_dev_p(MAKEDEV_CHECKNAME | MAKEDEV_WAITOK,
		    &ctl_dev,
		    &ctl_cdevsw,
		    0,
		    UID_ROOT,
		    GID_WHEEL,
		    0300,
		    DEV_DIR_NAME "/" DEV_FILE_CONTROL_NAME);
		if (error != 0) {
			printf("Could not create /dev/%s/%s\n", DEV_DIR_NAME, DEV_FILE_CONTROL_NAME);
			break;
		}
		error = make_dev_p(MAKEDEV_CHECKNAME | MAKEDEV_WAITOK,
		    &iter_dev,
		    &iter_cdevsw,
		    0,
		    UID_ROOT,
		    GID_WHEEL,
		    0600,
		    DEV_DIR_NAME "/" DEV_FILE_ITER_NAME);
		if (error != 0) {
			printf("Could not create /dev/%s/%s\n", DEV_DIR_NAME, DEV_FILE_ITER_NAME);
			destroy_dev(ctl_dev);
			break;
		}
		break;
	case MOD_UNLOAD:
		destroy_dev(ctl_dev);
		destroy_dev(iter_dev);
		break;
	default:
		error = EOPNOTSUPP;
		break;
	}
	return (error);
}


DEV_MODULE(lockdoc_test, lockdoc_test_loader, NULL);