#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <ctype.h>
#include <sys/stat.h>
#include <sys/mman.h>

#include "nand.h"

#define NAND_OOB_SZ		16
#define NAND_PAGE_SZ		512
#define NAND_SZ			(64*1024*1024)
#define NAND_PAGES		(NAND_SZ/NAND_PAGE_SZ)
#define NAND_PAGES_MASK		(NAND_PAGES-1)
#define NAND_IMAGE_SZ		(NAND_PAGES*(NAND_OOB_SZ+NAND_PAGE_SZ))

/* addressing only supports 8 bits of address within page */
#define NAND_ADDR_BLOCK_SHIFT	13
#define NAND_ADDR_PAGE_SHIFT	8
#define NAND_ADDR_PAGE_MASK	((1<<NAND_ADDR_PAGE_SHIFT)-1)
#define NAND_ADDR_OOB_MASK	(NAND_OOB_SZ-1)

#define NAND_PPB_SHIFT		(NAND_ADDR_BLOCK_SHIFT-NAND_ADDR_PAGE_SHIFT)
#define NAND_PAGE_PER_BLOCK	(1<<NAND_PPB_SHIFT)
#define NAND_BLOCKS		(NAND_SZ/(NAND_PAGE_PER_BLOCK*NAND_PAGE_SZ))

#define NAND_CLE		0x1	/* command latch enable */
#define NAND_ALE		0x2	/* address latch enable */
#define NAND_MODE_CMD		(NAND_CLE)
#define NAND_MODE_DATA		(0)
#define NAND_MODE_ADDR		(NAND_ALE)

#define NAND_CMD_READ_0		0x00
#define NAND_CMD_READ_1		0x01
#define NAND_CMD_PROGRAM	0x10
#define NAND_CMD_DUMMY_PROGRAM	0x11
#define NAND_CMD_READ_2		0x50
#define NAND_CMD_ERASE_0	0x60
#define NAND_CMD_ERASE_1	0xD0
#define NAND_CMD_STATUS		0x70
#define NAND_CMD_MP_STATUS	0x71
#define NAND_CMD_DATA_INPUT	0x80
#define NAND_CMD_READ_ID	0x90
#define NAND_CMD_RESET		0xff

#define NAND_STATUS_CMD_PASS	0x00
#define NAND_STATUS_CMD_FAIL	0x01
#define NAND_STATUS_READY	0x40
#define NAND_STATUS_BUSY	0x00
#define NAND_STATUS_WRITE_OK	0x80
#define NAND_STATUS_WRITE_NOK	0x00

#define MULTIPLANE

#define VERBOSE(x)

static struct nand {
    unsigned char ctrl;
    unsigned char io;
    unsigned char state;
    unsigned char cmd;
    unsigned char prevcmd;
    unsigned char status;
#ifdef MULTIPLANE
    unsigned char id[4];
#else
    unsigned char id[2];
#endif
    unsigned char idp;
    unsigned short bufp;
    unsigned int pageoff;
    unsigned int addr;
    unsigned int pagep;
    unsigned int readp;
    unsigned char buf[NAND_PAGE_SZ+NAND_OOB_SZ];
#ifdef MULTIPLANE
#define PLANES	4
#define PAGE_PLANE(paddr)	(((paddr)>>5)&3)
#define BLOCK_PLANE(baddr)	((baddr)&3)
    struct {
	unsigned char buf[PLANES][NAND_PAGE_SZ+NAND_OOB_SZ];
	unsigned int  addr[PLANES];
	unsigned char status;
	unsigned char mask;
    } mp;
#endif
    unsigned char *mem;
    unsigned char map[NAND_BLOCKS/8];
#define DEV_ADDR(page,a) ((page)*(NAND_PAGE_SZ+NAND_OOB_SZ)+(a))
} dev;

