#include "Adafruit_Keypad.h"

const byte KBD_ROWS = 3; // rows
const byte KBD_COLS = 9; // columns

char keys[KBD_ROWS][KBD_COLS] = {
  {'Q','W','E','R','T','Z','U','I','O'},
  {'A','S','D','F','G','H','J','K','*'},
  {'P','Y','X','C','V','B','N','M','L'}
};

byte rowPins[KBD_ROWS] = {10,11,12}; //connect to the row pinouts of the keypad
byte colPins[KBD_COLS] = {13,2,3,4,5,6,7,8,9}; //connect to the column pinouts of the keypad

Adafruit_Keypad kbd = Adafruit_Keypad( makeKeymap(keys), rowPins, colPins, KBD_ROWS, KBD_COLS);

#include <LiquidCrystal.h>
LiquidCrystal lcd(19, 18, 17, 16, 15, 14);

const int LCD_COLS = 20;
const int LCD_ROWS = 4;

bool debug_mode = false;

typedef struct {
  char *id;
  char *name;
  char wiring[26];
  char *notch;
} Rotor;

Rotor rotors [] = {
  // Enigma I, M3, M4
  // id            name          wiring                        notch
  // ==            ====          ======                        =====
  {  "ABC",        "ABC",        "ABCDEFGHIJKLMNOPQRSTUVWXYZ", ""   },  //  0 alphabet

  // 
  {  "I",          "I",          "EKMFLGDQVZNTOWYHXUSPAIBRCJ", "Q"  },  //  1 
  {  "II",         "II",         "AJDKSIRUXBLHWTMCQGZNPYFVOE", "E"  },  //  2
  {  "III",        "III",        "BDFHJLCPRTXVZNYEIWGAKMUSQO", "V"  },  //  3
  {  "IV",         "IV",         "ESOVPZJAYQUIRHXLNFTGKDCMWB", "J"  },  //  4
  {  "V",          "V",          "VZBRGITYUPSDNHLXAWMJQOFECK", "Z"  },  //  5

  // 
  {  "VI",         "VI",         "JPGVOUMFYQBENHZRDKASXLICTW", "ZM" },  //  6
  {  "VII",        "VII",        "NZJHGRCXMYSWBOUFAIVLPEKQDT", "ZM" },  //  7
  {  "VIII",       "VIII",       "FKQHTLXOCBJSPDZRAMEWNIUYGV", "ZM" },  //  8
  {  "beta",       "Beta",       "LEYJVCNIXWPBQMDRTAKZGFUHOS", ""   },  //  9  M4 only
  {  "gamma",      "Gamma",      "FSOKANUERHMBTIYCWLQPZXVGJD", ""   },  // 10  M4 only

  // reflectors
  {  "UKW-A",      "UKW A",      "EJMZALYXVBWFCRQUONTSPIKHGD", ""   },  // 11
  {  "UKW-B",      "UKW B",      "YRUHQSLDPXNGOKMIEBFZCWVJAT", ""   },  // 12
  {  "UKW-C",      "UKW C",      "FVPJIAOYEDRZXWGCTKUQSBNMHL", ""   },  // 13
  {  "UKW-B-thin", "UKW B thin", "ENKQAUYWJICOPBLMDXZVFTHRGS", ""   },  // 14  M4 only
  {  "UKW-C-thin", "UKW C thin", "RDOBJNTKVEHMLFCWZAXGYIPSUQ", ""   }   // 15  M4 only
};


typedef struct {
  char *name;
  char *description;
  Rotor rotor_types[];
  Rotor reflector_types[];
} EnigmaType;

EnigmaType enigma_types[] = {
  { 
    "M3",
    "Army",
    { rotors[1], rotors[2], rotors[3], rotors[4], rotors[5] },
    { rotors[11], rotors[12] }
  },
  
  { 
    "M4",
    "Navy",
    // work out thin rotors
    { rotors[1], rotors[2], rotors[3], rotors[4], rotors[5], rotors[6], rotors[7], rotors[8] },
    { rotors[11], rotors[12] }
  }
}; 

// entry wheel "EintrittswalzeEintrittswalze"
byte entry_wheel_type = 0;  // 

// rotor types
byte rotor_type [] = { 1, 2, 3 };

// reflector "UmkehrwalzeUmkehrwalze"
byte reflector_type = 12;
const int PLUGBOARD_SIZE = 10;
char plugboard [PLUGBOARD_SIZE][2] = { 
    {'A','B'}, {'.','.'}, {'.','.'}, {'.','.'}, {'R','Z'},
    {'C','D'}, {'.','.'}, {'.','.'}, {'.','.'}, {'G','H'}
  };

// track rotor positions
byte rotor_pos [] = { 0, 0, 0 };

