/*  =========================================================================================
 *   
 *    CS408 Environmental Monitoring Independent of Existing Infrastructure
 *    Copyright (C) 2021 Callum Inglis
 *
 *    This program is free software; you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation; either version 2 of the License, or
 *    (at your option) any later version.
 *    
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *    
 *    You should have received a copy of the GNU General Public License along
 *    with this program; if not, write to the Free Software Foundation, Inc.,
 *    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 *    
 *    Contact: Callum.Inglis.2018(at)uni.strath.ac.uk
 *    
 *  =========================================================================================
 *    
 *  Libraries Used
 *    ArduinoJson v6 Library - https://arduinojson.org/v6/doc/ (https://github.com/sandeepmistry/arduino-LoRa/blob/master/LICENSE)
 *    LoRa Library - https://github.com/sandeepmistry/arduino-LoRa/blob/master/API.md (https://github.com/bblanchon/ArduinoJson/blob/6.x/LICENSE.md)
 *    SHTSensor - https://github.com/Sensirion/arduino-sht (https://github.com/Sensirion/arduino-sht/blob/master/LICENSE)
 *    CCS811 - https://wiki.keyestudio.com/KS0457_keyestudio_CCS811_Carbon_Dioxide_Air_Quality_Sensor
 *
 *  =========================================================================================
 *  
 *  Add board to Arduino IDE - http://arduino.esp8266.com/stable/package_esp8266com_index.json
 *  
 */

#include <TimeLib.h>
#include <SPI.h>
#include <Wire.h>
#include <LoRa.h>
#include <SoftwareSerial.h>
#include <ArduinoJson.h>
#include "SHTSensor.h"
#include <CCS811.h>

// Debug / Polling / Transmitting Config
#define DEBUG 2
#define MIN_RAND_DELAY 500 // ms
#define MAX_RAND_DELAY 1250 // ms
#define MAX_TRANSMISSION_RETRIES 5 // No. Retries before recordings new values then retry
#define TX_RESERVATION_TIME 1 // How long do we require use of TX for sending packets? (Seconds)

// Poor mans memory-contrained hashmap alternative. Store index of polled count for each sensor in pollEventCountPerSensor
#define IDX_PPM10 0
#define IDX_PPM25 1
#define IDX_PPM100 2
#define IDX_TEMPERATURE 3
#define IDX_HUMIDITY 4
#define IDX_CO2 5

// Variable Values - May be updated by Gateway
int TX_AFTER_N_READINGS = 4; // No. of samples, after which Transmit Avg Readings
int POLLING_FREQUENCY = 30000; // ms Take value every POLLING_FREQUENCY ms

int pollEventCount = 0; // Number of times data has been sampled in this recording period
int pollEventCountPerSensor[6] = {0}; // Number of times each sensor has been polled in this recording period (Used for averaging values)

// Lora Config
#define ss 16 // Physical Pin 16 = D0 (Default is Physical Pin 5)
#define rst 0
#define dio0 15 // Physical Pin 15 = D8 (Default is Physical Pin 4)
static const int loraSpreadingFactor = 7;
static const int loraSignalBandwidth = 125E3;
static const int loraFrequency = 433E6;

// LoRa Message Tracking
String sensorID = "";
byte localAddress = 0xBB;
byte destination = 0xFF;
uint32_t msgCount = 0;

// Sensor Config
static const int RXPin = 4, TXPin = 3; // Serial PMS Config
SHTSensor sht; // Temp / Humidity Sensor
SoftwareSerial pmsSerial(2, 3); // Particle Sensor Serial
CCS811 co2Sensor; // CO2 Sensor

// Immediate Sensor Data
double temperature; // Sensor Values - Temp 
double humidity;    // Sensor Values - Humidity
double co2;         // Sensor Values - Co2
int32_t ppm10;     // Sensor Values - Particulate
int32_t ppm25;     // Sensor Values - Particulate
int32_t ppm100;    // Sensor Values - Particulate

// Partial Sensor Data
struct pms5003data {
  uint16_t framelen;
  int32_t pm10_standard, pm25_standard, pm100_standard;
  int32_t pm10_env, pm25_env, pm100_env;
  int32_t particles_03um, particles_05um, particles_10um, particles_25um, particles_50um, particles_100um;
  uint16_t unused;
  uint16_t checksum;
};

