#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <assert.h>

#include "stm32f10x.h"
#include "spi.h"
#include "flash.h"

#define FL_SELECT()	GPIO_ResetBits(GPIOA, GPIO_Pin_4)
#define FL_DESELECT()	GPIO_SetBits(GPIOA, GPIO_Pin_4)

static const char *flstattbl[] = {
    "BUSY",
    "WEL",
    "BP0",
    "BP1",
    "BP2",
    "BP3",
    "AAI",
    "BPL"
};

#define RW_IDLE		0
#define RW_RUNNING	1

static int writestate = RW_IDLE;
static int readstate = RW_IDLE;

void
flashcmd(int argc, char **argv) {
    uint8_t status, tmp, len;
    uint32_t addr;
    
    if (argc == 0) {
	fputs("No command specified\n", stdout);
	return;
    }
    
    if (!strcmp(argv[0], "str")) {
	status = flashreadstatus();
	fputs("Status = ", stdout);
	for (unsigned int i = 0; i < sizeof(flstattbl) / sizeof(flstattbl[0]); i++)
	    if (status & 1 << i) {
		fputs(flstattbl[i], stdout);
		fputs(" ", stdout);
	    }
	printf("(0x%02x)\n", status);
    } else if (!strcmp(argv[0], "stw")) {
	if (argc != 2) {
	    fputs("Incorrect number of arguments\n", stdout);
	    return;
	}
	tmp = atoi(argv[1]);
	flashwritestatus(tmp);
	status = flashreadstatus();
	printf("Wrote 0x%02x to status, now 0x%02x\n", tmp, status);
    } else if (!strcmp(argv[0], "er")) {
	if (argc != 2) {
	    fputs("Incorrect number of arguments\n", stdout);
	    return;
	}
	addr = atoi(argv[1]);
	flash4kerase(addr);
	printf("Erased 0x%x\n", (unsigned int)addr);
    } else if (!strcmp(argv[0], "rd")) {
	if (argc < 2) {
	    fputs("Incorrect number of arguments\n", stdout);
	    return;
	}
	
	addr = atoi(argv[1]);

	if (argc > 2)
	    len = atoi(argv[2]);
	else
	    len = 16;

	flashstartread(addr);
	
	for (uint32_t i = 0; i < len; i++)
	    printf("Read 0x%02x from 0x%06x\n", flashreadbyte(), (unsigned int)(addr + i));
	flashstopread();
	
	fputs("\n", stdout);
    } else if (!strcmp(argv[0], "wr")) {
	if (argc < 2) {
	    fputs("Incorrect number of arguments\n", stdout);
	    return;
	}

	addr = atoi(argv[1]);

	if (argc > 2)
	    len = atoi(argv[2]);
	else
	    len = 16;

	for (int i = 0; i < 16; i += 2) {
	    uint16_t data;
	    data = ((i + 1) << 8) | i;
	    printf("Writing 0x%04x to 0x%06x\n", data, (unsigned int)(addr + i));

	    if (i == 0)
		flashstartwrite(addr, data);
	    else
		flashwriteword(data);
	}
	flashstopwrite();
    } else if (!strcmp(argv[0], "id")) {
	printf("Flash ID = 0x%04hx (expect 0xbf41)\n", flashreadid());
    } else {
	fputs("Unknown sub command\n", stdout);
	return;
    }
}

void
flash4kerase(uint32_t addr) {
    flashenablewrite();			/* Enable writing */
    
    FL_SELECT();			/* Select device */

    SPI_WriteByte(FL_4KERASE);		/* Send command */
    SPI_WriteByte(addr >> 16);	 	/* Send address */
    SPI_WriteByte(addr >> 8);
    SPI_WriteByte(addr);
  
    FL_DESELECT();
    
    //fputs("4k erase ", stdout);
    flashwait();
}

void
flashwait(void) {
    uint8_t	cnt;
    
    /* Wait for not BUSY */
    for (cnt = 0; (flashreadstatus() & FL_BUSY) != 0; cnt++)
	;

    //printf("cnt = %d\n", cnt);
}

