#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <getopt.h>
#include <unistd.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <sys/stat.h>

#include <PR/bbfs.h>
#include "ecc256.h"
#include "flashif.h"

#define MAX_FSFILES	64

/* default number of blocks */
static int blocks = FLASHIF_DEFAULT_NUM_BLKS;
static int sysblocks = BB_FL_BYTE_TO_BLOCK(BB_SYSTEM_AREA_SIZE);

static char* skfile = "kernel";
static char* balfile = "bootlic";
static char* bafile = "bootapp";

static int verbose;

static BbFat16 fat;

#define VERBOSE(x)	do { if(verbose) printf x; } while(0)

flashif_t* fif;

/*
 * read all of the blocks on the flash, checking the disposiion of
 * the block by examining the spare area.
 */
static void
scan_bad_blocks(BbFat16* fat) {
    u8 spare[BB_FL_SPARE_SIZE];
    int i, status;
    if (verbose) printf("bad blocks: ");
    for(i = 0; i < blocks; i++) {
        if ((*fif->read_block)(fif->f, i, NULL, spare, NO_ECC) < 0)
	    status = BB_FAT_BAD;
	else
	    status = BB_FAT_AVAIL;
	fat->entry[i] = status;
        if (verbose && fat->entry[i] == BB_FAT_BAD) printf("%4d ", i);
    }
    if (verbose) printf("\n");
    /* non-existant blocks are bad */
    for(; i < blocks; i++)
    	fat->entry[i] = BB_FAT_BAD;
}

/*
 * copy reserved file to flash, starting at <start>, skipping bad blocks
 * along the way.  update the fat with reserved or bad block entries.
 * return the address of the next target block or -1 on error
 */
static int
write_rsvfile(u32 start, BbFat16* fat, const char* file, const u8* spare) {
    u8 buf[BB_FL_BLOCK_SIZE];
    int rv = -2, b = start, n, prev = 0, once = 0;
    FILE* fp = fopen(file, "r");

    VERBOSE(("write rsv file    %-16s @ %ld: ", file, start));
    if (!fp) return -1;
    while((n = fread(buf, 1, sizeof buf, fp)) > 0 || !once) {
    	if (n < BB_FL_BLOCK_SIZE)
	    memset(buf+n, 0, BB_FL_BLOCK_SIZE-n);
retry:
	if (b >= sysblocks) goto error;
	if ((*fif->write_block)(fif->f, b, buf, spare, YES_ECC) < 0) {
	    VERBOSE(("%4d!", b));
	    fat->entry[b] = BB_FAT_BAD;
	    b++;
	    goto retry;
	} else {
	    VERBOSE(("%4d ", b));
	    fat->entry[b] = BB_FAT_RESERVED;
//	    if (prev) fat->entry[prev] = b;
	    prev = b;
	    b++;
	}
	once = 1;
    }
    rv = b;
error:
    VERBOSE(("\n"));
    fclose(fp);
    return rv;
}

/*
 * copy reserved file to flash, starting at <start>, skipping bad blocks
 * along the way.  update the fat with reserved or bad block entries
 * and write the address of the successor block into the spare area
 * of each block. return the address of the next target block or -1
 * on error
 */
static int
write_rsvlinkfile(u32 start, BbFat16* fat, const char* file) {
    u8 buf[BB_FL_BLOCK_SIZE], spare[BB_FL_SPARE_SIZE];
    int rv = -2, b = start, n, prev = 0, p;
    FILE* fp = fopen(file, "r");
    struct stat sb;

    VERBOSE(("write rsv lk file %-16s @ %ld: ", file, start));
    if (!fp) return -1;
    if (fstat(fileno(fp), &sb) < 0) goto error;
    p = BB_FL_BYTE_TO_BLOCK(sb.st_size + BB_FL_BLOCK_SIZE-1);
    memset(spare, 0xff, sizeof spare);

    /* write in reverse order */
    while(--p >= 0) {
	fseek(fp, BB_FL_BLOCK_TO_BYTE(p), SEEK_SET);
        if ((n = fread(buf, 1, sizeof buf, fp)) < 0) break;
    	if (n < BB_FL_BLOCK_SIZE)
	    memset(buf+n, 0, BB_FL_BLOCK_SIZE-n);
retry:
	if (b >= sysblocks) goto error;
	if (prev) {
            if(prev>255) return -1;
	    spare[BB_FL_BLOCK_LINK_OFF] = prev;
	    spare[BB_FL_BLOCK_LINK_OFF+1] = prev;
	    spare[BB_FL_BLOCK_LINK_OFF+2] = prev;
	}
	if ((*fif->write_block)(fif->f, b, buf, spare, YES_ECC) < 0) {
	    VERBOSE(("%4d!", b));
	    fat->entry[b] = BB_FAT_BAD;
	    b++;
	    goto retry;
	} else {
	    VERBOSE(("%4d ", b));
	    fat->entry[b] = BB_FAT_RESERVED;
//	    if (prev) fat->entry[prev] = b;
	    prev = b;
	    b++;
	}
    }
    rv = b;
error:
    VERBOSE(("\n"));
    fclose(fp);
    return rv;
}

