#include <linux/module.h> 
#include <linux/init.h> 
#include <linux/slab.h> 
#include <linux/unistd.h> 
#include <linux/sched.h> 
#include <linux/fs.h> 
#include <linux/file.h> 
#include <linux/mm.h> 
#include <linux/vmalloc.h>
#include <linux/proc_fs.h>
#include <linux/mtd/mtd.h>
#include <linux/seq_file.h>
#include <asm/uaccess.h> 
#include "bspchip.h"

#define ST_IDLE		0
#define ST_WRITING	1
#define ST_UNLOCKED	2
#define MAGIC_KEY	"1321"

static int state;
static int file_read(struct file *fp,char *buf,int len) {
	if (fp->f_op && fp->f_op->read)
		return fp->f_op->read(fp, buf, len, &fp->f_pos);
	else
		return -1;
}

// on success, this function will reboot system.
static int fw_write(char *fwfile) {
	mm_segment_t oldfs;
	struct mtd_info *mtd=NULL;
	struct file *fp; 	
	struct kstat stat;
	struct erase_info ei;	
	int offset, part, imgFileSize;
	int nRead=0, nWritten=0;
	loff_t pos;	
	u8 *buf;
	unsigned long flags=0;
	char *strtmp;
	int retry = 3;
	
	extern void (*_machine_restart)(char *command);
	
	int rv = -EINVAL;
		
	oldfs = get_fs(); 
	set_fs(KERNEL_DS); 
	
	/*formt should be "offset;filename;part;imgFileSize" */
	do {
		char *pc, *s_offset;
		
		pc = strchr(fwfile, ';');
		if (!pc) goto ERROR;
		s_offset = fwfile;
		fwfile = &pc[1];
		*pc = '\0';
		offset = simple_strtol(s_offset, NULL, 0);
		
		// Mason Yu
		pc = strchr(fwfile, ';');
		if (!pc) goto ERROR;		
		s_offset = &pc[1];
		strtmp = &pc[1];
		*pc = '\0';
		part = simple_strtol(s_offset, NULL, 0);
		
		// Mason Yu
		pc = strchr(strtmp, ';');
		if (!pc) goto ERROR;
		s_offset = &pc[1];
		strtmp = &pc[1];
		*pc = '\0';
		imgFileSize = simple_strtol(s_offset, NULL, 0);
		
	} while (0);
	printk("%s(%d): fwfile=%s, offset=%d, part=%d, imgFileSize=%d\n", __func__, __LINE__, fwfile, offset, part, imgFileSize);
	
	// Mason Yu
	if (part == 1) 
		mtd = get_mtd_device_nm("rootfs");
	else if (part == 2)
		mtd = get_mtd_device_nm("fs-bak");
	if (IS_ERR(mtd)) {
		printk("rootfs/fs-bak not found!\n");
		goto ERROR;
	}
	
	//printk("MTD size: %lld erase size: %d, write size: %d\n", mtd->size, mtd->erasesize, mtd->writesize);
		
	fp = filp_open(fwfile, O_RDONLY, 0); 	
	//printk("%s(%d): fp=%p\n",__func__,__LINE__,fp);
	if(IS_ERR(fp))
		goto ERROR1;
	
	fp->f_pos = offset;
	
	rv = vfs_getattr(&fp->f_path, &stat); 
	//printk("%s(%d): rv=%d\n",__func__,__LINE__,rv);
	if (rv)
		goto ERROR2;		
	printk("%s: is %lld bytes\n", fwfile, stat.size);
	
	// Mason Yu
	//if ((stat.size - sizeof(imghdr) - offset) > mtd->size) {
	if ((stat.size - offset) > mtd->size) {
		printk("fw is larger than mtd(%s)\n", mtd->name);
		goto ERROR2;
	}
#ifdef CONFIG_MTD_NAND
	buf = kmalloc(mtd->erasesize+1, GFP_KERNEL);
#else
	buf = vmalloc(65536);
#endif
	if (!buf)
		goto ERROR2;
	
	printk("fw ready to write, pos=%lld\n", fp->f_pos);
	
	/* start to write */	
	local_irq_save(flags);
	
	/* hacking to disable watchdog during upgrade */
	REG32(BSP_WDTCTRLR) &= (~WDT_E);
	
RETRY:	
	memset(&ei, 0, sizeof(ei));
	ei.mtd = mtd;
	ei.addr = 0;
	ei.len  = (stat.size + mtd->erasesize - 1) & ~(mtd->erasesize - 1);
	printk("fw erase, addr=%llx, len=%llx , stat.size=%lld, mtd->erasesize=%x\n", ei.addr, ei.len, stat.size, mtd->erasesize);
	rv = mtd->_erase(mtd, &ei);
	printk("fw erase = %d\n", rv);
	
	if (rv) 
		goto ERROR3;
		
	pos = 0;	
	while (nWritten < imgFileSize) {
		size_t retlen;
		size_t rsize;
#ifdef CONFIG_MTD_NAND
		nRead = (mtd->erasesize >
	 		(imgFileSize - nWritten)) ? (imgFileSize -
				   		   nWritten) : mtd->erasesize;
#else
		nRead = (65536 >
			 (imgFileSize - nWritten)) ? (imgFileSize -
					       nWritten) : 65536;
#endif		
		rsize = file_read(fp, buf, nRead);
#ifdef CONFIG_MTD_NAND				
		if (rsize < mtd->erasesize)
			rsize = mtd->erasesize; 
#endif			
		printk("writing to %llx, size %x\n", pos, rsize);
		rv = mtd->_write(mtd, pos, rsize, &retlen, buf);
		if (rv || (rsize != retlen)) {
			printk("%s(%d): rv=%d rsize=%d retlen=%d\n",__func__,__LINE__,rv,rsize,retlen);
			goto ERROR3;
		}
		pos += rsize;
		nWritten += rsize;
	} 
	
	
	printk("restart\n");
	if (_machine_restart)
		_machine_restart(NULL);
	
	rv = 0;
ERROR3:
	if (retry) {
		printk("fw update failed: retrying.. %d\n",retry);
		retry--;		
		goto RETRY;
	}
#ifdef CONFIG_MTD_NAND	
	kfree(buf);
#else
	vfree(buf);
#endif	
ERROR2:	
	filp_close(fp, NULL);	
ERROR1:		
	put_mtd_device(mtd);
ERROR:	
	set_fs(oldfs);
	local_irq_restore(flags);
	return rv;
}

