#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 = 16;
const int LCD_ROWS = 2;

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
  {  "gamma",      "Gamma",      "FSOKANUERHMBTIYCWLQPZXVGJD", ""   },  // 10
  {  "UKW-A",      "UKW A",      "EJMZALYXVBWFCRQUONTSPIKHGD", ""   },  // 11 reflectors
  {  "UKW-B",      "UKW B",      "YRUHQSLDPXNGOKMIEBFZCWVJAT", ""   },  // 12
  {  "UKW-C",      "UKW C",      "FVPJIAOYEDRZXWGCTKUQSBNMHL", ""   },  // 13
  {  "UKW-B-thin", "UKW B thin", "ENKQAUYWJICOPBLMDXZVFTHRGS", ""   },  // 14
  {  "UKW-C-thin", "UKW C thin", "RDOBJNTKVEHMLFCWZAXGYIPSUQ", ""   }   // 15
};

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

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

// reflector "UmkehrwalzeUmkehrwalze"
byte reflector_type = 12;
char plugboard [10][2] = { 
    {'-','-'}, {'-','-'}, {'-','-'}, {'-','-'}, {'-','-'},
    {'-','-'}, {'-','-'}, {'-','-'}, {'-','-'}, {'-','-'}
  };

// 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(" ");
    }
  }

  // plugboard
  // TODO

  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 = 6;
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] = ' ';
  }
}

// 01234

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 setup () {
  kbd.begin();
  lcd.begin(16, 2);

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


void loop () {
  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);
}