struct pms5003data data;

/**
 * Setup Humidity / Temperature Sensor
 */
bool setupSHT() {
  Wire.begin();

  if (!sht.init()) {
    Serial.println("[-] SHT Init Failed");
    return false;
  }
  
  sht.setAccuracy(SHTSensor::SHT_ACCURACY_MEDIUM);
  return true;
}

double getTemperature() {
  double t = -1;

  if (sht.readSample()) {
    t = sht.getTemperature();
  }

  return t; // -1 on error
}

double getHumidity() {
  double h = -1;

  if (sht.readSample()) {
    h = sht.getHumidity();
  }

  return h; // -1 on error
}

/**
 * Setup CO2 Sensor
 */
 bool setupCO2() {
  if (co2Sensor.begin() != 0) {
    Serial.println("[-] CO2 Init Failed");
    return false;
  }

  /**
   * @brief Set measurement cycle
   * @param cycle:in typedef enum{
   *                  eClosed,      //Idle (Measurements are disabled in this mode)
   *                  eCycle_1s,    //Constant power mode, IAQ measurement every second
   *                  eCycle_10s,   //Pulse heating mode IAQ measurement every 10 seconds
   *                  eCycle_60s,   //Low power pulse heating mode IAQ measurement every 60 seconds
   *                  eCycle_250ms  //Constant power mode, sensor measurement every 250ms
   *                  }eCycle_t;
   */
  co2Sensor.setMeasCycle(co2Sensor.eCycle_250ms);
  return true;
}

/**
 * Get CO2 reading, returns -1 on error 
 */
double getCo2() {
  double c = -1;
  
  co2Sensor.writeBaseLine(0x847B);
  
  if(co2Sensor.checkDataReady() == true){
    c = co2Sensor.getCO2PPM();
  }

  return c; // -1 on error
}


// TODO DOC
void resetCounters() {
  ppm10 = 0;
  ppm25 = 0;
  ppm100 = 0;
  temperature = 0;
  humidity = 0;
  pollEventCount = 0;
}

bool setupLoRa() {
  LoRa.setPins(ss, rst, dio0);  
  if (!LoRa.begin(loraFrequency)) {
    Serial.println("[-] Fatal. Starting LoRa failed!");
    return false;
  }

  LoRa.setSpreadingFactor(loraSpreadingFactor);
  LoRa.setSignalBandwidth(loraSignalBandwidth);
  LoRa.setTxPower(20);

  Serial.println("[+] LoRa Initialized OK!");
  return true;
}

/* 
 * Return ESP Module ID, based on MAC address
 * Source: https://arduino.stackexchange.com/a/58678
 * 
 * 2021-11-07 - This method proved unreliable, switched to casting ESP.getChipID() as a string instead
 */
String getSensorUID() {
  String sensorID = "EMIEI-"; // [E]nviornmental [M]onitoring [I]ndepenet of [E]xisting [I]nfrastructure
  sensorID.concat(String(ESP.getChipId()));
  return sensorID;
}

/**
 * Get data from PMS Partical Sensor
 * @param Stream s PMS Serial Connection
 * @modifies struct data PMS Sensor Data
 * @returns boolean Data Read Success status
 */
boolean readPMSdata(Stream *s) {
  if (!s->available()) {
    return false;
  }
  
  // Read a byte at a time until we get to the special '0x42' start-byte
  if (s->peek() != 0x42) {
    s->read();
    return false;
  }
 
  // Now read all 32 bytes
  if (s->available() < 32) {
    return false;
  }
    
  uint8_t buffer[32];    
  uint32_t sum = 0;
  s->readBytes(buffer, 32);
 
  // get checksum ready
  for (uint8_t i=0; i<30; i++) {
    sum += buffer[i];
  }
  
  // The data comes in endian'd, this solves it so it works on all platforms
  uint16_t buffer_u16[15];
  for (uint8_t i=0; i<15; i++) {
    buffer_u16[i] = buffer[2 + i*2 + 1];
    buffer_u16[i] += (buffer[2 + i*2] << 8);
  }

  // Struct it
  memcpy((void *)&data, (void *)buffer_u16, 30);
 
  if (sum != data.checksum) {
    Serial.print("\n Checksum failure... ");
    Serial.print("Expected: "); Serial.print(sum);
    Serial.print(", Got: "); Serial.print(data.checksum);
    return false;
  }

  return true;
}