static int fw_proc_show(struct seq_file *s, void *data) {	
	switch(state) {
	case ST_IDLE: 	 return seq_printf(s, "idle"); 
	case ST_WRITING: return seq_printf(s, "writing"); 
	case ST_UNLOCKED:return seq_printf(s, "unlocked"); 
	}
	return 0;
}

static ssize_t fw_proc_write(struct file *file, const char __user * buffer, size_t count, loff_t * off) {
	char buf[128];
	if (copy_from_user(buf, buffer, sizeof(buf))) {
		return -EFAULT;
	}
	printk("%s(%d): buf = %s\n", __func__,__LINE__,buf);
	switch (state) {
	case ST_IDLE: 	 		
		if (!strncmp(buf, MAGIC_KEY, sizeof(MAGIC_KEY))) 
			state = ST_UNLOCKED;
		break;
	case ST_WRITING: 
		return -EFAULT;
	case ST_UNLOCKED:
		if (fw_write(buf))
			state = ST_IDLE;
		break;
	}
	
	return count;
}

static int fw_proc_open(struct inode *inode, struct file *file)
{
	return(single_open(file, fw_proc_show, NULL));
}

static const struct file_operations fwup_proc_fops = {
	.open           = fw_proc_open,
	.write		    = fw_proc_write,
	.read           = seq_read,
	.llseek         = seq_lseek,
	.release        = single_release,
};

extern struct proc_dir_entry *realtek_proc;

static int __init fwupdate_init(void) {
	struct proc_dir_entry *pe;
		
	pe = proc_create_data("fwupdate", S_IRUSR |S_IWUSR | S_IRGRP | S_IROTH, realtek_proc, &fwup_proc_fops, NULL);
	if (!pe) {
		return -EINVAL;
	}
	
	return 0;
}

static void __exit fwupdate_exit(void) {
}

module_init(fwupdate_init); 
module_exit(fwupdate_exit); 

MODULE_DESCRIPTION("FwUpdate"); 
MODULE_AUTHOR("Andrew Chang<yachang@realtek.com>"); 
MODULE_LICENSE("GPL");