#define ROTOR_COUNT (sizeof(rotor_type)/sizeof(rotor_type[0]))

// compute alphabet size
#define AZ (rotors[entry_wheel_type].wiring)
#define AZ_LEN sizeof(rotors[0].wiring)
#define AZ_MAX_INDEX (AZ_LEN-1)

int wrap_pos (int pos) {
  if (pos<0) {
    return wrap_pos(pos+AZ_LEN);  // TODO: optimization opportunity
  } else {
    return pos % AZ_LEN;
  }
}

void rotate (byte rotor) {
  rotor_pos[rotor] = wrap_pos(rotor_pos[rotor]+1);
}

bool at_notch (byte rotor) {
  return NULL != strchr(rotors[rotor_type[rotor]].notch, AZ[rotor_pos[rotor]]);
}

void print_rotor (int rotor) {
  Serial.print(at_notch(rotor) ? '(' : '[');
  Serial.print(AZ[rotor_pos[rotor]]);
  Serial.print(at_notch(rotor) ? ')' : ']');
}

void print_rotors () {
  for (int i=ROTOR_COUNT-1; i>=0; --i) {
    print_rotor(i);
  }
}

int stridx(char *str, char c) {
  return strchr(str, c) - str;
}

char encode (char c) {
  bool
    rotated1 = false,
    rotated2 = false;

  print_rotors();

  // rotate left rotor if middle rotor is at notch
  if (at_notch(1)) {
    rotated2 = true;
    rotate(2);
  }

  // rotate middle rotor only if both right and middle are at notch
  if (at_notch(0) || at_notch(1)) {
    rotated1 = true;
    rotate(1);
  }

  // always rotate the right rotor
  rotate(0);

  if (debug_mode) {
    // did any of the rotors (except for the rightmost) rotate?
    Serial.print(rotated2? '>' : '-');
    Serial.print(rotated1? '>' : '-');
    Serial.print('>');
  }
  
  print_rotors();

  if (debug_mode)
    Serial.print(' ');

  // going from rightmost rotor to the left
  int pin = stridx(AZ, c);
  int pos;
  char *wiring;
  for (int rotor = 0; rotor < ROTOR_COUNT; ++rotor) {
    pos = rotor_pos[rotor];
    wiring = rotors[rotor_type[rotor]].wiring;
    if (debug_mode)
      Serial.print(AZ[pin]); 

    pin = wrap_pos(stridx(AZ, wiring[wrap_pos(pin + pos)]) - pos);
    if (debug_mode) {
      Serial.print("->"); 
      Serial.print(AZ[pin]); 
      Serial.print(" "); 
    }
  }

  // reflector
  if (debug_mode) {
    Serial.print(" "); 
    Serial.print(AZ[pin]);
  }
  pin = stridx(AZ, rotors[reflector_type].wiring[pin]);
  if (debug_mode) {
    Serial.print("->");
    Serial.print(AZ[pin]);
    Serial.print("  "); 
  }
  
  // returning back from leftmost to the right
  for (int rotor = ROTOR_COUNT-1; rotor >= 0; --rotor) {
    if (debug_mode)
      Serial.print(AZ[pin]);
    pos = rotor_pos[rotor];
    wiring = rotors[rotor_type[rotor]].wiring;
    pin = wrap_pos(pin + pos);
    char ch = AZ[pin];
    ch = AZ[stridx(wiring, ch)];
    pin = stridx(AZ, ch);
    pin = wrap_pos(pin - pos);
    if (debug_mode) {
      Serial.print("->"); 
      Serial.print(AZ[pin]);
      Serial.print(" ");
    }
  }

  char plug_c = AZ[pin];

  // plugboard
  for (int i=0; i<PLUGBOARD_SIZE; ++i) {
    if (AZ[pin] == plugboard[i][0]) {
      return plugboard[i][1];
    }
  }

  return AZ[pin];
}

const char *plain  = "ABCDEFGHIJKLMNOPQRSTUVWXYZABCDEFGHIJKLMNOPQRSTUVWXYZABCDEFGHIJKLMNOPQRSTUVWXYZABCDEFGHIJKLMNOPQRSTUVWXYZABCDEFGHIJKLMNOPQRSTUVWXYZABCDEFGHIJKLMNOPQRSTUVWXYZ";
const char *cipher = "FUVEPUMWARVQKEFGHGDIJFMFXIMRENATHDMCEVOQHIUWXXGYSJADEGKHYJETLBLWVZNUXFNSPICQFGZCZJKYWLLGPXJKBYTNNEFYKQTCJOLGCWHGXUEYOQXDNIGIDEMBXACVPAVYUQCGPXILERRSJSBOOKJW";
int const PLAIN_LEN = strlen(plain);