uint16_t
flashreadid(void) {
    uint8_t	fac, dev;
    
    FL_SELECT();			/* Select device */

    SPI_WriteByte(FL_RDID);		/* Send command */
    SPI_WriteByte(0x00);		/* Send address cycles (ID data starts at 0) */
    SPI_WriteByte(0x00);
    SPI_WriteByte(0x00);
    fac = SPI_WriteByte(0x00);		/* Read ID */
    dev = SPI_WriteByte(0x00);
  
    FL_DESELECT();			/* De-select device */

    return fac << 8 | dev;
}

void
flashenablewrite(void) {
    FL_SELECT();			/* Select device */

    SPI_WriteByte(FL_WREN);		/* Send command */
  
    FL_DESELECT();			/* De-select device */
}

uint8_t
flashreadstatus(void) {
    uint8_t	status;
    
    FL_SELECT();			/* Select device */

    SPI_WriteByte(FL_RDSR);		/* Send command */
    SPI_WriteByte(0x00);		/* Send dummy byte for address cycle */
    status = SPI_WriteByte(0x00);	/* Read status */
  
    FL_DESELECT();			/* De-select device */

    return status;
}

void
flashwritestatus(uint8_t status) {
    /* Enable status write */
    FL_SELECT();			/* Select device */
    SPI_WriteByte(FL_EWSR);		/* Send command */
    SPI_WriteByte(0x00);		/* Send data byte */
    FL_DESELECT();

    /* Actually write status */
    FL_SELECT();			/* Re-select device for new command */			
    SPI_WriteByte(FL_WRSR);		/* Send command */
    SPI_WriteByte(status);		/* Send data byte */
    FL_DESELECT();			/* De-select device */
}

uint8_t
flashread(uint32_t addr) {
    uint8_t	data;
    
    FL_SELECT();			/* Select device */

    SPI_WriteByte(FL_READ);		/* Send command */
    SPI_WriteByte(addr >> 16);		/* Send address */
    SPI_WriteByte(addr >> 8);
    SPI_WriteByte(addr);
    data = SPI_WriteByte(0x00);		/* Read data */
  
    FL_DESELECT();			/* De-select device */

    return data;
}

void
flashwrite(uint32_t addr, uint8_t data) {
    flashwait();
    flashenablewrite();			/* Enable writes */

    FL_SELECT();			/* Select device */

    SPI_WriteByte(FL_BYTEPROG);		/* Send command */
    SPI_WriteByte(addr >> 16);		/* Send address */
    SPI_WriteByte(addr >> 8);
    SPI_WriteByte(addr);
    SPI_WriteByte(data);		/* Write data */
  
    FL_DESELECT();			/* De-select device */

}

/*
 * fStream reading looks like so
 *
 */

void
flashstartread(uint32_t addr) {
    assert(readstate == RW_IDLE);
    
    FL_SELECT();			/* Select device */

    SPI_WriteByte(FL_READ);		/* Send command */
    SPI_WriteByte(addr >> 16); 		/* Send address */
    SPI_WriteByte(addr >> 8);
    SPI_WriteByte(addr);

    readstate = RW_RUNNING;
}

uint8_t
flashreadbyte(void) {
    assert(readstate == RW_RUNNING);
    return SPI_WriteByte(0x00);		/* Read data */
}

void
flashstopread(void) {
    assert(readstate == RW_RUNNING);

    FL_DESELECT();

    readstate = RW_IDLE;
}

/* 
 * Auto increment writing looks like so
 *
 * Enable writing  				CS, WREN, nCS
 * Send start address & first data word		CS, AAI + addr + data, nCS
 * Send subsequent words			wait for nBUSY, CS, AAI + data, nCS
 * ...
 * Disable writing				CS, WRDI, nCS
 *
 * XXX: EBSY command links SO to flash busy state, I don't think the
 * STM32 could sample it without switching out of SPI mode.
 */
