/**
 * Sesame-Device
 * Christoph Hagen, 2022.
 *
 * The code for a simple door unlock mechanism where a servo pushes on an existing
 * physical button.
 */
#include <Arduino.h>
#include <WiFi.h>
#include <WiFiMulti.h>
#include <WiFiClientSecure.h>
#include <WebSocketsClient.h>
#include <EEPROM.h> // To mark used keys as expired
#include <ESP32Servo.h> // To control the servo


/* Settings */

// The WiFi network to connect to
constexpr const char* WIFI_SSID = "MyNetwork";
constexpr const char* WIFI_PWD = "MyPassword";

// The remote server to connect to
constexpr const char* SERVER_URL = "christophhagen.de";
constexpr const int SERVER_PORT = 443;
constexpr const char* SERVER_PATH = "/sesame/listen";
constexpr const char* SERVER_PSK = "access token";
#define USE_SSL // Use SSL for the Websocket connection to the server

// The interval to attempt to reconnect the socket to the server
constexpr unsigned long SOCKET_RECONNECT_TIME = 5000;

// Crypto setting for the security of keys
constexpr uint8_t KEY_STRENGTH = 128;
constexpr uint8_t KEY_BYTE_COUNT = KEY_STRENGTH / 8;

// The keys defined to allow opening
// Once a key is used, the corresponding byte in EEPROM is set to 1, 
// to prevent it from being used again.
const uint8_t keys[][KEY_BYTE_COUNT] = {
    {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
};
const uint16_t KEY_COUNT = sizeof keys / KEY_BYTE_COUNT;


/* Servo */

// Servo is Emax ES08MA II

// The time (in ms) to keep the door button pressed
constexpr uint32_t OPENING_DURATION = 2000;

// The timer to use to control the servo
constexpr int pwmTimer = 0;

// Pin wired up to the servo data line
constexpr int servoPin = 14;

// The Emax is a standard 50 Hz servo
constexpr int servoFrequency = 50;

// The microseconds to set the servo to the pressed state
constexpr int PRESS_STATE_INTERVAL = 1550;

// The microseconds to set the servo to the released state
constexpr int RELEASE_STATE_INTERVAL = 1700;


/* Global variables */

// Servo controller
Servo servo;

// PWM Module needed for the servo
ESP32PWM pwm;

// WiFi module to connect to the network
WiFiMulti WiFiMulti;

// WebSocket to connect to the control server
WebSocketsClient webSocket;

// Indicator that the socket is connected.
bool socketIsConnected = false;

// The index of the next valid key, which is the next key after the highest used key.
// If the index is larger or equal to the total key count, then no usable keys exist.
uint16_t nextKeyIndex = KEY_COUNT;

// Flag to signal that the door button should be pressed
bool shouldStartOpening = false;

// Indicator that the door button is pushed
bool buttonIsPressed = false;

// The time (in ms since start) when the door opening should end
uint32_t openingEndTime = 0;


/* Events */

/**
 * An event occuring due to a server request
 */
enum class SesameEvent {
    TextReceived = 1,
    UnexpectedSocketEvent = 2,
    InvalidPayloadSize = 3,
    InvalidKeyIndex = 4,
    InvalidKey = 5,
    KeyAlreadyUsed = 6,
    KeyWasSkipped = 7,
    KeyAccepted = 8,
};


/* Key management */

/**
 * Get the index of the next key which wasn't used yet
 *
 * A key is unused, if the EEPROM corresponding to the key is zero.
 * If all keys have been used, returns KEY_COUNT
 */
uint16_t indexOfNextUsableKey();

/**
 * Checks if a key was already marked as used.
 *
 * A key was used, if the EEPROM byte corresponding to the key is non-zero.
 * @param keyIndex The index of the key in the 'keys' array.
 * @return 0, if the key is unused, else non-zero.
 */
uint8_t keyWasAlreadyUsed(uint16_t keyIndex);

/**
 * Marks a key as used
 *
 * Sets the corresponding EEPROM byte to non-zero.
 * @param keyIndex The index of the key in the 'keys' array.
 */
void markKeyUsed(uint16_t keyIndex);

/**
 * Marks all keys unused.
 *
 * Sets the EEPROM data for each key index to zero.
 *
 * WARNING: Only to be used when a new set of keys has been set.
 * Otherwise replay attacks are possible with observed old keys.
 */
void markAllKeysUnused();

/**
 * Compares two keys in constant time
 * @return 1, if the keys are equal, else 0
 */
uint8_t keysAreEqual(const uint8_t* key1, const uint8_t* key2);


/* Servo management */

/**
 * Push the door opener button down by moving the servo arm.
 */
void pressButton();

/**
 * Release the door opener button by moving the servo arm.
 */
void releaseButton();


/* Logic */

/**
 * Send a response event to the server.
 *
 * Sends the event type as a single byte representing the raw event type.
 * @param event The event type
 */
void sendResponse(SesameEvent event) {
    uint8_t response[1];
    response[0] = static_cast<uint8_t>(event);
    webSocket.sendBIN(response, 1);
    //webSocket.sendTXT(text);
    Serial.printf("[INFO] Event %d\n", response[0]);
}

/**
 * Send the pre-shared key to the server to complete the socket connection.
 */
void authenticateDevice() {
    webSocket.sendTXT(SERVER_PSK);
}

/**
 * Process received binary data.
 *
 * Checks whether the received data is a valid and unused key,
 * and then signals that the motor should move.
 *
 * @param payload The pointer to the received data.
 * @param length The number of bytes received.
 * @return The event to signal to the server.
 */
SesameEvent handleReceivedRequest(const uint8_t* payload, size_t length) {
    if (length != KEY_BYTE_COUNT + 2) {
        return SesameEvent::InvalidPayloadSize;
    }
    uint16_t keyIndex = ((uint16_t) payload[0] << 8) + payload[1];
    if (keyIndex >= KEY_COUNT) {
        return SesameEvent::InvalidKeyIndex;
    }
    if (!keysAreEqual(payload + 2, keys[keyIndex])) {
        return SesameEvent::InvalidKeyIndex;
    }
    if (keyWasAlreadyUsed(keyIndex)) {
        return SesameEvent::KeyAlreadyUsed;
    }
    if (nextKeyIndex > keyIndex) {
        return SesameEvent::KeyWasSkipped;
    }
    markKeyUsed(keyIndex);
    nextKeyIndex = keyIndex + 1;
    // Move servo
    shouldStartOpening = true;
    return SesameEvent::KeyAccepted;
}

/**
 * Process received binary data.
 *
 * Checks whether the received data is a valid and unused key,
 * and then signals that the motor should move.
 * Sends the event id to the server as a response to the request.
 *
 * If the key is valid, then `shouldStartOpening` is set to true.
 *
 * @param payload The pointer to the received data.
 * @param length The number of bytes received.
 */
void processReceivedBytes(uint8_t * payload, size_t length) {
    SesameEvent event = handleReceivedRequest(payload, length);
    sendResponse(event);
}

/**
 * Callback for WebSocket events.
 *
 * Updates the connection state and processes received keys.
 *
 * @param payload The pointer to received data
 * @param length The number of bytes received
 */
void webSocketEvent(WStype_t type, uint8_t * payload, size_t length) {
    switch(type) {
        case WStype_DISCONNECTED:
            socketIsConnected = false;
            Serial.println("[INFO] Socket disconnected.");
            break;
        case WStype_CONNECTED:
            socketIsConnected = true;
            authenticateDevice();
            Serial.printf("[INFO] Socket connected to url: %s\n", payload);
            break;
        case WStype_TEXT:
            sendResponse(SesameEvent::TextReceived);
            break;
        case WStype_BIN:
            handleReceivedRequest(payload, length);
            break;
        case WStype_ERROR:
        case WStype_FRAGMENT_TEXT_START:
        case WStype_FRAGMENT_BIN_START:
        case WStype_FRAGMENT:
        case WStype_FRAGMENT_FIN:
            sendResponse(SesameEvent::UnexpectedSocketEvent);
            break;
    }
}

void setup() {
    Serial.begin(115200);
    Serial.setDebugOutput(true);
    Serial.println("[INFO] Device started");

    markAllKeysUnused();
    Serial.println("[WARN] All keys reset");
    
    ESP32PWM::allocateTimer(pwmTimer);
    servo.setPeriodHertz(servoFrequency);
    servo.attach(servoPin);
    Serial.println("[INFO] Servo configured");
    
    Serial.printf("[INFO] Key security: %d bit (%d byte keys)\n", KEY_STRENGTH, KEY_BYTE_COUNT);
    nextKeyIndex = indexOfNextUsableKey();
    uint8_t percentage = (KEY_COUNT - nextKeyIndex) * 100 / KEY_COUNT;
    Serial.printf("[INFO] %d of %d keys remaining (%d %%)\n", KEY_COUNT - nextKeyIndex, KEY_COUNT, percentage);
    Serial.printf("[INFO] Connecting to WiFi '%s'\n", WIFI_SSID);
    WiFiMulti.addAP(WIFI_SSID, WIFI_PWD);
    
    while(WiFiMulti.run() != WL_CONNECTED) {
        delay(100);
    }

    Serial.println("[INFO] WiFi connected");
    Serial.printf("[INFO] Opening socket %s%s on port %d\n", SERVER_URL, SERVER_PATH, SERVER_PORT);

    #ifdef USE_SSL
    webSocket.beginSSL(SERVER_URL, SERVER_PORT, SERVER_PATH);
    #else
    webSocket.begin(SERVER_URL, SERVER_PORT, SERVER_PATH);
    #endif
    webSocket.onEvent(webSocketEvent);
    
    // try again every 5000 ms if connection has failed
    webSocket.setReconnectInterval(SOCKET_RECONNECT_TIME);
}

void loop() {
    webSocket.loop();
    if (shouldStartOpening) {
        shouldStartOpening = false;
        openingEndTime = millis() + OPENING_DURATION;
        pressButton();
    }
    if (buttonIsPressed && millis() > openingEndTime) {
        releaseButton();
    }
}

/* Key management */

uint16_t indexOfNextUsableKey() {
    // Find the highest key which was previously used
    for (uint16_t keyIndex = 0; keyIndex < KEY_COUNT; keyIndex += 1) {
        if (EEPROM.read(KEY_COUNT - keyIndex)) {
            // The following key is the next unused
            return keyIndex + 1;
        }
    }
    // No key previously used
    return 0;
}

uint8_t keyWasAlreadyUsed(uint16_t keyIndex) {
    return EEPROM.read(keyIndex);
}

void markKeyUsed(uint16_t keyIndex) {
    EEPROM.write(keyIndex, 1);
}

void markAllKeysUnused() {
    for (uint16_t keyIndex = 0; keyIndex < KEY_COUNT; keyIndex += 1) {
        EEPROM.write(keyIndex, 0);
    }
}

uint8_t keysAreEqual(const uint8_t* key1, const uint8_t* key2) {
    uint8_t result = 0;
    for (uint8_t i = 0; i < KEY_BYTE_COUNT; i += 1) {
        result |= key1[i] ^ key2[i];
    }
    if (result) {
        return 0;
    }
    return 1;
}

/* Servo management */

void pressButton() {
    servo.write(PRESS_STATE_INTERVAL);
    buttonIsPressed = true;
}

void releaseButton() {
    servo.write(RELEASE_STATE_INTERVAL);
    buttonIsPressed = false;
}