/*
 * copy file system file to flash, using the fat for allocation,
 * appending an entry to the inode table.
 * return 0 on success, -1 on failure.
 */
static int
write_fsfile(BbFat16* fat, const char* file) {
    u8 buf[BB_FL_BLOCK_SIZE];
    int n, b, rv = -2, prev = 0, size = 0, incr, i;
    FILE* fp = fopen(file, "r");
    BbInode* in = 0;
    struct stat sb;
    const char* p;


    VERBOSE(("write fs file     %-16s ", file));
    if (!fp || fstat(fileno(fp), &sb) < 0) return -1;

    /* find a free inode */
    for(n = 0; n < BB_INODE16_ENTRIES; n++) {
        if (fat->inode[n].type == 0) {
	    in = fat->inode+n;
	    break;
	}
    }
    if (n == BB_INODE16_ENTRIES) goto error;

    if (sb.st_size > BB_BIG_FILE_THRESHOLD) {
	/* big files go low in the fat */
	b = sysblocks; incr = 1;
	VERBOSE(("B "));
    } else {
	/* small files at the other end */
	b = blocks-1; incr = -1;
	VERBOSE(("S "));
    }

    in->block = BB_FAT_LAST;
    while((n = fread(buf, 1, sizeof buf, fp)) > 0) {
    	if (n < BB_FL_BLOCK_SIZE)
	    memset(buf+n, 0, BB_FL_BLOCK_SIZE-n);
	size += n;
retry:
	/* allocate a block */
	//if (b >= blocks) goto error;
	while(b >= 0 && b < BB_FAT16_ENTRIES && fat->entry[b] != BB_FAT_AVAIL)
	    b += incr;
	if (b < 0 || b >= BB_FAT16_ENTRIES) goto error;
	if ((*fif->write_block)(fif->f, b, buf, 0, YES_ECC) < 0) {
	    VERBOSE(("%4d!", b));
	    fat->entry[b] = BB_FAT_BAD;
	    goto retry;
	} else {
	    VERBOSE(("%4d ", b));
	    fat->entry[b] = BB_FAT_LAST;
	    if (prev) fat->entry[prev] = b;
	    else in->block = b;
	    prev = b;
	}
    }
    if (!(p = strrchr(file, '/'))) p = file;
    else p++;
    memset(in->name, 0, sizeof in->name);
    for(i = 0; i < 8 && *p && *p != '.';)
	in->name[i++] = *p++;
    if (strrchr(p, '.'))
	strncpy(in->name+8, strrchr(p, '.')+1, 3);

    in->type = 1;
    in->size = (size+BB_FL_BLOCK_SIZE-1)&~(BB_FL_BLOCK_SIZE-1);
    rv = 0;
error:
    if (in && !in->type) memset(in, 0, sizeof *in);
    VERBOSE((" = %d\n", size));
    fclose(fp);
    return rv;
}

static int
write_fat(BbFat16* fat) {
    int rv0, rv1, i;
    u16 csum = 0, *p = (u16*)fat;

    memcpy(&fat->magic, BB_FAT16_MAGIC, 4);
    fat->seq = 1;

    /* fix endianness */
    for(i = 0; i < BB_FAT16_ENTRIES; i++)
	fat->entry[i] = htons(fat->entry[i]);
    for(i = 0; i < BB_INODE16_ENTRIES; i++) {
	fat->inode[i].block = htons(fat->inode[i].block);
	fat->inode[i].size = htonl(fat->inode[i].size);
    }
    fat->seq = htonl(fat->seq);

    /* update version stamp, checksum */
    for(i = 0; i < BB_FL_BLOCK_SIZE/sizeof(u16)-1; i++)
	csum += ntohs(p[i]);
    fat->cksum = htons(BB_FAT16_CKSUM - csum);

    rv0 = (*fif->write_block)(fif->f, blocks-1, fat, 0, YES_ECC);
    rv1 = (*fif->write_block)(fif->f, blocks-2, fat, 0, YES_ECC);
    if (rv0 < 0 && rv1 < 0)
    	return rv1;
    return 0;
}

static int
write_fail(int code, const char* file) {
    if (code < -1)
    	fprintf(stderr, "writeflash: write of %s failed\n", file);
    else 
        perror(file);
    return 1;
}