void
nand_init(void) {
    dev.ctrl = 0;
    dev.addr = 0;
    dev.cmd = dev.prevcmd = 0;
    dev.state = 0;
#ifdef MULTIPLANE
    dev.id[0] = 0xec;	/* maker */
    dev.id[1] = 0x76;	/* device */
    dev.id[2] = 0xa5;	/* don't care */
    dev.id[3] = 0xc0;	/* supports multiplane */
#else
    dev.id[0] = 0x98;	/* maker */
    dev.id[1] = 0x76;	/* device */
#endif
    dev.bufp = dev.idp = 0;
    dev.status = NAND_STATUS_CMD_PASS | NAND_STATUS_READY | NAND_STATUS_WRITE_OK;
#ifdef MULTIPLANE
    memset(&dev.mp, 0, sizeof dev.mp);
#endif
    memset(dev.map, 0, sizeof dev.map);
    if (getenv("SIMOS_FLASH")) {
	FILE* fp;
	unsigned char *p;
	unsigned int x;
        char line[1024];
        int num_read;
	struct stat sb;
	if ((fp = fopen(getenv("SIMOS_FLASH"), "r")) == NULL) {
	    perror(getenv("SIMOS_FLASH"));
	    exit(1);
	}
	num_read = fread(line, 1, sizeof line, fp);
	for(x = 0; x < num_read; x++)
	    if (!isascii(line[x])) goto binary;

	if (!(p = dev.mem = malloc(NAND_IMAGE_SZ))) {
	    perror("nand malloc");
	    exit(1);
	}
	memset(dev.mem, 0xff, NAND_IMAGE_SZ);
	fseek(fp, 0, 0);
	while(p < dev.mem+NAND_IMAGE_SZ) {
	    if(fscanf(fp, "%x", &x) == 1)
                *p++ = x;
            else if(fgets(line,sizeof line,fp)==NULL || line[0]!='/')
                break;
	}
	fclose(fp);
	return;
binary:
	if ((fp = freopen(getenv("SIMOS_FLASH"), "r+", fp)) == NULL) {
	    perror(getenv("SIMOS_FLASH"));
	    exit(1);
	}
	fstat(fileno(fp), &sb);
	if ( (dev.mem = 
              mmap(0, sb.st_size, PROT_READ|PROT_WRITE, MAP_SHARED, fileno(fp), 0))
             == (void*)-1) {
	    perror("nand mmap");
	    exit(1);
	}
    } else {
	if (!(dev.mem = malloc(NAND_IMAGE_SZ))) {
	    perror("nand malloc");
	    exit(1);
	}
	memset(dev.mem, 0xff, NAND_IMAGE_SZ);
    }
    if (getenv("SIMOS_FLASH_MAP")) {
	/* parse block map */
	char *map = getenv("SIMOS_FLASH_MAP");
	int o = 0, n, m;
	while(sscanf(map+o, " %d%n", &m, &n) == 1) {
	    o += n;
	    VERBOSE(printf("nand bad block = %d\r\n", m));
	    if (m < 0 || m >= NAND_BLOCKS) {
		printf("nand map bad block specification %d\r\n", m);
		exit(1);
	    }
	    dev.map[m/8] |= 1 << (m&7);
	}
	
    }
}

void
nand_final(void) {
    msync(dev.mem, NAND_IMAGE_SZ, MS_SYNC);
}

void
nand_ctrl(int w, unsigned char* v) {
    if (w)
	dev.ctrl = *v;
    else
	*v = dev.ctrl;
}

