/*  =========================================================================================
 *   
 *    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)
 *
 *  =========================================================================================
 */

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

// Debug / Polling / Transmitting Config
#define DEBUG 1
#define MIN_RAND_DELAY 500 // ms
#define MAX_RAND_DELAY 1250 // ms
#define POLLING_FREQUENCY 2500 // ms Take value every POLLING_FREQUENCY ms
#define MAX_TRANSMISSION_RETRIES 5 // No. Retries before recordings new values then retry
#define TX_AFTER_N_READINGS 10 // No. of samples, after which Transmit Avg Readings
#define TX_RESERVATION_TIME 1 // How long do we require use of TX for sending packets? (Seconds)
int pollEventCount = 0; // Number of times data has been sampled in this recording period

// 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;

// Serial PMS Config
static const int RXPin = 4, TXPin = 3;

// Temp / Humidity Sensor
SHTSensor sht;

// Particle Sensor Serial
SoftwareSerial pmsSerial(2, 3);

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

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

// Partial Sensor Data
struct pms5003data {
  uint16_t framelen;
  uint16_t pm10_standard, pm25_standard, pm100_standard;
  uint16_t pm10_env, pm25_env, pm100_env;
  uint16_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
}

/**
 * TODO getCo2()
 */
double getCo2() {
  return 0;
}


// 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);

  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);
      }
    }

    // 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) { return true; } // It all worked :)

    // 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 * 1.5, 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;
}

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

  pmsSerial.begin(9600); // Partical Sensor

  // Setup Hardware
  if (!setupSHT()) { while(1); } // Temp/Humidity - Die on Error
  if (!setupLoRa()) { while(1); } // Die on error
}

// Main
void loop() {
  // TODO Gather Sensor Data

  // Error reading PMS Data, ignore on this round
  if (!readPMSdata(&pmsSerial)) {
    data.pm10_standard = 0;
    data.pm25_standard = 0;
    data.pm100_standard = 0; 
  }
  
  pollEventCount++;
  
  if (DEBUG) {
    Serial.println();
    Serial.print(String(pollEventCount) + ") ");
    Serial.print("PM 1.0: "); Serial.print(data.pm10_standard);
    Serial.print("\t\tPM 2.5: "); Serial.print(data.pm25_standard);
    Serial.print("\t\tPM 10: "); Serial.print(data.pm100_standard);
    Serial.print("\t\tTemp: "); Serial.print(getTemperature());
    Serial.print("\t\tHumidity: "); Serial.print(getHumidity());
    Serial.print("\t\tCo2: "); Serial.print(getCo2());
  }

  // Add to Average
  temperature += getTemperature();
  humidity += getHumidity();
  co2 += getCo2();
  ppm10 = ppm10 + data.pm10_standard; // Standard (_standard) or Environmental (_env) readings for Particulate Data
  ppm25 = ppm25 + data.pm25_standard;
  ppm100 = ppm100 + data.pm100_standard;

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

    // Average Values over recording period
    double avgTemperature = temperature / pollEventCount;
    double avgHumidity = humidity / pollEventCount;
    double avgCo2 = co2 / pollEventCount;
    uint32_t avgPpm10 = ppm10 / pollEventCount;
    uint32_t avgPpm25 = ppm25 / pollEventCount;
    uint32_t avgPpm100 = ppm100 / pollEventCount;

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

  delay(POLLING_FREQUENCY);
}