/**
 * Provided sensor data, construct JSON object ready for transmission, Averaged over numSamples
 * 
 * @param messageID
 * @param numSamples Numer of samples averages
 * @param avgPpm10 Average Particuate readings
 * @param avgPpm25 Average Particuate readings
 * @param avgPpm100 Average Particuate readings
 * @param avgTemperature Average Temperature
 * @param avgHumidity Average Humidity
 * @param avgCo2 Average Co2
 * 
 * @return DynamicJsonDocument sensorData
 */
DynamicJsonDocument prepareSensorData(int messageID, int numSamples, uint32_t avgPpm10, uint32_t avgPpm25, uint32_t avgPpm100, double avgTemperature, double avgHumidity, double avgCo2) {
  DynamicJsonDocument doc(2048);

  JsonObject metadata = doc.createNestedObject("sensorMetadata");
  metadata["uid"] = getSensorUID();
  metadata["messageID"] = messageID;
  metadata["samplePeriod"] = numSamples;

  JsonObject data = doc.createNestedObject("data");

  JsonObject ppm = data.createNestedObject("ppm"); // Particulates
  ppm["p10"] = avgPpm10;
  ppm["p25"] = avgPpm25;
  ppm["p100"] = avgPpm100;

  JsonObject sht = data.createNestedObject("sht");  // Temp, Humidity
  sht["temperature"] = avgTemperature;
  sht["humidity"] = avgHumidity;

  JsonObject co2 = data.createNestedObject("co2"); // TODO Co2
  co2["tmp"] = avgCo2;

  return doc;
}

/**
 * Listen on Lora for other messages, returns true if no messages detected within [listenDuration] seconds
 * Use prior to transmissions to avoid interruption of other messages
 *
 * @param int listenDuration How long to listen on Lora RX in Seconds
 * @returns boolean If messages detected within listenDuration  
 */
boolean clearToSend(int listenDuration) {  
  int listenUntil = now() + listenDuration;

  // Listen until timeout expires
  while (listenUntil >= now()) {
    int packetSize = LoRa.parsePacket();
    
    if (packetSize) {
      Serial.println("[-] TX Busy - Not Clear To Send");
      return false; // Other message heard on Rx, infer that we can not transmit just now. 
    }

    delay(5);
  }

  return true; // We didn't hear anything, so continue
}

// Introduce random delay to avoid another collision
void waitRandomDelay(int minSeconds = 0) {
  int randDelay = (minSeconds * 1000) + random(MIN_RAND_DELAY, MAX_RAND_DELAY);
  if (DEBUG){
    Serial.println("[i] Delay for " + String(randDelay) + "ms");
  }
  delay(randDelay);
}

/**
 * Send short "clear to send?" packet
 * Tx - "I've got data to send!" + UID
 * @param messageID ID of this message / sample period
 */
void sendTxHello(int messageID) {
  Serial.println("[+] Transmit - \"Hello\"");
  
  DynamicJsonDocument txHello(2048);
  txHello["uid"] = getSensorUID();
  txHello["reservationTime"] = TX_RESERVATION_TIME;
  txHello["messageID"] = messageID;

  sendJsonPayloadWithLoRa(txHello);
}

/**
 * Send Payload
 * Tx - Sensor Payload
 * @param payload JSON Payload to be sent
 */
void sendTxPayload(DynamicJsonDocument payload) {
  Serial.println("[+] Transmit - Payload");
  sendJsonPayloadWithLoRa(payload);
}

/**
 * Listen for messages on TX, expecting JSON
 * 
 * @param listenDuration How long to listen on Lora RX in Seconds
 * @return StaticJsonDocument|null RX Payload in JSON, null on timeout reached or error
 */
