/** 
 * Anemone Memory Server 
 **/ 
#include "server.h"
/* Server kernel module initialization */
int server_init(void) {
	int ret;

	if((ret = mem_initialize()) < 0)
		goto out;
	if((ret = net_init(anemone_server_recv)) < 0)
		goto out;

	proc_init();

	if(plog) 
		printk(KERN_ALERT "Logging internal processing times.\n");
	if(debug)
		printk(KERN_ALERT "Memory Server loaded.\n");
out:
	return ret;
}

/* Memory storage specific initilazations */
int mem_initialize(void) {
	int err;

	if(!(percent = (percent * totalram_pages / 100))) {
		printk(KERN_ALERT "Error: Memory percentage limit too small.\n");
		mem_print();
		return -ENOMEM;
	}

	if(percent != totalram_pages)
		printk(KERN_ALERT 
		    "Restricting self to maximum %lu / %lu pages\n", percent,
		    totalram_pages);

	block_cache = kmem_cache_create("block_cache", sizeof(Block),
				   	0, 0, NULL, NULL);

	if(!block_cache) {
		printk(KERN_ALERT 
			"Anemone error: Storage cache allocation failed.\n");
		kmem_cache_destroy(block_cache);
		return -ENOMEM;
	}

	/*
	 * initialize the mapping hashtable between pages and the servers they are stored on
	 */
	if((err = init_table(&table, &hash_size, pagekey_hash, 
			      pagekey_equal, "table"))) {
		kmem_cache_destroy(block_cache);
		return err;
	}

	binary_hash_size = table.ibits;

	mem_print();

	init_timer(&broadcast_timer);
	broadcast_timer.function = broadcast_timeout;
	broadcast_timer.expires = jiffies +  SERVER_BROADCAST_TIMEOUT;
	broadcast_timer.data = SERVER_BROADCAST_TIMEOUT;
	add_timer(&broadcast_timer);

	init_timer(&session_timer);
	session_timer.function = session_timeout;
	session_timer.expires = jiffies +  SESSION_TIMEOUT;
	session_timer.data = SESSION_TIMEOUT;
	add_timer(&session_timer);
	
	printk(KERN_ALERT "Anemone Server Started.\n");

	return 0;
}

/* /proc procfile initiliazations */
void proc_init(void) {
	/* Procfile for printing out useful information */
	if(plog) {
		aprocfile = create_proc_entry(PROCNAME, S_IFREG | S_IRUGO,NULL);
		aprocfile->write_proc 	= toggle_logging;
		aprocfile->read_proc 	= server_report;
		aprocfile->data 	= report_data;
		aprocfile->owner 	= THIS_MODULE;
	}

	/* Procfile for debug toggling */
	aprocfile2 = create_proc_entry(PROCNAME2, S_IFREG | S_IRUGO, NULL); 
	aprocfile2->write_proc	= toggle_debug;

	/* Query server, asking whether or not a specific block exists */
	aprocfile3 = create_proc_entry(PROCNAME3, S_IFREG | S_IRUGO, NULL); 
	aprocfile3->write_proc 	= query_block;

	/* common values */
	aprocfile2->data = aprocfile3->data = report_data;
	aprocfile2->owner = aprocfile3->owner = THIS_MODULE;
	aprocfile2->read_proc = aprocfile3->read_proc = server_report;
}

/* Incoming packets are delivered to this routing immediately when they're
 * handled by the NIC's device driver if the protocol ID matches 0x9000
 */
int anemone_server_recv(sk_buff * skb) {
	Message *sendmsg,*msg = (Message *) skb->data;
	PageKey	key;
	sk_buff	* reply = NULL;

	/*
	 * Make sure there were no hardware-related
	 * transmission errors.
	if(!valid_checksum(skb->data))
		goto failed;
	 */

	if(plog) 
		do_gettimeofday(&start); 


	memcpy(key.client_mac, skb->mac.raw + ETH_ALEN, ETH_ALEN);
	key.offset 	= msg->offset;
	key.vpid 	= msg->vpid;

	if(debug) {
		printk(KERN_ALERT 
			  "Message Received: offset %d, lower: %lu, count %d\n",
			                  msg->offset, key.offset, msg->count);
		printk(KERN_ALERT " Source: %02x:%02x:%02x:%02x:%02x:%02x\n",
			mac_ntoa(key.client_mac));
	}

	if(msg->count != PAGE_SIZE && debug)
		printk(KERN_ALERT 
		    " Warning: non-PAGE_SIZE request: %d bytes\n", msg->count);

	/*
	 * Respond to the request 
	 */
	switch(msg->type) {
		case TYPE_SERVER_BROADCAST:
			goto bad;
		case TYPE_CLIENT_BROADCAST:
			update_session_status(key.client_mac, msg->session_id);
			goto good;
		case PAGE_IN:
			reply = page_in(msg, &key, skb);
			reads++;
//			if ( net_ratelimit() ) {
			if ( MOD(reads, 17) == 0 ) {
				printk(KERN_ALERT " get page in request!!!\n");
			}
			break;
		case PAGE_OUT:
			reply = page_out(msg, &key, skb);
			writes++;
//			if ( net_ratelimit() ) {
			if ( MOD(writes, 17) == 0 ) {
				printk(KERN_ALERT " get page out request!!!\n");
			}
			break;
		default:
			printk(KERN_ALERT " Invalid Message received.\n");
			goto bad;
	}

	if(reply) {
		sendmsg = (Message *) (reply->data + LINK_HDR);
		fill_status_info(sendmsg);
		while (sendmsg->total_mem <= 0 ) {
			fill_status_info(sendmsg);
			printk(KERN_ALERT " Reply cannot get memory status right,redo getting memory status.\n");
		}
			
		anemone_send(reply, key.client_mac);
	} else
		if (debug) printk(KERN_ALERT " Reply not warranted.\n");
good:
	kfree_skb(skb);
	return NET_RX_SUCCESS;
bad:
	kfree_skb(skb);
	return NET_RX_DROP;
}	

