#include <WiFi.h>
#include <Update.h>
#include <WebServer.h>
#include <esp_sleep.h>
#include <Preferences.h>

#define WIFI_SSID "Wokwi-GUEST"
#define WIFI_PASS ""
#define NTP_CYCLE 300  // seconds
#define NTP_TIMEOUT 30  // seconds
#define NTP_SRV "pool.ntp.org"
#define SLEEP_CYCLE 60  // seconds
#define UPDATE_PIN 2

const char* htmlUpdate = R"literal(
<form method='POST' action='/update' enctype='multipart/form-data'>
   <input type='file' name='update'>
   <input type='submit' value='Update'>
</form>
)literal"; // htmlUpdate

RTC_DATA_ATTR uint32_t last_update = 0;
RTC_DATA_ATTR uint32_t drift_save = 0;
RTC_DATA_ATTR uint8_t ntp_count = 0;
RTC_DATA_ATTR uint8_t ntp_errs = 0;

bool restartNow = false;
WebServer server(80);
Preferences prefs;

timeval ntp2tv(uint8_t* packet, uint8_t tvloc) {
  uint64_t sec = 0;
  uint64_t fractional = 0;
  for (int x=0; x<4; x++) sec += packet[tvloc+x] << ((3-x)*8); //get the secs
  sec -= 2208988800; // 70 years (1900-1970)
  for (int x=0; x<4; x++) fractional += packet[tvloc+4+x] << ((3-x)*8); // get the fraction
  fractional = fractional * 1E6 / (1 << 32);
  return (timeval){(long int)sec, (long int)fractional};
}

timeval sntpTime() {
  WiFiUDP udp;
  uint8_t sendPacket[48] = {0};
  sendPacket[0] = 0x1B; 
  // put the current time in the packet?
  udp.begin(123);
  udp.beginPacket(NTP_SRV, 123);
  uint32_t origination = millis();
  udp.write(sendPacket, 48);
  udp.endPacket();
  uint32_t ntpLoops = 150;
  uint32_t ntpDelay = 100;
  for (int z=0; z < ntpLoops; z++) {
    delay(ntpDelay);
    if (udp.parsePacket()) { // we've got a packet
        byte recvPacket[48];
        if (udp.read(recvPacket, 48) == 48) {
          struct timeval recvTime;
          gettimeofday(&recvTime, NULL);
          uint32_t travel = (millis() - origination)/2; // we have to assume symmetric latency
          udp.stop();
          timeval ntpTime = ntp2tv(recvPacket, 40);
          if (travel > ntpLoops * ntpDelay) travel = 0; // extraneous
          if (last_update > 0) {
            int32_t drift;
            drift = (recvTime.tv_sec - ntpTime.tv_sec) * 1000;
            drift += (recvTime.tv_usec - ntpTime.tv_usec) / 1000;
            Serial.printf("Drift: %d\n", drift);
            if (abs(drift) > NTP_CYCLE * 10) {
              Serial.println("Excessive drift");
            } else {              
              prefs.begin("ntp");
              uint32_t ntpcount = prefs.getUInt("count",0);
              ntpcount++;
              uint32_t drift = prefs.getUInt("drift",0);
              uint32_t nowDiff = 0; //This is the diff in setTime and now
              drift = drift + nowDiff / ntpcount;
            }
          }
          return ntpTime;
        }
    }
  }
  return (timeval){0, 0};
}

void WiFiEvent(WiFiEvent_t event, arduino_event_info_t info){
  switch(event){
    case ARDUINO_EVENT_WIFI_STA_DISCONNECTED:
      Serial.println("Disconnected from station, attempting reconnection");
      WiFi.reconnect();
      break;
  }
}

void wifiInit() {
  WiFi.onEvent(WiFiEvent);
  WiFi.begin(WIFI_SSID,WIFI_PASS);
  WiFi.waitForConnectResult();
  if (WiFi.waitForConnectResult() != WL_CONNECTED) {
    log_e("Unable to connect to WiFi");
    return;
  }
  log_i("WiFi connected");
}

void handleUpdate() {
  server.send(200, "text/html", htmlUpdate);
}

void handledoUpdate() {
    HTTPUpload& upload = server.upload();
    if (upload.status == UPLOAD_FILE_START) {
      disableCore0WDT();
      log_v("\nUpdate: %s\n", upload.filename.c_str());
      if (!Update.begin(UPDATE_SIZE_UNKNOWN)) { //start with max available size
        Update.printError(Serial);
      }
    } else if (upload.status == UPLOAD_FILE_WRITE) {
      /* flashing firmware to ESP*/
      if (Update.write(upload.buf, upload.currentSize) != upload.currentSize) {
        Update.printError(Serial);
      }
    } else if (upload.status == UPLOAD_FILE_END) {
      if (Update.end(true)) { //true to set the size to the current progress
        log_v("Update Success: %u\nRebooting...", upload.totalSize);        
      } else {
        Update.printError(Serial);
      }
      server.sendHeader("Refresh", "15");  
      server.sendHeader("Location","/");
      server.send(302);
      delay(1000);
      restartNow = true;
    }
}

void updater() {
  wifiInit();
  server.on("/", HTTP_ANY, [](){handleUpdate();});  
  server.on("/update", HTTP_ANY, [](){handleUpdate();}, [](){handledoUpdate();});
  server.begin();
  log_i("WebServer listening at http://%s", WiFi.localIP().toString());
}

void syncTime() {
  wifiInit();
  Serial.printf("Time: %lu\n", time(NULL));
  timeval mytime = sntpTime();
  if (mytime.tv_sec == 0) {
    ntp_errs++;
    log_e("Unable to acquire NTP time");
    return;
  }
  settimeofday((const timeval*)&mytime, NULL);
  if (mytime.tv_usec < 1000000 - 10000) { // block off until the next second tick
    mytime.tv_sec += 1;
    delayMicroseconds(1000000UL - mytime.tv_usec - 10000);
  }
  uint32_t sys_offset = 0;
  sys_offset = mytime.tv_sec - millis() / 1000UL;
  last_update = time(NULL);
  Serial.printf("NTP time updated:%lu\n", last_update);
  Serial.printf("sys_offset: %d\n", sys_offset);

}

void setup() {
  setCpuFrequencyMhz(240);
  Serial.begin(115200);
  pinMode(UPDATE_PIN, INPUT);
  if (digitalRead(UPDATE_PIN)) {
    updater();
    return;
  }
  
  if (time(NULL) > last_update + NTP_CYCLE) syncTime();
  if (last_update == 0) syncTime();
  
  // do your data collection here

/*  
  esp_sleep_enable_timer_wakeup(SLEEP_CYCLE * 1000000ULL);
  Serial.println("Going to sleep now");
  Serial.flush(); 
  esp_deep_sleep_start();
*/  
}

void loop() {
  if (restartNow) abort();
  server.handleClient();
  if (time(NULL) > last_update + NTP_CYCLE) syncTime();
  delay(10);
}