void
nand_io(int w, unsigned char* v) {
    if (!w && (dev.ctrl & (NAND_CLE|NAND_ALE)) != NAND_MODE_DATA) {
    	fprintf(stderr, "nand: illegal/unsupported read operation\r\n");
	return;
    }
    switch(dev.ctrl & (NAND_CLE|NAND_ALE)) {
    case NAND_MODE_CMD:
	dev.prevcmd = dev.cmd;
        switch(dev.cmd = *v) {
	case NAND_CMD_READ_0:
VERBOSE(printf("nand read0\r\n"));
	    dev.pageoff = 0;
	    break;
	case NAND_CMD_READ_1:
VERBOSE(printf("nand read1\r\n"));
	    dev.pageoff = 256;
	    break;
	case NAND_CMD_READ_2:
VERBOSE(printf("nand read2\r\n"));
	    dev.pageoff = NAND_PAGE_SZ;
	    break;
	case NAND_CMD_READ_ID:
VERBOSE(printf("nand readid\r\n"));
	    dev.idp = 0;
	    break;
	case NAND_CMD_STATUS:
VERBOSE(printf("nand status\r\n"));
	    break;
#ifdef MULTIPLANE
	case NAND_CMD_MP_STATUS:
VERBOSE(printf("nand multiplane status\r\n"));
	    break;
#endif
	case NAND_CMD_PROGRAM: {
	    int i, bad = 0;
VERBOSE(printf("nand program page %x\r\n", dev.addr>>NAND_ADDR_PAGE_SHIFT));
	    /*XXXblythe check previous command?*/
	    for(i = 0; i < NAND_PAGE_SZ+NAND_OOB_SZ; i++)
		dev.mem[DEV_ADDR((dev.addr>>NAND_ADDR_PAGE_SHIFT)&NAND_PAGES_MASK, i)] &= dev.buf[i];
	    if (dev.map[(dev.addr>>NAND_ADDR_BLOCK_SHIFT)/8] & (1<<((dev.addr>>NAND_ADDR_BLOCK_SHIFT)&7)))
		bad |= 1 << BLOCK_PLANE(dev.addr>>NAND_ADDR_BLOCK_SHIFT);
#ifdef MULTIPLANE
	    dev.mp.mask &= ~(1<<PAGE_PLANE(dev.addr>>NAND_ADDR_PAGE_SHIFT));
	    /* do up to 3 other planes planes */
	    if (dev.mp.mask) {
		int j;
VERBOSE(printf("nand multipage %x\r\n", dev.mp.mask));
		for(j =0; j < PLANES; j++) {
		    if (!(dev.mp.mask & (1<<j))) continue;
		    dev.mp.mask &= ~(1<<j);
		    for(i = 0; i < NAND_PAGE_SZ+NAND_OOB_SZ; i++)
			dev.mem[DEV_ADDR((dev.mp.addr[j]>>NAND_ADDR_PAGE_SHIFT)&NAND_PAGES_MASK, i)] &= dev.mp.buf[j][i];
		    if (dev.map[(dev.mp.addr[j]>>NAND_ADDR_BLOCK_SHIFT)/8] & (1<<((dev.mp.addr[j]>>NAND_ADDR_BLOCK_SHIFT)&7)))
			bad |= 1 << BLOCK_PLANE(dev.mp.addr[j]>>NAND_ADDR_BLOCK_SHIFT);
		}
	    }
	    dev.mp.status = bad << 1;
#endif
	    if (bad)
		dev.status |= NAND_STATUS_CMD_FAIL;
	    else
		dev.status &= ~NAND_STATUS_CMD_FAIL;
	    }
	    break;
#ifdef MULTIPLANE
	case NAND_CMD_DUMMY_PROGRAM: {
	    int i;
VERBOSE(printf("nand dummy program\r\n"));
	    memcpy(dev.mp.buf[PAGE_PLANE(dev.addr>>NAND_ADDR_PAGE_SHIFT)], dev.buf, sizeof dev.buf);
	    dev.mp.addr[PAGE_PLANE(dev.addr>>NAND_ADDR_PAGE_SHIFT)] = dev.addr;
	    dev.mp.mask |= 1 << PAGE_PLANE(dev.addr>>NAND_ADDR_PAGE_SHIFT);
	    /* check consistency */
	    for(i = 0; i < PLANES; i++) {
		if (!(dev.mp.mask&(1<<i))) continue;
		if ((dev.addr&0xff) != (dev.mp.addr[i]&0xff))
		    fprintf(stderr, "nand: inconsistent page address for multiplane program\r\n");
	    }}
	    break;
#endif
	case NAND_CMD_ERASE_0:
VERBOSE(printf("nand erase0\r\n"));
#ifdef MULTIPLANE
	    if (dev.prevcmd == NAND_CMD_ERASE_0) {
		dev.mp.addr[BLOCK_PLANE(dev.addr>>NAND_PPB_SHIFT)] = dev.addr;
		dev.mp.mask |= 1 << BLOCK_PLANE(dev.addr>>NAND_PPB_SHIFT);
	    }
#endif
	    break;
	case NAND_CMD_ERASE_1:
VERBOSE(printf("nand erase1\r\n"));
	    if (dev.prevcmd == NAND_CMD_ERASE_0) {
		int i, bad = 0;
		int page = dev.addr;
		page &= NAND_PAGES_MASK;
		page &= ~(NAND_PAGE_PER_BLOCK-1);
VERBOSE(printf("erase block @page %x\r\n", page));
		for(i = 0; i < NAND_PAGE_PER_BLOCK; i++)
		    memset(dev.mem+DEV_ADDR(page+i,0), 0xff, NAND_PAGE_SZ+NAND_OOB_SZ);
		if (dev.map[(dev.addr>>NAND_PPB_SHIFT)/8] & (1 << ((dev.addr>>NAND_PPB_SHIFT)&7)))
		    bad |= 1<<BLOCK_PLANE(dev.addr>>NAND_PPB_SHIFT);
#ifdef MULTIPLANE
		dev.mp.mask &= ~(1<<BLOCK_PLANE(dev.addr>>NAND_PPB_SHIFT));
		if (dev.mp.mask > 0) {
		    int j;
		    for(j = 0; j < PLANES; j++) {
			if (!(dev.mp.mask & (1<<j))) continue;
			dev.mp.mask &= ~(1<<j);
			page = dev.mp.addr[j];
			page &= NAND_PAGES_MASK;
			page &= ~(NAND_PAGE_PER_BLOCK-1);
VERBOSE(printf("erase mblock @page %x\r\n", page));
			for(i = 0; i < NAND_PAGE_PER_BLOCK; i++)
			    memset(dev.mem+DEV_ADDR(page+i,0), 0xff, NAND_PAGE_SZ+NAND_OOB_SZ);
			if (dev.map[(dev.mp.addr[j]>>NAND_PPB_SHIFT)/8] & (1 << ((dev.mp.addr[j]>>NAND_PPB_SHIFT)&7)))
			    bad |= 1<<BLOCK_PLANE(dev.mp.addr[j]>>NAND_PPB_SHIFT);
		    }
		}
		dev.mp.status = bad << 1;
#endif
		if (bad)
		    dev.status |= NAND_STATUS_CMD_FAIL;
		else
		    dev.status &= ~NAND_STATUS_CMD_FAIL;
	    } else
		dev.status |= NAND_STATUS_CMD_FAIL;
#ifdef MULTIPLANE
	    dev.mp.mask = 0;
#endif
	    break;
	case NAND_CMD_RESET:
VERBOSE(printf("nand reset\r\n"));
	    nand_init();
	    break;
	case NAND_CMD_DATA_INPUT:
VERBOSE(printf("nand data input\r\n"));
	    memset(dev.buf, 0xff, sizeof dev.buf);
	    dev.bufp = 0;
	    break;
	default:
	}
	dev.state = 0;
	break;
    case NAND_MODE_DATA:
	if (w) {
	    dev.buf[dev.bufp++] = *v;
	    if (dev.bufp > NAND_PAGE_SZ+NAND_OOB_SZ) abort();
	} else {
	    switch(dev.cmd) {
	    case NAND_CMD_READ_0:
	    case NAND_CMD_READ_1:
	    case NAND_CMD_READ_2:
		*v = dev.mem[DEV_ADDR(dev.pagep,dev.readp)];
		dev.readp++;
		if (dev.readp >= NAND_PAGE_SZ+NAND_OOB_SZ) {
		    dev.pagep++;
		    dev.pagep &= NAND_PAGES_MASK;
		    if (dev.pageoff != NAND_PAGE_SZ) dev.pageoff = 0;
		    dev.readp = dev.pageoff;
		}
		break;
	    case NAND_CMD_READ_ID:
	    	*v = dev.id[dev.idp++];
VERBOSE(printf("nand read id = 0x%x\r\n", *v));
#ifdef MULTIPLANE
		dev.idp &= 3;
#else
		dev.idp &= 1;
#endif
		break;
#ifdef MULTIPLANE
	    case NAND_CMD_MP_STATUS:
VERBOSE(printf("nand read mp status = 0x%x\r\n", dev.status | dev.mp.status));
	    	*v = dev.status|dev.mp.status;
		break;
#endif
	    case NAND_CMD_STATUS:
VERBOSE(printf("nand read status = 0x%x\r\n", dev.status));
	    	*v = dev.status;
		break;
	    }
	}
	dev.state = 0;
	break;
    case NAND_MODE_ADDR:
        if (dev.state == 0) dev.addr = 0;
	dev.addr |= *v << (8*dev.state++);
	if (dev.cmd == NAND_CMD_READ_0 || dev.cmd == NAND_CMD_READ_1 ||
		dev.cmd == NAND_CMD_READ_2) {
	    dev.pagep = (dev.addr >> NAND_ADDR_PAGE_SHIFT)&NAND_PAGES_MASK;
	    dev.readp = dev.pageoff + (dev.addr&((dev.pageoff == NAND_PAGE_SZ) ? NAND_ADDR_OOB_MASK : NAND_ADDR_PAGE_MASK));
	} else if (dev.cmd == NAND_CMD_DATA_INPUT)
	    dev.readp = dev.pageoff + (dev.addr&((dev.pageoff == NAND_PAGE_SZ) ? NAND_ADDR_OOB_MASK : NAND_ADDR_PAGE_MASK));
//printf("addr = %x pagep %x readp %x\r\n", dev.addr, dev.pagep, dev.readp);
	break;
    default:
    }
}