/*
 * Assignment 1 part I
 * Created by: Gaylord Mbuyi and James Yurick
 * LBL A2
*/

#include <Arduino.h>

uint32_t mul_mod(uint32_t a, uint32_t b, uint32_t m) {
  uint32_t result = 0;
  uint32_t x = b % m;

  for (int i = 0; i < 32; ++i) {
    uint32_t myBit = 31 - i;
    //Serial.print(i);
    //Serial.print(": ");
    //Serial.println(((a << myBit) >> 31));
    if (((a << myBit) >> 31) == 1) {
      result = (result + x) % m;
    }
    x = (x << 1) % m;
  }

  return result;
}

/** Calculates base^power mod m */
uint32_t pow_mod(uint32_t base, uint32_t power, uint32_t m) {
  uint32_t result = 1;
  uint32_t x = base % m;
  for (int i = 0; i < 32; ++i ) {
    if ((power & (1 << i)) != 0) {
      result = mul_mod(result, x, m);
    }
    x = mul_mod(x, x, m);
  }

  return result;
}

/** Converts zero-terminated string into an unsigned long */
int str2int( char* buf ) {
  return atoi(buf);
}

/** Reads a number from the keyboard by using the Serial port */
int get_number() {
  const int LINE_LENGTH = 128;
  int ret = 0; // number to be returned
  char line[LINE_LENGTH+1];
  int i = 0;
  while (i<LINE_LENGTH) {
    int byte = Serial.read();
    if (byte!=-1) {
      if (byte==10 || byte==13) { // <ENTER> was received
        break;
      }
      line[i] = byte;
      i = i + 1;
      Serial.print( (char) byte );
    }
  }
  line[i] = 0;
  // calculate ret
  ret = str2int(line);
  Serial.println();

  return ret;
}

/** Generates random private key using analog port 1 */
int generate_key() {
  uint32_t result = 0;
  uint32_t base2 = 1;
  // generate 32-bit random number
  for (int i = 0; i < 32; ++i) {
    uint16_t keybit = analogRead(1); // get number from analog port 1
    keybit = (keybit << 15) >> 15; // get least significant bit
    // convert bits into integer
    if (keybit == 1) {
      result = result + base2;
    }
    base2 = base2 * 2;

    delay(50);
  }
  return result;
}

void print_status(){
  typedef enum { client, server } client_serv;
  client_serv status;
  if (digitalRead(13) == HIGH) {
    status= server;
    Serial.println("I'm a server");
  }
  else {
    status= client;
    Serial.println("I'm a client");
  }
  Serial.println(status);
}

/** Waits for a certain number of bytes on Serial3 or timeout
 * @param nbytes: the number of bytes we want
 * @param timeout: timeout period (ms); specifying a negative number
 *                turns off timeouts (the function waits indefinitely
 *                if timeouts are turned off).
 * @return True if the required number of bytes have arrived.
 */
bool wait_on_serial3( uint8_t nbytes, long timeout ) {
  unsigned long deadline = millis() + timeout;//wraparound not a problem
  while (Serial3.available()<nbytes && (timeout<0 || millis()<deadline))
  {
    delay(1); // be nice, no busy loop
  }
  return Serial3.available()>=nbytes;
}

/** Writes an uint32_t to Serial3, starting from the least-significant
 * and finishing with the most significant byte.
 */
void uint32_to_serial3(uint32_t num) {
  Serial3.write((char) (num >> 0));
  Serial3.write((char) (num >> 8));
  Serial3.write((char) (num >> 16));
  Serial3.write((char) (num >> 24));
}

/** Reads an uint32_t from Serial3, starting from the least-significant
 * and finishing with the most significant byte.
 */
uint32_t uint32_from_serial3() {
  uint32_t num = 0;
  num = num | ((uint32_t) Serial3.read()) << 0;
  num = num | ((uint32_t) Serial3.read()) << 8;
  num = num | ((uint32_t) Serial3.read()) << 16;
  num = num | ((uint32_t) Serial3.read()) << 24;
  return num;
}