void test () {
  Serial.begin(115200);
  Serial.print("Testing "); Serial.print(PLAIN_LEN); Serial.println(" chars.");

  bool failed = false;

  debug_mode = true;

  for (byte i=0; i<PLAIN_LEN; ++i) {
    if (i<10) Serial.print(' ');
    if (i<100) Serial.print(' ');
    Serial.print(i, DEC);
    Serial.print(": ");
    char encoded = encode(plain[i]);
    if (cipher[i] != encoded) {
      failed = true;
      Serial.print(" (Expected: ");
      Serial.print(cipher[i]); 
      Serial.print(")");
    }
    Serial.println("");
  }

  if (failed) {
    Serial.println("\n"
      "\n* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"
      "\nOne or more characters did not encrypt correctly. See messages above."
      "\n* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"
    );
  }
}


const int BUF_SIZE = LCD_COLS - ROTOR_COUNT*3 - 1;
char input_buf [BUF_SIZE];
char encoded_buf [BUF_SIZE];
int buf_pos = 0;

void init_buf (char *buf) {
  for (int i=0; i<BUF_SIZE; ++i) {
    buf[i] = ' ';
  }
}

void add_to_buf (char *buf, char c) {
  for (int i=0; i<BUF_SIZE-1; ++i) {
    buf[i] = buf[i+1];
  }
  buf[BUF_SIZE-1] = c;
}

void disp_buf (char *buf, int lcd_row) {
  lcd.setCursor(LCD_COLS-BUF_SIZE,lcd_row);
  for (int i=0; i<BUF_SIZE; ++i) {
    lcd.print(buf[i]);
  }
}

void disp_rotor (int rotor) {
  lcd.setCursor((ROTOR_COUNT-rotor-1)*3,0);
  lcd.print(at_notch(rotor) ? '(' : '[');
  lcd.print(AZ[rotor_pos[rotor]]);
  lcd.print(at_notch(rotor) ? ')' : ']');
}

void disp_rotors () {
  for (int i=ROTOR_COUNT-1; i>=0; --i) {
    disp_rotor(i);
  }
}

void disp_plugboard () {
  if (LCD_ROWS > 2) {
    for (int i=0; i<PLUGBOARD_SIZE/2; ++i) {
      lcd.setCursor(i*3+(LCD_COLS-PLUGBOARD_SIZE/2*3+1),2);
      lcd.print(plugboard[i][0]);
      lcd.print(plugboard[i][1]);
      if (i<PLUGBOARD_SIZE/2-1) {
        lcd.print(' ');
      }
      lcd.setCursor(i*3+(LCD_COLS-PLUGBOARD_SIZE/2*3+1),3);
      lcd.print(plugboard[i+PLUGBOARD_SIZE/2][0]);
      lcd.print(plugboard[i+PLUGBOARD_SIZE/2][1]);
      if (i<PLUGBOARD_SIZE/2-1) {
        lcd.print(' ');
      }
    }
  }
}

// config:

// reset
// enigma type
// for each rotor
// - ringsetting
// - rotor selection
// plugboard
// - clear
// - set pair

void setup () {
//  Serial.begin(115200);
  kbd.begin();
  lcd.begin(LCD_COLS, LCD_ROWS);

  lcd.setCursor(0,1);
  lcd.print("Enigma");
  lcd.setCursor(0,2);
  lcd.print("M3");

  test();
  init_buf(input_buf);
  init_buf(encoded_buf);
  disp_rotors();
  disp_plugboard();
}


bool demo_mode = false;
bool config_mode = true;

void enigma () {
  kbd.tick();

  while (kbd.available()) {
    keypadEvent e = kbd.read();
    if (e.bit.EVENT == KEY_JUST_PRESSED) {
      char plain_c = (char)e.bit.KEY;
      char encoded_c = encode(plain_c);

      disp_rotors();

      add_to_buf(input_buf, plain_c);
      add_to_buf(encoded_buf, encoded_c);
      disp_buf(input_buf, 0);
      disp_buf(encoded_buf, 1);
    }
  }

  delay(10);
}

void demo () {
  
}

enum ConfigOptionType { YesNo, EnigmaList, RotorList, ThinRotorList, Plugboard, Scalar };

typedef struct {
  char *prompt;
  ConfigOptionType type;
  void *target;
  void (*skip) ();
  void (*blink) ();
} ConfigOption;

EnigmaType enigma_type;

ConfigOption config_options[] = {
  { "Demo?", YesNo, &demo_mode, NULL, NULL },
  { "Enigma Type?", EnigmaList, &enigma_type, NULL, NULL },
  { "Right Rotor?", RotorList, &rotor_type[0], NULL, NULL },
  { "Middle Rotor?", RotorList, &rotor_type[1], NULL, NULL },
  { "Left Rotor?", RotorList, &rotor_type[2], NULL, NULL },
  { "Thin Rotor?", ThinRotorList, &rotor_type[2], NULL, NULL },
  { "Plugboard?", Plugboard, &plugboard, NULL, NULL }
};

