#include "ROMWireProtocol.h"

#include "XModem.h"

#define DATA_LEN 128
#define CONTAINER_LEN 512
#define ID_LEN 1

enum {
  CMD_INVALID,
  CMD_WRITE,
  CMD_TEST,
  CMD_DUMP,
  CMD_ERASE
};


char line[120];
byte* sendBuffer[64];
byte cmd = NULL;

XModem xmodem;

uint16_t addressPage = 0;

bool receiveBlockHandler(void *blk_id, size_t idSize, byte *data, size_t dataSize) {
  unsigned int id = *((int *) blk_id);
  //unsigned long page = (id - 1) * DATA_LEN;
  //addressPage += page;
  for (int i = 0; i < dataSize; i++) {
    unsigned long index = i + addressPage;
    programData(data[i], index);
    //programData(id, i);
  }

  addressPage += DATA_LEN;

  return true;
}

void blockLookupHandler(void *blk_id, size_t idSize, byte *send_data, size_t dataSize) {
  //uint8_t blkId = blk_id;
  //uint16_t id = *((uint16_t *) blk_id);
  // TODO: This was close to working
  uint8_t id = *((uint8_t *) blk_id);
  //unsigned long page = (id - 1) * DATA_LEN;
  //addressPage += page;
  for (int i = 0; i < dataSize; i++) {
    unsigned long index = i + addressPage;
    send_data[i] = readData(index);
  }

  addressPage += DATA_LEN;
}

void setup() {
  // Debug trigger
  pinMode(10, OUTPUT);
  digitalWrite(10, LOW);
  // put your setup code here, to run once:
  setCtrlPins();

  setAddrPinsOut();
  setDataPinsOut();
  //setDataPinsOut();

  Serial.begin(4800, SERIAL_8N1);
  xmodem.begin(Serial, XModem::ProtocolType::XMODEM);
  xmodem.setRecieveBlockHandler(receiveBlockHandler);
  xmodem.setBlockLookupHandler(blockLookupHandler);
  xmodem.setIdSize(ID_LEN);
  //xmodem.allowNonSequentailBlocks(true);
  delay(2);
  Serial.println("Started Serial COM.");
}


/*
 * Function grabbed from TommyPROM library
 * Copyright Tom Nisbet
 */
char *readLine(char *buffer, int len) {
  for (int ix = 0; ix < len; ix++) {
    buffer[ix] = 0;
  }

  char c = ' ';
  int ix = 0;
  do {
    if (Serial.available()) {
      c = Serial.read();
      if ((c == '\b') && (ix > 0)) {
        --ix;
      }
      buffer[ix++] = c;
      Serial.write(c);
    }
  } while ((c != '\n') && (c != '\r') && (ix < len));

  buffer[ix - 1] = 0;
  return buffer;
}

/*
 * Function modified from TommyPROM library
 * Copyright Tom Nisbet
 */
byte parseCommand(char c) {
  byte cmd = CMD_INVALID;

  if ((c >= 'A') && (c <= 'Z')) {
    c |= 0x20;
  }

  switch (c) {
    case 'w': cmd = CMD_WRITE; break;
    case 't': cmd = CMD_TEST; break;
    case 'd': cmd = CMD_DUMP; break;
    case 'e': cmd = CMD_ERASE; break;
    default: cmd = CMD_INVALID; break;
  }

  return cmd;
}



void loop() {
  // put your main code here, to run repeatedly:
  //programData(0x3F, 0x0000);

  //byte data = readData(0x0000);
  // Serial.print("data: ");
  //Serial.print("Data returned: ");
  //Serial.println(data, HEX);
  // Serial.println("\n");
  //delay(1);

  if (cmd == NULL) {
    Serial.print("\n>");
    readLine(line, sizeof(line));
    Serial.flush();
    cmd = parseCommand(line[0]);
  }

  switch (cmd) {
    case CMD_WRITE:
      setDataPinsOut();
      addressPage = 0;
      Serial.println("Burn new ROM: Waiting for an xmodem transfer..");
      if (xmodem.receive() == true) {
        delay(1000);
        Serial.println("Finished burning firmware.");
        cmd = NULL;
      }
      break;
    case CMD_TEST:
      Serial.println("Testing with a read/write cycle");
      programData(0x3F, 0x0000);
      Serial.println("Wrote: 0x3F");
      setDataPinsIn();
      delay(500);
      Serial.print("Received: ");
      Serial.println(readData(0x0000), HEX);
      cmd = NULL;
      break;
    case CMD_DUMP:
      //setDataPinsIn();
      //delay(500);
      Serial.println("Waiting to send dump via xmodem...");
      dumpROM();
      // for(int i = 0; i < 20; i++) {
      //   Serial.println(readData(0x1000), HEX);
      // }
      //cmd = NULL;
      break;
    case CMD_ERASE:
      eraseROM();
      cmd = NULL;
      break;

    case CMD_INVALID:
    default:
      Serial.println("NOR Flasher v1.0 Alpha");
      Serial.println();
      Serial.println("Commands: ");
      Serial.println("w\twrite new firmware");
      Serial.println("d\tdump firmware");
      Serial.println("t\tdo a test write and then read");
      Serial.println("e\terase chip");
      Serial.println();
      break;
  }
}

void dumpROM() {
  setDataPinsIn();
  addressPage = 0;

  struct XModem::bulk_data container;

  //byte *id = (byte *) malloc(2);
  byte id[CONTAINER_LEN];
  size_t data_len[CONTAINER_LEN];
  byte * data_arr[CONTAINER_LEN];

  // unsigned long long temp = 1;
  // for(size_t i = 0; i < ID_LEN; ++i) {
  //   id[ID_LEN - i - 1] = (byte *) (temp & 0xFF);
  //   temp >>= 8;
  // }

  //size_t * data_len = malloc(2);
  for(int i = 0; i < CONTAINER_LEN; i++) {
    id[i] = i + 1;
    data_len[i] = DATA_LEN;
    data_arr[i] = NULL;
  }

  container.data_arr = data_arr;
  container.len_arr = data_len;//(size_t*) 512;
  container.id_arr = id;
  container.count = CONTAINER_LEN;

  xmodem.send_bulk_data(container);

  free(id);
  free(data_len);

  delay(500);
  Serial.println("Finished sending ROM dump.");
  cmd = NULL;
  // for(int i = 0; i < 512; i++) {
  //   for(int d = 0; d < 128; d++) {
  //     sendBuffer[d] = readData(d * i);
  //   }
  //   xmodem.send(sendBuffer, 64);
  // }
  // if(xmodem.lookup_send(1) == true) {
  //   delay(1000);
  //   Serial.println("Finished transferring firmware dump.");
  //   cmd = NULL;
  // }
  // for (int i = 0; i < 10; i++) {
  //   byte data = readData(0x0000 | i);
  //   //Serial.println(data, HEX);
  // }
}

void eraseROM() {
  setDataPinsOut();

  writeByte(0xAA, 0x5555);
  writeByte(0x55, 0x2AAA);
  writeByte(0x80, 0x5555);
  writeByte(0xAA, 0x5555);
  writeByte(0x55, 0x2AAA);
  writeByte(0x10, 0x5555);

  delay(100);
  setDataPinsIn();
}