#include <SPI.h>
#include "FS.h"
#include "SPIFFS.h"

#define CS 5
#define FORMAT false
SPIClass * vspi = NULL;

SPISettings settings(20000000, MSBFIRST, SPI_MODE0);
//SPISettings settings(1000000, MSBFIRST, SPI_MODE0);

byte jedec[3];

void readBuffer(uint32_t addr, uint8_t *buff, uint32_t len) {
  uint8_t ad[3];
  uint8_t out;
  ad[0] = (addr & 0xFF0000) >> 16;
  ad[1] = (addr & 0x00FF00) >> 8;
  ad[2] = (addr & 0x0000FF) >> 0;

  digitalWrite(CS, LOW);
  vspi->beginTransaction(settings);
  vspi->transfer(0x03); // Read Data command
  vspi->transfer(ad[0]);
  vspi->transfer(ad[1]);
  vspi->transfer(ad[2]);
  for (uint32_t i = 0; i < len; i++) {
    *buff = vspi->transfer(0x00);
    buff++;
  }
  vspi->endTransaction();
  digitalWrite(CS, HIGH);
}

void printHexPad32(uint32_t val) {
  for (int i = 7; i > 0; i--) {
    uint32_t v = (1 << (i * 4));
    if (val < v) {
      Serial.print("0");
    } else {
      break;
    }
  }
  Serial.print(val, HEX);
}

void printHexPad8(uint8_t val) {
  if (val < 0x10) {
    Serial.print("0");
  }
  Serial.print(val, HEX);
}

void printASCII(uint8_t v) {
  if (v >= 0x20 && v <= 0x7E) {
    Serial.print((char)v);
  } else {
    Serial.print(".");
  }
}

void printBuff(uint8_t *buff, uint32_t len, uint32_t mainoffset) {
  Serial.println("-----------------------------");
  for (uint32_t i = 0; i < len; i++) {
    if (i % 16 == 0) {
      printHexPad32(mainoffset + i);
      Serial.print(": ");
    }
    printHexPad8(buff[i]);
    Serial.print(" ");
    if (i % 16 == 15) {
      uint32_t off = (i / 16) * 16;
      for (uint32_t p = 0; p < 16; p++) {
        printASCII(buff[off+p]);
      }
      Serial.println("");
    }
  }
  Serial.println("-----------------------------");
}

#define CHUNK 256

uint8_t databuff[CHUNK];
uint8_t pattern[CHUNK];

uint32_t flashSize = 512 * 1024;

uint8_t readStatus() {
  uint8_t status;
  digitalWrite(CS, LOW);
  vspi->beginTransaction(settings);
  vspi->transfer(0x05);
  status = vspi->transfer(0);
  vspi->endTransaction();
  digitalWrite(CS, HIGH);
  return status;
}

void waitReady() {
  uint8_t status = 1;
//  Serial.println("Waiting device not BUSY");
  while (status) {
    delay(1);
    status = readStatus();
    status &= 1; // Only first bit
  }
//  Serial.println("Device is free!");
}

void sectorErase(uint32_t addr) {
  uint8_t ad[3];
  uint8_t out;
  ad[0] = (addr & 0xFF0000) >> 16;
  ad[1] = (addr & 0x00FF00) >> 8;
  ad[2] = (addr & 0x0000FF) >> 0;

  writeEnable();
  digitalWrite(CS, LOW);
  vspi->beginTransaction(settings);
  vspi->transfer(0x20); // Read Data command
  vspi->transfer(ad[0]);
  vspi->transfer(ad[1]);
  vspi->transfer(ad[2]);
  vspi->endTransaction();
  digitalWrite(CS, HIGH);
  waitReady();
}

#define SRP (1 << 7)
#define TB  (1 << 5)
#define BP2 (1 << 4)
#define BP1 (1 << 3)
#define BP0 (1 << 2)
#define WEL (1 << 1)

void unprotect() {
  writeEnable();
  uint8_t ok = 0;
  while(!ok) {
    uint8_t status = readStatus();
    status &= ~(SRP | TB | BP2 | BP1 | BP0); // Reset protections
    status &= ~(WEL); // Avoid checking Write Enable that is reseted after operation
    
    digitalWrite(CS, HIGH);
    delay(1);
    digitalWrite(CS, LOW);
    vspi->beginTransaction(settings);
    vspi->transfer(0x01); 
    vspi->transfer(status);
    vspi->endTransaction();
    digitalWrite(CS, HIGH);
    delay(1);
    waitReady();
  
    uint8_t status2 = readStatus();
    if (status != status2) {
      Serial.print("Error writting status! Expected ");
      Serial.print(status, HEX);
      Serial.print(" got ");
      Serial.println(status2, HEX);
    } else {
      ok = 1;
    }
  }
}

void writeEnable() {
//  Serial.println("Write Enable");
  digitalWrite(CS, LOW);
  vspi->beginTransaction(settings);
  vspi->transfer(0x06); 
  vspi->endTransaction();
  digitalWrite(CS, HIGH);
  delay(10);
  uint8_t status = readStatus();
  status &= 2;
  if (status != 2) {
    Serial.println("ERROR setting Write Enable!");
  }
  waitReady();
}

void chipErase() {
  writeEnable();
  Serial.println("Chip Erase");
  digitalWrite(CS, LOW);
  vspi->beginTransaction(settings);
  vspi->transfer(0xC7); 
  vspi->endTransaction();
  digitalWrite(CS, HIGH);
  delay(10);
  waitReady();
}