int config_option_index = 0;

void config () {
  
  // demo? y/n -> demo_mode
  // enigma type? menu -> enigma_type
  // rotor[
}

void loop () {
  if (config_mode) {
    config();
  } else if (demo_mode) {
    demo();
  } else {
    enigma();
  }
}


/*
  Eintrittwalze entry wheel (stator)
  Greek wheels Beta and Gamma [for obvious reasons] used in the M4
  Grundstellung ground setting
  Ringstellung  ring setting
  Stecker plug
  Steckerbrett  plugboard [sometimes referred to as "Stecker board"]
  Umkehrwalze reflector
  Walzen  wheels
  Walzenlage  wheel order
 */

 /*
https://cryptocellar.org/Enigma/index.html
https://cryptocellar.org
http://www.mlb.co.jp/linux/science/genigma/enigma-referat/node4.html#SECTION00043000000000000000
http://www.macs.hw.ac.uk/~foss/valentin/Naval%20Enigma.html
https://github.com/melick/enigma
https://www.codesandciphers.org.uk/enigma/example1.htm
https://www.codesandciphers.org.uk/enigma/index.htm
https://www.codesandciphers.org.uk/enigma/enigma2.htm
  */


/*

Codebook

+------------------------------------------------+
| 30 | I V IV    | VAX | AX CE DK IM TG WF | LQT |
+------------------------------------------------+
| 29 | IV V II   | JNF | AL KI MT NB UH ZR | SYO |
+------------------------------------------------+
| 28 | III IV II | FHT | AW IS JR KB YH ZX | TJE |
+------------------------------------------------+
| 27 | V III I   | OQY | AZ CV DN FJ TU WP | IIF |
+------------------------------------------------+
| 26 | IV I V    | GRT | BG ER IJ LF OC ZP | GST |
+------------------------------------------------+
| 25 | V II IV   | HCE | BR CJ QX UF YO ZT | OJT |
+------------------------------------------------+
| 24 | III V IV  | FAK | IQ RH UZ WM XJ YC | BKW |
+------------------------------------------------+
| 23 | II III I  | QMD | CQ IH KY RW VL XT | QSU |
+------------------------------------------------+
| 22 | II IV III | JSI | AN FH JV ML RO WD | KNQ |
+------------------------------------------------+
| 21 | II IV I   | JMY | AP BM GD NS QV RC | HWO |
+------------------------------------------------+
| 20 | IV I III  | GMX | AB TD WE XH YP ZF | BOD |
+------------------------------------------------+
| 19 | V III II  | COA | FR IT PE QD SH XZ | EGM |
+------------------------------------------------+
| 18 | V II I    | VTU | FJ GR KC LS NX WA | MTJ |
+------------------------------------------------+
| 17 | II I III  | KKX | BV CO LG MK SY WN | JLY |
+------------------------------------------------+
| 16 | II I III  | MJO | GQ IL KF OD PS UC | NBV |
+------------------------------------------------+
| 15 | II V IV   | MFX | CL DS HK JX MW YN | IEU |
+------------------------------------------------+
| 14 | I V II    | TTN | DA IH JS KU TP WZ | MYF |
+------------------------------------------------+
| 13 | IV V I    | TEV | AP CM LB NZ UI XS | KJP |
+------------------------------------------------+
| 12 | III II IV | RXS | BO CL DY GR SI WA | EAF |
+------------------------------------------------+
| 11 | II III I  | MXN | CZ KB ME NL QP TS | BLW |
+------------------------------------------------+
| 10 | I IV II   | OHP | FU GR LJ OA QX VI | LEB |
+------------------------------------------------+
|  9 | III I V   | PQY | AG MC NU QI TX ZV | NWT |
+------------------------------------------------+
|  8 | III I V   | CJN | AN EC OK UI WT YB | BTD |
+------------------------------------------------+
|  7 | II IV I   | HJI | CM HX KW SJ TI UQ | MPD |
+------------------------------------------------+
|  6 | II I III  | HID | AP BM GK TY VJ XI | JIJ |
+------------------------------------------------+
|  5 | II V I    | ESJ | FV JK OT PY SE ZH | XMB |
+------------------------------------------------+
|  4 | IV I III  | JJN | AE NH OU VR XZ YQ | WTV |
+------------------------------------------------+
|  3 | V II I    | MPB | BU LS ND PG TA XM | OFS |
+------------------------------------------------+
|  2 | III II V  | NXR | BR GV HN MD WL ZI | KSU |
+------------------------------------------------+
|  1 | I IV V    | PIE | MC OG SH VI WL YQ | RBU |
+------------------------------------------------+

 */