#include <SPI.h>
#include <GD.h>

static PROGMEM prog_uint32_t crc_table[16] = {
    0x00000000, 0x1db71064, 0x3b6e20c8, 0x26d930ac,
    0x76dc4190, 0x6b6b51f4, 0x4db26158, 0x5005713c,
    0xedb88320, 0xf00f9344, 0xd6d6a3e8, 0xcb61b38c,
    0x9b64c2b0, 0x86d3d2d4, 0xa00ae278, 0xbdbdf21c
};

unsigned long crc_update(unsigned long crc, byte data)
{
    byte tbl_idx;
    tbl_idx = crc ^ (data >> (0 * 4));
    crc = pgm_read_dword_near(crc_table + (tbl_idx & 0x0f)) ^ (crc >> 4);
    tbl_idx = crc ^ (data >> (1 * 4));
    crc = pgm_read_dword_near(crc_table + (tbl_idx & 0x0f)) ^ (crc >> 4);

    return crc;
}

#define FLOCAL 2
#define FREMOTE 10

#define SEL_local() digitalWrite(FLOCAL, LOW)
#define UNSEL_local() digitalWrite(FLOCAL, HIGH)
#define SEL_remote() digitalWrite(FREMOTE, LOW)
#define UNSEL_remote() digitalWrite(FREMOTE, HIGH)

#define spix(n) SPI.transfer(n)

static void spipage(int n)
{
  spix(n >> 7);
  spix(n << 1);
  spix(0);
}

static unsigned long pagecrc(int n)
{
  SPI.transfer(0x03);
  spipage(n);

  unsigned long crc = ~0;
  int len = 264;
  while (len--) {
    crc = crc_update(crc, SPI.transfer(0));
  }
  crc = ~crc;
  return crc;
}

static unsigned long local_pagecrc(int n)
{
  SEL_local();
  unsigned long r = pagecrc(n);
  UNSEL_local();
  return r;
}

static unsigned long remote_pagecrc(int n)
{
  SEL_remote();
  unsigned long r = pagecrc(n);
  UNSEL_remote();
  return r;
}

static byte remote_status()
{
  SEL_remote();
  spix(0xd7);   // read SPI flash status
  byte status = spix(0);
  UNSEL_remote();
  return status;
}

static void UNSEL_remote_wait()
{
  UNSEL_remote();
  while ((remote_status() & 0x80) == 0)
    ;
}

static void pgcmd(byte cmd, int page)
{
  SEL_local();
  spix(cmd);
  spipage(page);
}

static void pgcmd_remote(byte cmd, int page)
{
  SEL_remote();
  spix(cmd);
  spipage(page);
}

static byte img[4 * 264];

static void failmsg(char *msg)
{
  for (;;) {
    GD.waitvblank();
    GD.wr16(BG_COLOR, RGB(32, 0, 0));
    GD.putstr(0, 18, msg);
    delay(500);
    GD.waitvblank();
    GD.wr16(BG_COLOR, RGB(0, 0, 0));
    GD.putstr(0, 18, "                                                  ");
    delay(100);
  }
}

static void write_remote(byte *src, int pg, int pages)
{
  int i;
  while (pages--) {
    pgcmd_remote(0x84, pg);
    for (i = 0; i < 264; i++)
      //SPDR = *src++;
      spix(*src++);
    UNSEL_remote();

    pgcmd_remote(0x88, pg);
    UNSEL_remote_wait();

    pgcmd_remote(0x60, pg);
    UNSEL_remote_wait();
    if (remote_status() & (1 << 6))
      failmsg("Compare failed");

    pg++;
  }
}

static void read_spi(byte *dst, int sz)
{
#if 0
    byte *lim = dst + sz - 1;
    byte v = spix(0);
    while (dst != lim) {
      SPDR = 0;
      *dst++ = v;
      v = SPDR;
    }
    *dst = v;
#else
  while (sz--)
    *dst++ = spix(0);
#endif
}

