#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/interrupt.h>
#include <linux/proc_fs.h>
#include <linux/sched.h>
#include <linux/keyboard.h>
#include <asm/uaccess.h>
#include <asm/io.h>
#include <linux/slab.h>

#define PROC_FILE_NAME "kbd_demo"

void* lookup_node(struct rb_root *root, void *search_target, int cmp(const void*, const void*), size_t node_offset){
	// if not found, return 0
	// if found, return the address of the item
	struct rb_node **spot = &(root->rb_node), *parent = NULL;
	while(*spot) {
		void *spot_data = (*spot) /* Yields a struct rb_node* */ + node_offset;  
		parent = *spot;
		if(!cmp(spot_data, search_target)){ // We found it!
			return spot_data;
		} else if(cmp(spot_data, search_target) > 0)
			spot = &((*spot)->rb_right);
		else
			spot = &((*spot)->rb_left);
	}
	return 0;	
}

/* Returns 0 if the data was already there.  Node is not inserted in this case */
void* insert_node(struct rb_root *root, void *data, int cmp(const void*, const void*), size_t datasize, size_t node_offset){
	struct rb_node **spot = &(root->rb_node), *parent = NULL;
	void *new_node;
	struct rb_node* node;
	while(*spot) {
		void *spot_data = (char*)(*spot) + node_offset;  
		parent = *spot;
		if(!cmp(spot_data, data)){ // We found it!
			return 0;
		} else if(cmp(spot_data, data) > 0)
			spot = &((*spot)->rb_right);
		else
			spot = &((*spot)->rb_left);
	}

	// This next bit adds a new node.  We'll have exited above if we found our node already in the tree
	new_node = kmalloc(datasize, 1);
	memcpy(new_node, data, datasize);
	node = new_node + node_offset;
	rb_link_node(node, parent, spot);
	rb_insert_color(node, root);
	return new_node;
}

struct keycount_node {
	struct rb_node node; // 24 bytes, on x86-64, 0 bytes into the struct
	int keycode;	// 4 bytes, 24 bytes into the struct
	unsigned int keycount; // 4 bytes, 28 bytes into the struct
};

int cmp_keycount_node(const void* a, const void* b){
	const struct keycount_node *an = a;
	const struct keycount_node *bn = b;
	printk("Comparing nodes for keycodes %d and %d\n", an->keycode, bn->keycode);
	if(an->keycode > bn->keycode)
		return 1;
	else if(an->keycode == bn->keycode)
		return 0;
	return -1;
}

struct rb_root root = RB_ROOT;

struct notifier_block nb;

ssize_t read_simple(struct file *filp,char *buf,size_t count,loff_t *offp ) 
{
	struct keycount_node *pos, *n;
	size_t sofar = 0;
	if(*offp)
		return 0;
	// kinda like for(auto *pos : rb_root) in C++11
	rbtree_postorder_for_each_entry_safe(pos, n, &root, node){
		sofar += snprintf(buf + sofar, count - sofar, "Keycode %d : %u\n", pos->keycode, pos->keycount);
	}
	*offp = sofar;
	return sofar;
}

struct proc_ops proc_fops = {
	proc_read: read_simple,
};

int kb_notifier_fn(struct notifier_block *nb, unsigned long action, void* data){
	struct keyboard_notifier_param *kp = (struct keyboard_notifier_param*)data;
	struct keycount_node comparison = {
		.keycode = kp->value,
		.keycount = 1
	};
	struct keycount_node *tree_entry = lookup_node(&root, &comparison, cmp_keycount_node, 0);
	if(tree_entry){ // It was found!
		tree_entry->keycount++;
		return 0;
	}
	insert_node(&root, &comparison, cmp_keycount_node, sizeof(struct keycount_node), 0);
	return 0;
}

/*
int kb_notifier_fn(struct notifier_block *nb, unsigned long action, void* data, size_t datasize, ){
	struct keyboard_notifier_param *kp = (struct keyboard_notifier_param*)data;
	// search and insert if needed
	struct rb_node **new = &(root.rb_node), *parent = NULL;
	while(*new) {
		struct keycount_node *this = container_of(*new, struct keycount_node, node);  
		parent = *new;
		if(this->keycode == kp->value){ // We found it!
			this->keycount++;
			printk("Added to %d:%u\n", this->keycode, this->keycount);
			return 0;
		} else if(this->keycode < kp->value)
			new = &((*new)->rb_right);
		else
			new = &((*new)->rb_left);
	}
	{
  	new_node = kmalloc(sizeof(struct keycount_node), 1);
	new_node->keycode = kp->value;
	new_node->keycount = 1;
	rb_link_node(&new_node->node, parent, new);
	rb_insert_color(&new_node->node, &root);
	printk("Added new node %d:%u\n", new_node->keycode, new_node->keycount);
	}
	return 0;
}
*/
int init (void) {
	nb.notifier_call = kb_notifier_fn;
	register_keyboard_notifier(&nb);
	proc_create(PROC_FILE_NAME,0,NULL,&proc_fops);
	return 0;
}

void cleanup(void) {
	struct keycount_node *pos, *n;

	unregister_keyboard_notifier(&nb);
	remove_proc_entry(PROC_FILE_NAME,NULL);
	
	printk("Beginning Tree Deallocation\n");
	rbtree_postorder_for_each_entry_safe(pos, n, &root, node){
		printk("Keycode %d : %u\n", pos->keycode, pos->keycount);
		kfree(pos);
	}
}

MODULE_LICENSE("GPL"); 
module_init(init);
module_exit(cleanup);