#include "app.h"

#include <HTTPClient.h>
#include <LittleFS.h>
#include <Update.h>

static const uint8_t PACKAGE_HEADER_SIZE = 32;
static const char PACKAGE_MAGIC[8] = {'T', 'S', 'L', 'U', 'P', 'D', '1', 0};

enum PackageStage : uint8_t {
  PKG_HEADER,
  PKG_FILESYSTEM,
  PKG_FIRMWARE,
  PKG_DONE,
  PKG_ERROR
};

struct PackageState {
  PackageStage stage;
  uint8_t header[PACKAGE_HEADER_SIZE];
  size_t headerRead;
  uint32_t filesystemSize;
  uint32_t firmwareSize;
  uint32_t remaining;
  bool filesystemEnded;
  String error;
};

static PackageState packageState;

static uint32_t readLe32(const uint8_t *p) {
  return (uint32_t)p[0] | ((uint32_t)p[1] << 8) | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

static void resetPackageState() {
  packageState.stage = PKG_HEADER;
  packageState.headerRead = 0;
  packageState.filesystemSize = 0;
  packageState.firmwareSize = 0;
  packageState.remaining = 0;
  packageState.filesystemEnded = false;
  packageState.error = "";
}

static bool failPackage(const String &message) {
  packageState.stage = PKG_ERROR;
  packageState.error = message;
  Update.abort();
  if (packageState.filesystemEnded) LittleFS.begin(false);
  return false;
}

static bool beginPackagePart(uint32_t size, int command, const char *label) {
  if (command == U_SPIFFS) LittleFS.end();
  if (!Update.begin(size, command)) {
    return failPackage(String(label) + " update begin failed: " + Update.errorString());
  }
  packageState.remaining = size;
  return true;
}

static bool finishPackagePart(const char *label) {
  if (!Update.end(true)) {
    return failPackage(String(label) + " update failed: " + Update.errorString());
  }
  return true;
}

static bool parsePackageHeader() {
  if (memcmp(packageState.header, PACKAGE_MAGIC, sizeof(PACKAGE_MAGIC)) != 0) {
    return failPackage("Invalid update package magic");
  }
  uint32_t headerSize = readLe32(packageState.header + 8);
  packageState.filesystemSize = readLe32(packageState.header + 12);
  packageState.firmwareSize = readLe32(packageState.header + 16);
  if (headerSize != PACKAGE_HEADER_SIZE) return failPackage("Unsupported update package header");
  if (!packageState.filesystemSize || !packageState.firmwareSize) return failPackage("Update package must contain filesystem and firmware images");
  packageState.stage = PKG_FILESYSTEM;
  return beginPackagePart(packageState.filesystemSize, U_SPIFFS, "Filesystem");
}

static bool feedPackageBytes(const uint8_t *data, size_t length) {
  while (length && packageState.stage != PKG_ERROR && packageState.stage != PKG_DONE) {
    if (packageState.stage == PKG_HEADER) {
      size_t n = min(length, (size_t)PACKAGE_HEADER_SIZE - packageState.headerRead);
      memcpy(packageState.header + packageState.headerRead, data, n);
      packageState.headerRead += n;
      data += n;
      length -= n;
      if (packageState.headerRead == PACKAGE_HEADER_SIZE && !parsePackageHeader()) return false;
    } else {
      size_t n = min(length, (size_t)packageState.remaining);
      if (Update.write((uint8_t *)data, n) != n) {
        return failPackage(String(packageState.stage == PKG_FILESYSTEM ? "Filesystem" : "Firmware") + " write failed");
      }
      packageState.remaining -= n;
      data += n;
      length -= n;

      if (packageState.remaining == 0) {
        if (packageState.stage == PKG_FILESYSTEM) {
          if (!finishPackagePart("Filesystem")) return false;
          packageState.filesystemEnded = true;
          packageState.stage = PKG_FIRMWARE;
          if (!beginPackagePart(packageState.firmwareSize, U_FLASH, "Firmware")) return false;
        } else {
          if (!finishPackagePart("Firmware")) return false;
          packageState.stage = PKG_DONE;
        }
      }
    }
  }
  if (length && packageState.stage == PKG_DONE) return failPackage("Trailing bytes in update package");
  return packageState.stage != PKG_ERROR;
}

static bool streamPackageFromUrl(const String &url, String &error) {
  resetPackageState();
  HTTPClient http;
  http.begin(url);
  int code = http.GET();
  if (code != HTTP_CODE_OK) {
    http.end();
    error = "Update package URL returned HTTP " + String(code);
    return false;
  }

  uint8_t buffer[1024];
  WiFiClient *stream = http.getStreamPtr();
  int expectedLength = http.getSize();
  int receivedLength = 0;
  uint32_t lastRead = millis();
  while (expectedLength < 0 || receivedLength < expectedLength) {
    int available = stream->available();
    if (available > 0) {
      int wanted = min(available, (int)sizeof(buffer));
      if (expectedLength >= 0) wanted = min(wanted, expectedLength - receivedLength);
      int n = stream->readBytes(buffer, wanted);
      if (n > 0) {
        receivedLength += n;
        lastRead = millis();
        if (!feedPackageBytes(buffer, n)) break;
      }
    } else {
      if (!http.connected()) break;
      if (millis() - lastRead > 30000) {
        failPackage("Update package download timed out");
        break;
      }
      delay(10);
    }
  }
  http.end();

  if (packageState.stage != PKG_DONE) {
    error = packageState.error.length() ? packageState.error : "Incomplete update package";
    return false;
  }
  return true;
}

void handleOtaCheck() {
  if (!authorize()) return;
  HTTPClient http;
  http.begin(updateUrl);
  int code = http.GET();
  int size = http.getSize();
  http.end();
  sendJson(code > 0 && code < 400 ? 200 : 502, jsonOk("\"version\":\"" + String(APP_VERSION) + "\",\"url\":\"" + jsonEscape(updateUrl) + "\",\"httpStatus\":" + String(code) + ",\"contentLength\":" + String(size)));
}

void handleOtaRun() {
  if (!authorize()) return;
  appendLog(LOG_SECURITY_AUDIT, "package URL update started");
  String error;
  if (!streamPackageFromUrl(updateUrl, error)) return sendJson(502, jsonError(error));
  sendJson(200, jsonOk("\"restart\":true,\"filesystemUpdated\":true,\"firmwareUpdated\":true"));
  delay(500);
  ESP.restart();
}

void handleUpdateUploadDone() {
  if (!authorize()) return;
  bool ok = packageState.stage == PKG_DONE;
  String error = packageState.error.length() ? packageState.error : "Incomplete update package";
  sendJson(ok ? 200 : 500, ok ? jsonOk("\"restart\":true,\"filesystemUpdated\":true,\"firmwareUpdated\":true") : jsonError(error));
  if (ok) {
    delay(500);
    ESP.restart();
  }
}

void handleUpdateUploadChunk() {
  HTTPUpload &upload = server.upload();
  if (upload.status == UPLOAD_FILE_START) {
    if (!authorize()) return;
    appendLog(LOG_SECURITY_AUDIT, "package upload started");
    resetPackageState();
  } else if (upload.status == UPLOAD_FILE_WRITE) {
    feedPackageBytes(upload.buf, upload.currentSize);
  }
}