void programPage(uint32_t addr, uint8_t *data) {
  uint8_t ad[3];
  ad[0] = (addr & 0xFF0000) >> 16;
  ad[1] = (addr & 0x00FF00) >> 8;
  ad[2] = (addr & 0x0000FF) >> 0;
  
  writeEnable();
  digitalWrite(CS, LOW);
  vspi->beginTransaction(settings);
  vspi->transfer(0x02);
  vspi->transfer(ad[0]);
  vspi->transfer(ad[1]);
  vspi->transfer(ad[2]);
  for (int i = 0; i < 256; i++) {
    vspi->transfer(data[i]); 
  }
  vspi->endTransaction();
  digitalWrite(CS, HIGH);
  delay(1);
  waitReady();
}

char *wololo = "ACNDEFHEJRHSUDIFKEMLOPEC";

void TestChip() {
  Serial.println("Unprotect");
  unprotect();
//  Serial.println("Erasing chip");
//  chipErase();
  Serial.println("Programming test buffer");
  for (int i = 0; i < CHUNK; i++) {
    pattern[i] = (i & 0xFF);
  }
  int numPages = flashSize / CHUNK;
  for (int i = 0; i < numPages; i++) {
    Serial.print("Testing page ");
    Serial.println(i, HEX);
    sectorErase(i*CHUNK);
    programPage(i*CHUNK, pattern);
    readBuffer(i * CHUNK, databuff, CHUNK);
    int diff = compare(databuff, pattern, 256);
    if (diff != -1) {
      Serial.print("Expected at offset ");
      printHexPad32(i * CHUNK);
      Serial.println("");
      printBuff(pattern, CHUNK, i * CHUNK);
      Serial.println("");
      Serial.println("Got: ");
      printBuff(databuff, CHUNK, i * CHUNK);
      break;
    }
  }
  Serial.println("FINISH");
}

void setup() {
  Serial.begin(115200);
  Serial.setDebugOutput(true);
  Serial.println("ON!!!");
  
  pinMode(23, OUTPUT);
  pinMode(19, INPUT);
  pinMode(18,OUTPUT);
  pinMode(CS, OUTPUT);
  pinMode(13, OUTPUT);
  pinMode(22, OUTPUT);
  digitalWrite(CS, HIGH);
  digitalWrite(13, HIGH);
  digitalWrite(22, HIGH);

  if(FORMAT) {
    SPIFFS.format();
  }
  
  if(!SPIFFS.begin(true)){
      Serial.println("SPIFFS Mount Failed");
      return;
  }

  memcpy(databuff, wololo, 24);
  
  vspi = new SPIClass(VSPI);
  vspi->begin();
  vspi->setDataMode(0);
  vspi->setBitOrder(MSBFIRST);

  
  delay(100);

  uint32_t offset = 0;
  Serial.println("Send anything to start.");
  while(Serial.available() <= 0);

  TestChip();
//  unprotect();
//  chipErase();
//  programPage(0, databuff);
//  PrintFlash();
//  DumpFlash();
//  WriteFlash();
//  VerifyFlash();
}

int compare(uint8_t *buff0, uint8_t *buff1, uint32_t len) {
  for (uint32_t i = 0; i < len; i++) {
    if (buff0[i] != buff1[i]) {
      return i;
    }
  }

  return -1;
}

void VerifyFlash() {
  File file = SPIFFS.open("/FLSHSAVE.BIN", FILE_READ);

  if (!file) {
    Serial.println("Error opening file!");
    while(true) {
      delay(1000);
    }
  }
  
  uint32_t addr = 0;
  int diff = -1;
  Serial.println("VERIFING");
  while(addr < flashSize) {
//    printHexPad32(addr);
//    Serial.println("");
    file.read(databuff, 256);
    readBuffer(addr, databuff + 512, 256);
    diff = compare(databuff, databuff + 512, 256);
    if (diff != -1) {
      Serial.print("\nCorrupted at index ");
      Serial.print(addr + diff, HEX);
      Serial.println("");
      break;
    }
    addr += 256;
  }
 
  file.close();
  Serial.println("DONE");
}

void WriteFlash() {
  unprotect();
  chipErase();
  
  File file = SPIFFS.open("/FLSHSAVE.BIN", FILE_READ);

  if (!file) {
    Serial.println("Error opening file!");
    while(true) {
      delay(1000);
    }
  }
  
  uint32_t addr = 0;
  Serial.println("FLASHING");

  while(addr < flashSize) {
    printHexPad32(addr);
    Serial.println("");
    file.read(databuff, 256);
    programPage(addr, databuff);
    addr += 256;
  }

  file.close();
  Serial.println("DONE");
}


void PrintFlash() {
  uint32_t offset = 0;
  
  Serial.println("READING FLASH");
  while (offset < flashSize) {
    readBuffer(offset, databuff, CHUNK);
    printBuff(databuff, CHUNK, offset);
    offset += CHUNK;
    break;
  }
}

void DumpFlash() {
  File file = SPIFFS.open("/SAVE.BIN", FILE_WRITE);
  if (!file) {
    Serial.println("Error opening file!");
    while(true) {
      delay(1000);
    }
  }

  uint32_t offset = 0;
  
  Serial.println("READING FLASH");
  while (offset < flashSize) {
    readBuffer(offset, databuff, CHUNK);
    file.write(databuff, CHUNK);
    printHexPad32(offset);
    printBuff(databuff, CHUNK, offset);
    Serial.println("");
    offset += CHUNK;
  }
  file.close();
}

void loop() {
  delay(1000);
}