uint32_t server_mod(uint32_t our_public_key) {
Serial.println("The func server_mod is called");
  // choose private key

//~ Serial.print("Our public key: "); Serial.println(our_public_key);
//~ Serial.print("Their public key: ");
//~ int their_public_key = get_number(); // get from other computer (was B)
//uint32_t shared_key = (pow_mod(their_public_key, private_key,prime)) % 256; // B^a % p and convert it to klow
uint32_t ckey;
//~ Serial.print("Shared key: "); Serial.println(shared_key);

  typedef enum {
    listen, waitingForKey, waitForAck, dataExchange
    } Status;
    Status myStatus = listen;

    Serial.print("My status before while is:"); Serial.println(myStatus);
    while (myStatus != dataExchange) {
      if (myStatus == listen) {
        Serial.println("My status is listen");
        if ((wait_on_serial3(5, 1000) == true) && (Serial3.read() == 'C')) {
          Serial.print("My status in first if is:"); Serial.println(myStatus);
          myStatus = waitingForKey;
        }
      }
      else if ((wait_on_serial3(4, 1000) == true) && (myStatus == waitingForKey)) {
        Serial.println("My status is waitingForKey");
        ckey = uint32_from_serial3();
        Serial.print("Just read ckey: "); Serial.println(ckey);
        Serial3.write('A');
        Serial.println("Just wrote A: ");
        uint32_to_serial3(our_public_key);
        Serial.print("Our_public_key: "); Serial.println(our_public_key);
        myStatus = waitForAck;
      }
      else if ((wait_on_serial3(1, 1000) == true) && (myStatus == waitForAck)) {
        Serial.println("My status is waitForACK");
        char a = Serial3.read();
        if (a == 'C') {
          Serial.println("Just read C");
          myStatus = waitingForKey;
        }
        else if (a == 'A') {
          myStatus = dataExchange;
          Serial.println("My status is dataExchange in else if");
        }
      }
      else {
        myStatus = listen;
      }
    }
    Serial.println("My status is dataExchange");
    Serial.print("Our public key (skey): "); Serial.println(our_public_key);
    Serial.print("ckey: "); Serial.println(ckey);
return ckey;
}

/** Implements the Park-Miller algorithm with 32 bit integer arithmetic
 * @return ((current_key * 48271)) mod (2^31 - 1);
 * This is linear congruential generator, based on the multiplicative
 * group of integers modulo m = 2^31 - 1.
 * The generator has a long period and it is relatively efficient.
 * Most importantly, the generator's modulus is not a power of two
 * (as is for the built-in rng),
 * hence the keys mod 2^{s} cannot be obtained
 * by using a key with s bits.
 * Based on:
 * http://www.firstpr.com.au/dsp/rand31/rand31-park-miller-carta.cc.txt
 */
uint32_t next_key(uint32_t current_key) {
  const uint32_t modulus = 0x7FFFFFFF; // 2^31-1
  const uint32_t consta = 48271; // we use that this is <=15 bits
  uint32_t lo = consta*(current_key & 0xFFFF);
  uint32_t hi = consta*(current_key >> 16);
  lo += (hi & 0x7FFF)<<16;
  lo += hi>>15;
  if (lo > modulus) lo -= modulus;
  return lo;
}

int main() {
  // Initialize Arduino functionality
  init();

  // Attach USB for applicable processors
  #ifdef USBCON
     USBDevice.attach();
  #endif

  Serial.begin(9600); // Serial0: communication with the PC
  Serial3.begin(9600); // Serial3: communication with the other Arduino

  //~ print_status();

  uint32_t prime = 2147483647; // prime in Diffie-Hellman
  uint32_t generator = 16807; // generator in Diffie-Hellman
  uint32_t private_key = generate_key(); // generate random key
  uint32_t our_public_key = pow_mod(generator, private_key, prime ); // A = g^a % p
  uint32_t their_public_key = server_mod(our_public_key);
  uint32_t shared_key = (pow_mod(their_public_key, private_key, prime)); // B^a % p and convert it to klow
  Serial.print("share_key before while: "); Serial.println(shared_key);
  Serial.println("*******");

  // chatting starts here
  int incomingByte = 0;
  uint32_t shared_key_received = shared_key;
  uint32_t shared_key_send = shared_key;
  while (true) {
    // grab byte from other Arduino
    if (Serial3.available() > 0) {
      incomingByte = Serial3.read();
      shared_key_received = (next_key(shared_key_received) << 24) >> 24;
      //Serial.print("shared_key in first if: "); Serial.println(shared_key);
      uint32_t decryptedByte =
        incomingByte ^ shared_key;
      if (decryptedByte == 10 || decryptedByte == 13) { // Enter received
        Serial.write('\n');  // LF
        Serial.write('\r');  // CR
      }
      else
        Serial.write((char)decryptedByte); // print to the PC screen
    }
    // grab byte from the PC
    if (Serial.available() > 0) {
      incomingByte = Serial.read();
      shared_key_send = (next_key(shared_key_send) << 24) >> 24;
      //Serial.print("shared_key in scnd if: "); Serial.println(shared_key);
      uint32_t byteToSend = incomingByte ^ shared_key;
      //Serial.print("The char is: "); Serial.println((char)byteToSend);
      Serial3.write((char)byteToSend); // send byte to other Arduino
    }

    //Serial.println(shared_key);
  }

  Serial.end();
  Serial3.end();

  return 0;
}