/* Handle a Page-IN request from a client. Retreive from main memory send. */
sk_buff * page_in(Message * m, PageKey * key, sk_buff * orig) {
	u8		* addr;
	ulong		allocate = LINK_HDR + MSGSIZE;
	Block 	 	* block;
	Message  	* reply;
	sk_buff  	* skb;
	ethhdr 		* eth;

	if(debug) {
		printk(KERN_ALERT " Page-In. Retreiving Page from table...\n");
		print_key(key);
	}

	block = (Block *) lookup_table(&table, key);

	if(block)
		allocate += m->count;

	skb = dev_alloc_skb(allocate);

	/* 
	 * During high memory usage, I hope this doesn't happen.
	 * If it does, I'll have to deal with it.
	 */
	if(!skb)
		BUG();
	
	allocate 	-= m->count;
	eth 		= (ethhdr *) skb->data;
	reply 		= (Message *) (eth + 1);
	reply->status 	= REPLY_BAD;

	if(block) {
		addr = kmap_atomic(block->page, KM_SKB_DATA_SOFTIRQ);

		/* The subtraction here is always zero, I believe. Just get rid of it */
		memcpy(reply + 1, addr + (m->offset - key->offset), m->count);

		reply->status = REPLY_GOOD;

		allocate += m->count;

		kunmap_atomic(addr, KM_SKB_DATA_SOFTIRQ);
	} else
		printk(" Retrieve Failed. The page is not here: %lu.\n", (ulong) key->offset);
			
	reply->type 	= PAGE_IN_REPLY;
	reply->offset 	= m->offset;
	reply->vpid 	= m->vpid;
	reply->seq	= m->seq;
	reply->count 	= m->count;


	if(plog) {
		do_gettimeofday(&end);
		timersub(&end, &start, &diff);
		reply->serv_internal_time = USEC(diff);	
	}

	skb->len = allocate;

	return skb;
}

/* Handle a Page-OUT request from a client */
sk_buff * page_out(Message * m, PageKey * key, struct sk_buff * orig) {
	u8		* addr = NULL;
	Block 	 	* block = NULL;
	sk_buff  	* skb = NULL;
	Message  	* reply;
	struct page 	* pg;

	if(debug) {
		printk(KERN_ALERT " Page-Out. Checking Hashtable...\n");
		print_key(key);
	}

	if((skb = dev_alloc_skb(headers))) {
		/* Assume failure. Then contradict later. */
		reply 		= (Message *) (skb->data + LINK_HDR);
		reply->status 	= REPLY_BAD;
	} else {
		/* Not enough memory to allocate reply packet. 
		 * Use the original instead to send back error. 
		 */
		skb 		= orig;
		reply		= m;
		reply->status 	= REPLY_FULL;
		atomic_inc(&skb->users);
		goto nomem;
	}

	/* Do we already know about the page at this particular offset */
	if(!(block = (Block *) lookup_table(&table, key))) {
		if((totalram_pages - nr_free_pages()) > percent) {
			if(debug)
			       printk(KERN_ALERT 
			       		" Threshold reached. Cannot allocate.\n");
			reply->status = REPLY_FULL;
			goto nomem;
		}

		if(debug)
			printk(KERN_ALERT " Not Found. Trying to allocate...\n");

		/* Allocate a new page for storage of this new request */
		if((pg = alloc_pages(__GFP_HIGHMEM | GFP_ATOMIC, 0))) {
			block = (Block *)kmem_cache_alloc(block_cache, GFP_ATOMIC);
			if(block) {
				block->page = pg;
				block->key = *key;
				num_blocks++;
				insert_table(&table, &block->key, block);
			} else {
				__free_page(pg);
				block = NULL;
				reply->status = REPLY_FULL;
				goto nomem;
			}
		} else {
			if(debug)
				printk(KERN_ALERT
				     " Server full. %lu pages stored. Rejecting.\n",
				     num_blocks);
			reply->status = REPLY_FULL;
		}
	} else if(debug)
		printk(KERN_ALERT " Old page found.\n");
	 
	/* Copy the page into main memory */
	if(block) {
		addr = kmap_atomic(block->page, KM_SKB_DATA_SOFTIRQ);

		if(addr) {
			/* Again, isn't this subtraction always zero? */
			memcpy(addr + (m->offset - key->offset), m + 1, m->count);
			reply->status = REPLY_GOOD;
		} else
			BUG();

		kunmap_atomic(addr, KM_SKB_DATA_SOFTIRQ);
	} 

nomem:
	/* Setup the Acknowledgement */
	reply->type 	= PAGE_OUT_REPLY;
	reply->offset 	= m->offset;
	reply->vpid 	= m->vpid;
	reply->seq	= m->seq;
	reply->count	= m->count;

	if(plog) {
		do_gettimeofday(&end);
		timersub(&end, &start, &diff);
		reply->serv_internal_time = USEC(diff);	
	}

	skb->len = headers;

	return skb;
}