StaticJsonDocument<1024> listenForAndConsumeMessage(int listenDuration) {
  int listenUntil = now() + listenDuration;
  StaticJsonDocument<1024> json;

  // Listen until timeout expires
  while (listenUntil >= now()) {
    int packetSize = LoRa.parsePacket();

    if (!packetSize) {
      delay(3);
      continue;
    }

    // Read in packet, ensure we only bring in anything after { and before } inclusive
    String incoming = "";
    char temp;
    
    while (LoRa.available()) {
      temp = (char)LoRa.read();

      if (incoming.length() == 0 && temp == '{') { // Opening {
        incoming = "{";

      } else if (temp == '}') { // Closing } 
        incoming.concat("}");
        break;

      } else if (incoming.length() > 0) { // Anything else that's valid
        incoming.concat(temp);
      }
    }

    Serial.println(incoming);
    // Deserialize - TODO Error Handling https://arduinojson.org/v6/api/json/deserializejson/
    deserializeJson(json, incoming);
    break;
  }

  return json;
}

/**
 * Listen for response to TxHello for [listenDuration] seconds, check Sensor UID & Message ID match and return true if we have clear to transmit payload.
 * @param listenDuration How long to listen on Lora RX in Seconds
 * @param messageID ID of message we are expecting to receive
 * @return boolean Clear to Transmit 
 */
boolean listenForTxHelloAccept(int listenDuration, int messageID) {
  Serial.println("[+] Transmit - \"Hello\" - Listening for \"Hello\" Ack & Clear to Send");

  StaticJsonDocument<1024> json = listenForAndConsumeMessage(listenDuration);

  // Timeout, likely TX Hello was not recieved or ack got lost
  if (json.isNull()) {
    Serial.println("[-] Transmit - \"Hello\" - Timeout while waiting for Clear to Send\n\n");
    return false;
  }

  const bool okToTransmit = json["okTransmit"];
  const String authIsForUid = json["uid"];
  const int authIsForMessageID = json["messageID"];
  const String gatewayUid = json["gatewayUid"];

  // Verify txHello.okTransmit is True and Sensor UID & Message IDs Match
  if (DEBUG) {
    if (authIsForUid == getSensorUID()) { 
      Serial.println("[+] Transmit - \"Hello\" - Sensor UID Match!"); 
    } else { 
      Serial.println("[-] Transmit - \"Hello\" - Sensor UID Mis-Match! " + String(authIsForUid) + " vs " + String(getSensorUID())); 
    }

    if (authIsForMessageID == messageID) {
      Serial.println("[+] Transmit - \"Hello\" - Message ID Match!"); 
    } else { 
      Serial.println("[-] Transmit - \"Hello\" - MessageID Mis-Match!"); 
    }
  }

  // Ok To Trasmit, Sensor UID Match & Message ID Match
  bool clearToSend = (okToTransmit 
                      && authIsForUid == getSensorUID() 
                      && authIsForMessageID == messageID);

  if (DEBUG) {
    if (clearToSend) {
      Serial.println("[+] Transmit - \"Hello\" - Recieved \"Clear to Transmit\" Payload");
    } else {
      Serial.println("[-] Transmit - \"Hello\" - Can Not Transmit At This Time");
    }
  }

  return clearToSend;
}

// TODO DOC
boolean listenForTxPayloadAccept(int listenDuration, int messageID) {
  Serial.println("[.] Transmit - Payload - Listening for Ack");

  StaticJsonDocument<1024> json = listenForAndConsumeMessage(listenDuration);

  // Timeout, likely TX Payload was not recieved or ack got lost.
  if (json.isNull()) {
    Serial.println("[-] Transmit - Payload - Ack Timeout Reached - Assuming Message Was Not Delivered\n\n");
    return false;
  }

  const bool ackStatus = json["ackStatus"];
  const String authIsForUid = json["uid"];
  const int authIsForMessageID = json["replyMsgID"];
  const String gatewayUid = json["gatewayUid"];
  

  // Verify Sensor UID Match
  if (authIsForUid == getSensorUID()
      && authIsForMessageID == messageID) {

    Serial.println("[+] Transmit - Payload - Ack Recieved: " + String(ackStatus) + "\n");
    
    if (ackStatus) {  // It all worked :)
      // Check for any updated config values
      const int newTxAfterNReadings = json["txAfterNReadings"];
      if (newTxAfterNReadings != NULL && newTxAfterNReadings != TX_AFTER_N_READINGS) {
        TX_AFTER_N_READINGS = newTxAfterNReadings;
        Serial.println("[+] Tx After N Readings Updated, now: " + String(TX_AFTER_N_READINGS) + " samples");
      }
      
      const int newPollingFrequency = json["pollingFrequency"];
      if (newPollingFrequency != NULL && newPollingFrequency != POLLING_FREQUENCY) {
        POLLING_FREQUENCY = newPollingFrequency;
        Serial.println("[+] Polling Frequency Updated, now: " + String(POLLING_FREQUENCY) + "ms");
      }
      
      return true; 
    }

    // TODO Retransmit, recover, etc
    Serial.println("[-] Transmit - Payload - Ack Failed - TODO Setup Retransmission");
    return false;
  }

  // TODO Else UID Mis-Match so we wait for next message
  Serial.println("[-] Transmit - Payload - Ack Message ID or Sensor ID Mis-Match");
  return false;
}