static unsigned long clone(int pg)
{
  pgcmd_remote(0x50, pg);  // erase
  UNSEL_remote();

  pgcmd(0x03, pg);
  read_spi(img, 4 * 264);
  UNSEL_local();
  UNSEL_remote_wait();
  write_remote(img, pg, 4);

  pgcmd(0x03, pg + 4);
  read_spi(img, 4 * 264);
  UNSEL_local();
  write_remote(img, pg + 4, 4);
}

#define PAGES 600   // only first part of flash matters

static unsigned long master_sum;

static unsigned long sum_pages()
{
  int i;
  unsigned long sum = 0;

#if 0
  byte p = spix(0);
  for (i = 0; i < (264 * (PAGES / 8)) - 1; i++) {
    SPDR = 0; sum += p; p = SPDR;
    SPDR = 0; sum += p; p = SPDR;
    SPDR = 0; sum += p; p = SPDR;
    SPDR = 0; sum += p; p = SPDR;
    SPDR = 0; sum += p; p = SPDR;
    SPDR = 0; sum += p; p = SPDR;
    SPDR = 0; sum += p; p = SPDR;
    SPDR = 0; sum += p; p = SPDR;
  }
  sum += p;
#else
  for (i = 0; i < (264 * 8 * PAGES); i++)
    sum += spix(0);
#endif
  return sum;
}

#define BUTTON_PIN 7

void setup()
{
  pinMode(BUTTON_PIN, INPUT);
  digitalWrite(BUTTON_PIN, HIGH);

  Serial.begin(1000000);  // JCB

  GD.begin();
  GD.ascii();
  GD.wr(IOMODE, 'F');
  GD.putstr(0, 1, "Clone utility - factory use only");

  pinMode(FLOCAL, OUTPUT);
  pinMode(FREMOTE, OUTPUT);

  UNSEL_local();
  UNSEL_remote();

  {
    pgcmd(0x03, 0);
    unsigned long sum = sum_pages();
    UNSEL_local();
    master_sum = sum;
  }
  GD.putstr(0, 18, "Press button to start...");
}

static int count = 0;

void loop()
{
  char msg[50];

  while (digitalRead(BUTTON_PIN))
    ;

  long t0 = millis();
  GD.fill(0, ' ', 4096);
  GD.putstr(0, 1, "Clone utility - factory use only");
  sprintf(msg, "%4d", count);
  GD.putstr(44, 1, msg);

  // wait for flash on REMOTE
  byte status;
  do {
    status = remote_status();
    
    sprintf(msg, "flash status %02x", status);
    GD.putstr(0, 3, msg);
  } while (((status & 0xc1) != 0x80) && ((status & 0x3c) == 0));

  GD.putstr(0, 5, "cloning");

  // disable sector protection
  SEL_remote();
  spix(0x3d);
  spix(0x2a);
  spix(0x7f);
  spix(0x9a);
  UNSEL_remote();

  // copy all pages from LOCAL to REMOTE
  int p;
  for (p = 0; p < PAGES; p += 8) {
    byte x = (p * 40 / PAGES);
    GD.wr(64 * 5 + 10 + x, 127);

    clone(p);
  }

  pgcmd_remote(0x82, 1000);
  img[0] = lowByte(count);
  img[1] = highByte(count);
  img[2] = lowByte(t0);
  img[3] = highByte(t0);
  unsigned short t0u = t0 >> 16;
  img[4] = lowByte(t0u);
  img[5] = highByte(t0u);
  int i;
  for (i = 0; i < 264; i++)
    spix(img[i]);
  UNSEL_remote_wait();

  pgcmd_remote(0x03, 0);
  unsigned long sum = sum_pages();
  UNSEL_remote();
  if (sum != master_sum) {
    failmsg("Verify failed");
  }

  GD.putstr(0, 9, "SUCCESS");
  count++;
  // GD.screenshot(0);
  GD.putstr(0, 18, "Change flash then press button...");
  Serial.print("took "); Serial.println(millis() - t0, DEC);  // JCB
}