void
flashstartwrite(uint32_t addr, uint16_t data) {
    assert(writestate == RW_IDLE);
    
    flashenablewrite();			/* Enable writes */

    FL_SELECT();			/* Select device */

    SPI_WriteByte(FL_AAIWP);		/* Send command */
    SPI_WriteByte(addr >> 16);
    SPI_WriteByte(addr >> 8);
    SPI_WriteByte(addr & 0xff); 	/* Send address */

    SPI_WriteByte(data & 0xff);		/* Write LSB */
    SPI_WriteByte(data >> 8);		/* Write MSB */

    FL_DESELECT();

    writestate = RW_RUNNING;
}

void
flashwriteword(uint16_t data) {
    assert(writestate == RW_RUNNING);
    
    //fputs("write word ", stdout);
    flashwait();			/* Wait until not busy */

    FL_SELECT();			/* Select device */

    SPI_WriteByte(FL_AAIWP);		/* Send command */
    SPI_WriteByte(data & 0xff);		/* Write LSB */
    SPI_WriteByte(data >> 8);		/* Write MSB */

    FL_DESELECT();			/* De-select device */
}

void
flashstopwrite(void) {
    assert(writestate == RW_RUNNING);

    //fputs("flash stop write start ", stdout);
    flashwait();			/* Wait until not busy */

    FL_SELECT();			/* Select device */

    SPI_WriteByte(FL_WRDI);		/* Send command */

    FL_DESELECT();			/* Deselect device */

    //fputs("flash stop write end ", stdout);
    flashwait();			/* Wait until not busy */

    writestate = RW_IDLE;
}

int
flashreadblock(uint32_t addr, uint32_t len, void *_data) {
    uint8_t	*data = _data;
    uint32_t	flashcrc, ramcrc;

    /* Must be a multiple of 4 due to CRC check */
    assert(len % 4 == 0);

    flashstartread(addr);
   CRC_ResetDR();
    for (int i = len; i > 0; i--) {
	*data = flashreadbyte();
	CRC_CalcCRC(*data);
	data++;
    }

    flashcrc = flashreadbyte();
    flashcrc |= flashreadbyte() << 8;
    flashcrc |= flashreadbyte() << 16;
    flashcrc |= flashreadbyte() << 24;

    flashstopread();

    ramcrc = CRC_GetCRC();
    
    /* printf("RAM CRC 0x%08x Flash CRC 0x%08x\n", (uint)ramcrc, (uint)flashcrc); */
    if (ramcrc == flashcrc)
	return 1;
    else
	return 0;
}

uint32_t
flashcrcblock(uint32_t addr, uint32_t len) {
    assert(len % 4 == 0);

   CRC_ResetDR();

    flashstartread(addr);
    for (int i = len; i > 0; i--)
	CRC_CalcCRC(flashreadbyte());

    flashstopread();

    return CRC_GetCRC();
}

int
flashwriteblock(uint32_t addr, uint32_t len, void *_data) {
    uint16_t	*data = _data;
    uint32_t	crc, vcrc;
    
    //printf("Writing %u bytes to 0x%06x\n", (uint)len, (uint)addr);
    
    /* Ensure data is
     * - 16 bit aligned
     * - a multiple of 32 bits in length (for CRCs, the flash only need 16 bits)
     * - not longer than a sector 
     */
    assert(addr % 2 == 0);
    assert(len % 4 == 0);
    assert(len <= 4096);
    
    /* Disable write protect */
    flashwritestatus(0x00);

    /* Erase sector */
    flash4kerase(addr);
 
   CRC_ResetDR();

    /* Write data */
   for (uint i = 0; i < len / 2; i++) {
	if (i == 0)
	    flashstartwrite(addr, *data);
	else
	    flashwriteword(*data);
	CRC_CalcCRC(*data);
	data++;
    }

   /* Calculate CRC */
   crc = CRC_GetCRC();
   
   //printf("CRC is 0x%08x\n", (uint)crc);
   
   /* Write CRC */
   flashwriteword(crc);
   flashwriteword(crc >> 16);
   
   flashstopwrite();

   /* Read back and check CRC */
   vcrc = flashcrcblock(addr, len);
   if (vcrc != crc)
       return 1;
   else
       return 0;
}