/**
 * Send JSON Payload over LoRa
 * @param payload JSON Payload to be send
 */
void sendJsonPayloadWithLoRa(DynamicJsonDocument payload) {
  LoRa.beginPacket();
  serializeJson(payload, LoRa);
  LoRa.endPacket();
}

// TODO - Finish Acks & Naks
boolean transmitData(DynamicJsonDocument payload) {

  // TODO MAX_RETRIES

  // Listen for other communication
  //    Rx - Listen for other messages, if no messages heard then continue
  if (!clearToSend(0.5)) {
    waitRandomDelay(); // Wait for short time before retrying
    while (!clearToSend(0.5)) { waitRandomDelay(); } // TODO MAX_TIMEOUT
  }

  // Send TX Hello
  sendTxHello(msgCount);

  // Await TX Hello Auth - Expect: Timeout | Not Auth | Accept + Clear To Transmit
  if (!listenForTxHelloAccept(TX_RESERVATION_TIME * 4, msgCount)) { 
    return false; // Can't transmit just now, we will retry
  }

  // Send TX Payload
  sendTxPayload(payload);

  // Await TX Payload Ack - Expect: Timeout | Nack | Accept + Data Match
  if (!listenForTxPayloadAccept(2, msgCount)) {
    return false; // TODO Ack Failed Setup a retry here!
  }

  // TODO Update Sensor Config

  // TODO Clear values & Continue
  resetCounters();
  Serial.println("Packet Sent Succesfully\n");
  Serial.println("---------------+++++-------------------");
  msgCount++;

  return true;
}

/**
 * Add recorded values to the total, do not store false/failed/invalid values, e.g. -1
 */
double addValueToTotal(double avgValue, double newValue, int idx_sensor_count) {
  if (newValue == -1) {
    return avgValue;
  }

  pollEventCountPerSensor[idx_sensor_count] += 1;
  return avgValue + newValue;
}

double averageValueByNumberPolls(double value, int idx_sensor_count) {
  // Ensure we do not divide by 0
  if (pollEventCountPerSensor[idx_sensor_count] == 0) {
    return -1;
  }

  // Otherwise divide total by number of times we sampled this value
  return value / pollEventCountPerSensor[idx_sensor_count];
}

void setup() {
  delay(1000);
  
  Serial.begin(115200); // Console Debug
  Serial.println("\n\n[+] Transmitter Node");

  pmsSerial.begin(9600); // Partical Sensor

  // Setup Hardware TODO Handle broken sensor!
  if (!setupSHT()) { /*while(1);*/ } // Temp/Humidity - Die on Error
  if (!setupCO2()) { /*while(1);*/ }
  if (!setupLoRa()) { /*while(1);*/ } // Die on error
}

