#ifndef __HTTP_UPDATE_SERVER_H
#define __HTTP_UPDATE_SERVER_H

#include <StreamString.h>
#include <Update.h>
#include <WebServer.h>

static const char serverIndex[] PROGMEM =
  R"(<!DOCTYPE html>
     <html lang='en'>
     <head>
         <meta charset='utf-8'>
         <meta name='viewport' content='width=device-width,initial-scale=1'/>
     </head>
     <body>
     <form method='POST' action='' enctype='multipart/form-data'>
         Firmware:<br>
         <input type='file' accept='.bin,.bin.gz' name='firmware'>
         <input type='submit' value='Update Firmware'>
     </form>
     <form method='POST' action='' enctype='multipart/form-data'>
         FileSystem:<br>
         <input type='file' accept='.bin,.bin.gz,.image' name='filesystem'>
         <input type='submit' value='Update FileSystem'>
     </form>
     </body>
     </html>)";
static const char successResponse[] PROGMEM = "<META http-equiv=\"refresh\" content=\"15;URL=/\">Update Success! Rebooting...";
static const char *csrfHeaders[2] = {"Origin", "Host"};

class HTTPUpdateServer {
public:
  HTTPUpdateServer(bool serial_debug = false) {
    _serial_output = serial_debug;
    _server = NULL;
    _username = emptyString;
    _password = emptyString;
    _authenticated = false;
  }

  void setup(WebServer *server) {
    setup(server, emptyString, emptyString);
  }

  void setup(WebServer *server, const String &path) {
    setup(server, path, emptyString, emptyString);
  }

  void setup(WebServer *server, const String &username, const String &password) {
    setup(server, "/update", username, password);
  }

  void setup(WebServer *server, const String &path, const String &username, const String &password) {

    _server = server;
    _username = username;
    _password = password;

    // collect headers for CSRF verification
    _server->collectHeaders(csrfHeaders, 2);

    // handler for the /update form page
    _server->on(path.c_str(), HTTP_GET, [&]() {
      if (_username != emptyString && _password != emptyString && !_server->authenticate(_username.c_str(), _password.c_str())) {
        return _server->requestAuthentication();
      }
      _server->send_P(200, PSTR("text/html"), serverIndex);
    });

    // handler for the /update form POST (once file upload finishes)
    _server->on(
      path.c_str(), HTTP_POST,
      [&]() {
        if (!_authenticated) {
          if (_username == emptyString || _password == emptyString) {
            _server->send(200, F("text/html"), String(F("Update error: Wrong origin received!")));
            return;
          }
          return _server->requestAuthentication();
        }
        if (Update.hasError()) {
          _server->send(200, F("text/html"), String(F("Update error: ")) + _updaterError);
        } else {
          _server->client().setNoDelay(true);
          _server->send_P(200, PSTR("text/html"), successResponse);
          delay(100);
          _server->client().stop();
          ESP.restart();
        }
      },
      [&]() {
        // handler for the file upload, gets the sketch bytes, and writes
        // them through the Update object
        HTTPUpload &upload = _server->upload();

        if (upload.status == UPLOAD_FILE_START) {
          _updaterError.clear();
          if (_serial_output) {
            Serial.setDebugOutput(true);
          }

          _authenticated = (_username == emptyString || _password == emptyString || _server->authenticate(_username.c_str(), _password.c_str()));
          if (!_authenticated) {
            if (_serial_output) {
              Serial.printf("Unauthenticated Update\n");
            }
            return;
          }

          String origin = _server->header(String(csrfHeaders[0]));
          String host = _server->header(String(csrfHeaders[1]));
          String expectedOrigin = String("http://") + host;
          if (origin != expectedOrigin) {
            if (_serial_output) {
              Serial.printf("Wrong origin received! Expected: %s, Received: %s\n", expectedOrigin.c_str(), origin.c_str());
            }
            _authenticated = false;
            return;
          }

          if (_serial_output) {
            Serial.printf("Update: %s\n", upload.filename.c_str());
          }
          if (upload.name == "filesystem") {
            if (!Update.begin(UPDATE_SIZE_UNKNOWN, U_FLASHFS)) {  //Instead of SPIFFS.totalBytes(). Fix https://github.com/espressif/arduino-esp32/issues/9967
              if (_serial_output) {
                Update.printError(Serial);
              }
            }
          } else {
            uint32_t maxSketchSpace = (ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000;
            if (!Update.begin(maxSketchSpace, U_FLASH)) {  //start with max available size
              _setUpdaterError();
            }
          }
        } else if (_authenticated && upload.status == UPLOAD_FILE_WRITE && !_updaterError.length()) {
          if (_serial_output) {
            Serial.printf(".");
          }
          if (Update.write(upload.buf, upload.currentSize) != upload.currentSize) {
            _setUpdaterError();
          }
        } else if (_authenticated && upload.status == UPLOAD_FILE_END && !_updaterError.length()) {
          if (Update.end(true)) {  //true to set the size to the current progress
            if (_serial_output) {
              Serial.printf("Update Success: %zu\nRebooting...\n", upload.totalSize);
            }
          } else {
            _setUpdaterError();
          }
          if (_serial_output) {
            Serial.setDebugOutput(false);
          }
        } else if (_authenticated && upload.status == UPLOAD_FILE_ABORTED) {
          Update.end();
          if (_serial_output) {
            Serial.println("Update was aborted");
          }
        }
        delay(0);
      }
    );
  }

  void updateCredentials(const String &username, const String &password) {
    _username = username;
    _password = password;
  }

protected:
  void _setUpdaterError() {
    if (_serial_output) {
      Update.printError(Serial);
    }
    StreamString str;
    Update.printError(str);
    _updaterError = str.c_str();
  }

private:
  bool _serial_output;
  WebServer *_server;
  String _username;
  String _password;
  bool _authenticated;
  String _updaterError;
};

#endif