int
main(int argc, char* argv[]) {
    int c, i;
    int n_fsfiles = 0;
    char* fsfile[MAX_FSFILES];
    u32 cursor = 0;
    int rv, balblock[2];
    u8 spare[BB_FL_SPARE_SIZE];

    while ((c = getopt(argc, argv, "b:k:l:s:v")) != EOF) {
        switch(c) {
	case 's':
	    blocks = strtoul(optarg, 0, 0);
	    if (blocks <= 0 || (blocks&(blocks-1)) != 0) goto usage;
	    blocks = BB_FL_BYTE_TO_BLOCK(blocks*1024*1024);
	    break;
	case 'k':
	    skfile = optarg; break;
	case 'l':
	    balfile = optarg; break;
	case 'b':
	    bafile = optarg; break;
	case 'v':
	    verbose = 1; break;
	case '?':
usage:
	    fprintf(stderr, "Usage: writeflash -s MBytes -k kernel"
		    " -b bootapp -l boot_license [fsfile0 [fsfile1 ...]]\n");
	    return 1;
	}
    }

    while(optind < argc)
    	fsfile[n_fsfiles++] = argv[optind++];


    if (!(fif = new_fileif()))
	return 1;
    if (!(*fif->open)(&fif->f, "flash.img", 1, blocks)) {
	fprintf(stderr, "flash device open failed\n");
	return 1;
    }
    blocks = (*fif->blocks)(fif->f);
    VERBOSE(("blocks %d\n", blocks));


restart:
    /* initialize the fat with bad block information */
    scan_bad_blocks(&fat);
    for(i = 0; i < BB_FAT16_BLOCKS; i++) {
	if (fat.entry[blocks-i-1] == BB_FAT_AVAIL)
	    fat.entry[blocks-i-1] = BB_FAT_RESERVED;
    }
    for(i = 0; i < BB_SYSTEM_AREA_SIZE/BB_FL_BLOCK_SIZE; i++) {
	if (BB_FAT16_NEXT(&fat,i) != BB_FAT_BAD)
	    BB_FAT16_NEXT(&fat,i) = BB_FAT_RESERVED;
    }

    /* write secure kernel */
    if ((rv = write_rsvfile(cursor, &fat, skfile, 0)) < 0) {
	return write_fail(rv, skfile);
    }
    cursor = rv;

    /* write bootapp license 0 */
    if ((rv = write_rsvfile(cursor, &fat, balfile, 0)) < 0) {
	return write_fail(rv, balfile);
    }
    cursor = rv;
    balblock[0] = rv - 1;

    /* write bootapp license 1 */
    if ((rv = write_rsvfile(cursor, &fat, balfile, 0)) < 0) {
	return write_fail(rv, balfile);
    }
    cursor = rv;
    balblock[1] = rv - 1;

    /* write bootapp copy 0 */
    if ((rv = write_rsvlinkfile(cursor, &fat, bafile)) < 0) {
	return write_fail(rv, bafile);
    }
    cursor = rv;

    /* rewrite bootapp license 0. patching the next block pointer */
    memset(spare, 0xff, sizeof spare);
    if((rv-1)>255) return 1;
    spare[BB_FL_BLOCK_LINK_OFF] = rv-1;
    spare[BB_FL_BLOCK_LINK_OFF+1] = rv-1;
    spare[BB_FL_BLOCK_LINK_OFF+2] = rv-1;
    spare[BB_FL_SEQ_OFF] = 0;
    VERBOSE(("update link @ %d -> %d\n", balblock[0], rv-1));
    if ((rv = write_rsvfile(balblock[0], &fat, balfile, spare)) < 0) {
	write_fail(rv, balfile);
	goto restart;
    }

    /* write bootapp copy 1 */
    if ((rv = write_rsvlinkfile(cursor, &fat, bafile)) < 0) {
	return write_fail(rv, bafile);
    }
    cursor = rv;

    /* rewrite bootapp license 1. patching the next block pointer */
    memset(spare, 0xff, sizeof spare);
    if((rv-1)>255) return 1;
    spare[BB_FL_BLOCK_LINK_OFF] = rv-1;
    spare[BB_FL_BLOCK_LINK_OFF+1] = rv-1;
    spare[BB_FL_BLOCK_LINK_OFF+2] = rv-1;
    spare[BB_FL_SEQ_OFF] = 1;
    VERBOSE(("update link @ %d -> %d\n", balblock[1], rv-1));
    if ((rv = write_rsvfile(balblock[1], &fat, balfile, spare)) < 0) {
	write_fail(rv, balfile);
	goto restart;
    }

    /* write file system files */
    for(i = 0; i < n_fsfiles; i++) {
        if ((rv = write_fsfile(&fat, fsfile[i])) < 0) {
	    return write_fail(rv, fsfile[i]);
	}
    }
    if ((rv = write_fat(&fat)) < 0) {
        fprintf(stderr, "writeflash: fat write failed\n");
	return 1;
    }
    if (verbose) {
	int free = 0;
	for(i = 0; i < blocks; i++)
	    if (fat.entry[i] == BB_FAT_AVAIL) free++;
	printf("space remaining %d blocks %.2f MB\n", free, BB_FL_BLOCK_TO_BYTE(free)/(1024.*1024.));
    }
    (*fif->close)(fif->f);
    return 0;
}