void loop() {  
  delay(POLLING_FREQUENCY);
  pollEventCount++;

  // Get values from sensors
  double instantTemperature = getTemperature(); // Prevent polling of sensors too frequently when debugging
  double instantHumidity = getHumidity();
  double instantCo2 = getCo2();

  int instantPPM10, instantPPM25, instantPPM100;
  if (!readPMSdata(&pmsSerial)) { // Error reading PMS Data, ignore on this round
    instantPPM10 = instantPPM25 = instantPPM100 = -1; 
    
  } else {
    instantPPM10 = data.pm10_standard;  // Standard (_standard) or Environmental (_env) readings for Particulate Data
    instantPPM25 = data.pm25_standard;  
    instantPPM100 = data.pm100_standard;
  }
  
  if (DEBUG) {
    Serial.println(); Serial.print(String(pollEventCount) + ") "); Serial.print("PM 1.0: "); Serial.print(instantPPM10); Serial.print("\t\tPM 2.5: "); Serial.print(instantPPM25); Serial.print("\t\tPM 10: "); Serial.print(instantPPM100);
    Serial.print("\t\tTemp: "); Serial.print(instantTemperature); Serial.print("\t\tHumidity: "); Serial.print(instantHumidity); Serial.print("\t\tCo2: "); Serial.print(instantCo2);
  }

  // Add to running total, we will divide by the number of times samples later on to get an average over sample period
  temperature = addValueToTotal(temperature,  instantTemperature, IDX_TEMPERATURE);
  humidity    = addValueToTotal(humidity,     instantHumidity,    IDX_HUMIDITY);
  co2         = addValueToTotal(co2,          instantCo2,         IDX_CO2);
  ppm10       = addValueToTotal(ppm10,        instantPPM10,       IDX_PPM10);
  ppm25       = addValueToTotal(ppm25,        instantPPM25,       IDX_PPM25);
  ppm100      = addValueToTotal(ppm100,       instantPPM100,      IDX_PPM100);

  // If we should now transmit
  if (pollEventCount >= TX_AFTER_N_READINGS) {

    // Average Values over recording period
    double  avgTemperature  = averageValueByNumberPolls(temperature,  IDX_TEMPERATURE);
    double  avgHumidity     = averageValueByNumberPolls(humidity,     IDX_HUMIDITY);
    double  avgCo2          = averageValueByNumberPolls(co2,          IDX_CO2);
    int     avgPpm10        = averageValueByNumberPolls(ppm10,        IDX_PPM10);
    int     avgPpm25        = averageValueByNumberPolls(ppm25,        IDX_PPM25);
    int     avgPpm100       = averageValueByNumberPolls(ppm100,       IDX_PPM100);

    if (DEBUG) {
      Serial.println(""); Serial.print("Avg ppm10: "); Serial.print(avgPpm10); Serial.print("\t\tAvg ppm25: "); Serial.print(avgPpm25); Serial.print("\t\tAvg ppm100: "); Serial.print(avgPpm100);
      Serial.print("\t\tAvg Temp: "); Serial.print(avgTemperature); Serial.print("\t\tAvg Humidity: "); Serial.print(avgHumidity); Serial.print("\t\tAvg Co2: "); Serial.print(avgCo2);
      Serial.print("\t\tChip ID: "); Serial.println(getSensorUID()); Serial.println("");
    }

    // Prepare Data For Send
    DynamicJsonDocument sensorData = prepareSensorData(msgCount, pollEventCount, avgPpm10, avgPpm25, avgPpm100, avgTemperature, avgHumidity, avgCo2);
    
    // Transmit
    if (transmitData(sensorData)) {
      return; // It all worked, values reset, now record new values
    }

    // Transmission failed, handle re-tries
    int numRetries = 1;
    waitRandomDelay(TX_RESERVATION_TIME*2); // Introduce random delay to avoid another collision
    
    while (!transmitData(sensorData) && numRetries < MAX_TRANSMISSION_RETRIES){
      numRetries++;
      Serial.println("[-] Failed to send packet, retrying. Attempt " + String(numRetries) + " of " + String(MAX_TRANSMISSION_RETRIES) + "\n");
      waitRandomDelay(TX_RESERVATION_TIME*2); // Introduce random delay to avoid another collision
    }

    // We were able to transmit after retries, values reset, now record new values
    if (numRetries < MAX_TRANSMISSION_RETRIES) {
     return;
    }

    // Failed to transmit - Don't Clear Counters, record more values then try to retransmit on next send
    Serial.println("[-] Failed to send packet, max retries reached. Aborting");
    return;
  }
}