void mem_cleanup(void) {
	Block 	* block;
	int 	freed_blocks = 0;

	del_timer(&broadcast_timer);
	del_timer(&session_timer);
	
	printk(KERN_ALERT "Trying to free %lu pages.\n", num_blocks);

	rewind_table(&table);

	while((block = next_in_table(&table))) {
		if(!remove_current_in_table(&table))
			BUG();

		__free_page((struct page *) block->page);
		kmem_cache_free(block_cache, block);
		freed_blocks++;
	}

	if(table.count) 
		printk(KERN_ALERT "Error: Entries still exist in hashtable!\n");

	printk(KERN_ALERT "Cleaning up hashtable.\n");

	clean_table(&table);

	if(kmem_cache_destroy(block_cache))
		printk(KERN_ALERT "Error: All memory pages not deallocated.\n");

	printk(KERN_ALERT "Freed %d total blocks from memory.\n", freed_blocks);

	mem_print();
}

void mem_print(void) {
	printk(KERN_ALERT "Total free pages: %lu, %lu MB\n", 
		nr_free_pages(), MB(nr_free_pages()));
	printk(KERN_ALERT "Total RAM: %lu, %lu MB\n", 
		totalram_pages, MB(totalram_pages));
}
	
void server_cleanup(void) {
	net_exit();

	mem_cleanup();	
		
	printk(KERN_ALERT "Removing proc files.\n");
	if(aprocfile)
		remove_proc_entry(PROCNAME, NULL);
	remove_proc_entry(PROCNAME2, NULL);
	remove_proc_entry(PROCNAME3, NULL);

	printk(KERN_ALERT "Anemone Server Stopped.\n");
}	

/*
 * /proc/sinfo Write function
 */
int toggle_logging(struct file* file, const char* buffer, 
		   ulong count, void* data) {
	if(plog) {
		printk(KERN_ALERT "Anemone: Turning packet logging OFF.\n");
		plog = 0;
	} else {
		printk(KERN_ALERT "Anemone: Turning packet logging ON.\n");
		plog = 1;
	}

	return count;
}
int toggle_debug(struct file* file, const char* buffer, 
		 ulong count, void* data) {
	if(debug) {
		printk(KERN_ALERT "Anemone: Turning debugging OFF.\n");
		debug = false;
	} else {
		printk(KERN_ALERT "Anemone: Turning debugging ON.\n");
		debug = true;
	}

	return count;
}

/*
 * receive client status
 */
void session_timeout(ulong timer){
	Session *p;
	list_head *x, *y;
	timeval curr, sub;
	ulong flags;
	if (timer != SESSION_TIMEOUT)
		return;
	local_irq_save(flags);
	local_irq_disable();
	
	list_for_each_safe(x, y, &session_list) {
		p = list_entry(x, struct Session, queue);

		do_gettimeofday(&curr);
		timersub(&curr, &p->timestamp, &sub);
		/* Give client 3 retry times */
		if (USEC(sub) > (3 * SESSION_TIMEOUT * 10000) ) {
			delete_session(p->mac);
			list_del(x);
			kfree(p);
			printk(KERN_ALERT "Delete Session from list with mac addr = %02x:%02x:%02x:%02x:%02x:%02x.\n", mac_ntoa(p->mac));
		}
	}

	local_irq_restore(flags);
	mod_timer(&session_timer,  jiffies + SESSION_TIMEOUT);
}

void delete_session(u8* mac) {
	Block 	* block;

	rewind_table(&table);

	while((block = next_in_table(&table))) {
		if(memcmp(block->key.client_mac, mac, ETH_ALEN) == 0) {
			if(!remove_current_in_table(&table))
				BUG();
			__free_page((struct page *) block->page);
			kmem_cache_free(block_cache, block);
		}
	}

}
	
void update_session_status(u8* clientmac, ulong sessionid){
	Session *session,*p = NULL;
	struct list_head *pos;
	bool found = false;

	list_for_each(pos, &session_list) {
		p = list_entry(pos, struct Session, queue);
		if (( memcmp(clientmac, p->mac, ETH_ALEN) == 0) && (sessionid == p->sessionid)){
			do_gettimeofday(&p->timestamp);
			found = true;
		}
	}

	if (!found) {
		session = (struct Session*)kmalloc(sizeof(struct Session), GFP_ATOMIC);
		do_gettimeofday(&session->timestamp);
		memcpy(session->mac, clientmac, ETH_ALEN);
		session->sessionid = sessionid;
		printk(KERN_ALERT "\nNew Session setup mac=%x\n", session->mac[5]);
		list_add_tail(&session->queue, &session_list);
	}
}


/*
 * Broadcast server status
 */
void broadcast_timeout(ulong timer){
	if (timer != SERVER_BROADCAST_TIMEOUT) 
		return;
	broadcast_status();
	mod_timer(&broadcast_timer,  jiffies + SERVER_BROADCAST_TIMEOUT);
}

void broadcast_status(){	
	Message   *m;
	struct sk_buff *skb;
	ethhdr *eth;	

	u8	dest[ETH_ALEN];
	skb = dev_alloc_skb(headers);

	if(!skb)
		BUG();

	skb->len = headers;

	eth = (ethhdr *) skb->data;
	m = (Message *) (eth + 1);
	m->seq = 0;
	m->status = 0;
	m->count = 0;
	m->offset = 0;
	m->vpid = 0;
	m->serv_internal_time = 0;
	m->fragbyte = 0;		
	m->fragment = 0;
	m->fragleft = 0;
	m->type	= TYPE_SERVER_BROADCAST;
	fill_status_info(m);
	if(debug)
		printk(KERN_ALERT "Broadcast avail_mem = %d pages\n", m->avail_mem);

	inet_mac(BROADCAST_MAC, dest);

	anemone_send(skb, dest);

}

void fill_status_info(Message * reply) {	
	reply->avail_mem = nr_free_pages();
	reply->total_mem = totalram_pages;
	if((totalram_pages - nr_free_pages()) > percent) {
		printk(KERN_ALERT " Threshold reached. Cannot serve for client.\n");
		reply->server_status = STATUS_UNAVAIL;
	} else {
		reply->server_status = STATUS_GOOD;
	}
}

/* Unimportant function */
int query_block(struct file* file, const char* buffer, 
		ulong count, void* data) {
	char stuff[100], *end;
	u64 offset;
	Message * msg;
	bool found = false;

	if(count <= 100) {
		stuff[count]='\0';
		if(copy_from_user(stuff, buffer, count)) {
			printk(KERN_ALERT "User copy failed.\n");
			return 0;
		}
		offset = simple_strtoull(stuff, &end, 10);
		rewind_table(&table);	
		do {
			if((msg = next_in_table(&table))) {
				if(offset == msg->offset) {
					printk(KERN_ALERT "Yes, it's here.\n");
					found = true;
					break;
				}
			}
		} while(msg);

		if(!found)
			printk(KERN_ALERT "No, it's not here.\n");

		return count;
	}
	return 0;
}


/*
 * /proc/sdebug. Reports server usage statistics.
 */
int server_report(char *buf, char **start, off_t off, int cnt, int * eof, void * data)
{
	int len = 0;

	if(off > 0)
		return 0;

	*start = report_data;

	len += sprintf(report_data + len, "Interface Drops:\t%d\n", if_drops); 
	len += sprintf(report_data + len, "Server usage:\t\t%lu (%lu MB)\n",
			table.count, MB(table.count)); 
	len += sprintf(report_data + len, "Free RAM:\t%lu, %lu MB\n", 
			nr_free_pages(), MB(nr_free_pages()));
	len += sprintf(report_data + len, "Total RAM:\t%lu, %lu MB\n", 
			totalram_pages, MB(totalram_pages));
	len += sprintf(report_data + len, "Page-IN:\t\t%lu\n", reads);
	len += sprintf(report_data + len, "Page-OUT:\t\t%lu\n", writes);

	*eof = 1;

	return len;
}

module_init(server_init);
module_exit(server_cleanup);