#include <Arduino.h>
#include <WiFi.h>
#include <AsyncTCP.h>
#include <ESPAsyncWebServer.h>
#include <Adafruit_NeoPixel.h>
#include <HTTPClient.h>
#include <Update.h>
#include <ArduinoJson.h>
#include <Preferences.h>
#include "mbedtls/sha256.h"
 
const char* ssid     = "UPB-Guest";
const char* password = "";
 
// Pull-OTA settings (host these on your local machine)
const char* fw_version       = "1.0.0";
const char* ota_manifest_url = "http://192.168.0.104:8000/manifest.json"; // adjust to your LAN host
const char* ota_telemetry_url = "http://192.168.0.104:8000/telemetry";     // optional: logs update results
constexpr unsigned long OTA_CHECK_INTERVAL_MS = 30UL * 1000UL;
unsigned long lastOtaCheckMs = 0;
 
constexpr uint8_t NEOPIXEL_PIN = 3;
constexpr uint16_t NEOPIXEL_COUNT = 1;
Adafruit_NeoPixel statusPixel(NEOPIXEL_COUNT, NEOPIXEL_PIN, NEO_GRB + NEO_KHZ800);
 
AsyncWebServer server(80);
Preferences prefs;
 
String getPrefStringSafe(const char* key) {
  if (!prefs.isKey(key)) {
    return "";
  }
  return prefs.getString(key, "");
}
 
struct OtaManifest {
  String version;
  String url;
  String sha256;
};
 
struct TelemetryState {
  String lastResult;   // success | failed | noop | ""
  String lastError;    // sha_mismatch | manifest_fail | ... | ""
  String fromVersion;
  String toVersion;
  String message;      // human-readable summary
} telemetry;
 
void loadTelemetryFromNvs() {
  prefs.begin("ota", false);
  telemetry.lastResult = getPrefStringSafe("result");
  telemetry.lastError = getPrefStringSafe("error");
  telemetry.fromVersion = getPrefStringSafe("from");
  telemetry.toVersion = getPrefStringSafe("to");
  telemetry.message = getPrefStringSafe("msg");
}
 
void storeTelemetryToNvs() {
  prefs.putString("result", telemetry.lastResult);
  prefs.putString("error", telemetry.lastError);
  prefs.putString("from", telemetry.fromVersion);
  prefs.putString("to", telemetry.toVersion);
  prefs.putString("msg", telemetry.message);
}
 
void recordTelemetry(const char* result, const char* error, const String& fromVer, const String& toVer, const String& msg) {
  telemetry.lastResult = result ? result : "";
  telemetry.lastError = error ? error : "";
  telemetry.fromVersion = fromVer;
  telemetry.toVersion = toVer;
  telemetry.message = msg;
  storeTelemetryToNvs();
}
 
String toHex(const uint8_t* data, size_t len) {
  static const char* hex = "0123456789abcdef";
  String out;
  out.reserve(len * 2);
  for (size_t i = 0; i < len; ++i) {
    out += hex[(data[i] >> 4) & 0x0F];
    out += hex[data[i] & 0x0F];
  }
  return out;
}
 
String urlEncode(const String& input) {
  String encoded;
  for (size_t i = 0; i < input.length(); i++) {
    char c = input[i];
    if (isalnum(static_cast<unsigned char>(c)) || c == '-' || c == '_' || c == '.' || c == '~') {
      encoded += c;
    } else if (c == ' ') {
      encoded += "%20";
    } else {
      char buf[4];
      snprintf(buf, sizeof(buf), "%%%02X", static_cast<unsigned char>(c));
      encoded += buf;
    }
  }
  return encoded;
}
 
void sendTelemetry(const char* status, const String& detail) {
  if (ota_telemetry_url == nullptr || strlen(ota_telemetry_url) == 0) {
    return;
  }
 
  HTTPClient http;
  String url = String(ota_telemetry_url) +
               "?status=" + status +
               "&version=" + fw_version +
               "&msg=" + urlEncode(detail.substring(0, 60)); // keep it short
  if (http.begin(url)) {
    http.GET(); // best-effort; ignore response
    http.end();
  }
}
 
bool fetchManifest(OtaManifest& manifest) {
  HTTPClient http;
  if (!http.begin(ota_manifest_url)) {
    sendTelemetry("manifest_begin_fail", "http.begin failed");
    recordTelemetry("failed", "manifest_begin_fail", fw_version, "", "Manifest fetch failed");
    return false;
  }
 
  const int code = http.GET();
  if (code != HTTP_CODE_OK) {
    sendTelemetry("manifest_http_fail", "code " + String(code));
    recordTelemetry("failed", "manifest_http_fail", fw_version, "", "HTTP " + String(code));
    http.end();
    return false;
  }
 
  StaticJsonDocument<512> doc;
  DeserializationError err = deserializeJson(doc, http.getString());
  http.end();
  if (err) {
    sendTelemetry("manifest_parse_fail", err.f_str());
    recordTelemetry("failed", "manifest_parse_fail", fw_version, "", "Manifest parse failed");
    return false;
  }
 
  manifest.version = doc["version"] | "";
  manifest.url = doc["url"] | "";
  manifest.sha256 = doc["sha256"] | "";
 
  if (manifest.version.isEmpty() || manifest.url.isEmpty() || manifest.sha256.isEmpty()) {
    sendTelemetry("manifest_missing", "missing field");
    recordTelemetry("failed", "manifest_missing", fw_version, "", "Manifest missing field");
    return false;
  }
 
  return true;
}
 
bool downloadAndUpdate(const String& firmwareUrl, const String& expectedSha256, const String& targetVersion) {
  WiFiClient client;
  HTTPClient http;
  if (!http.begin(client, firmwareUrl)) {
    sendTelemetry("fw_begin_fail", "http.begin");
    recordTelemetry("failed", "fw_begin_fail", fw_version, targetVersion, "Firmware HTTP begin failed");
    return false;
  }
 
  const int code = http.GET();
  if (code != HTTP_CODE_OK) {
    sendTelemetry("fw_http_fail", "code " + String(code));
    recordTelemetry("failed", "fw_http_fail", fw_version, targetVersion, "HTTP " + String(code));
    http.end();
    return false;
  }
 
  const int contentLength = http.getSize();
  if (contentLength <= 0 && contentLength != -1) {
    sendTelemetry("fw_size_fail", "invalid length");
    recordTelemetry("failed", "fw_size_fail", fw_version, targetVersion, "Invalid content length");
    http.end();
    return false;
  }
 
  if (!Update.begin(contentLength > 0 ? contentLength : UPDATE_SIZE_UNKNOWN)) {
    sendTelemetry("update_begin_fail", Update.errorString());
    recordTelemetry("failed", "update_begin_fail", fw_version, targetVersion, "Update.begin failed");
    http.end();
    return false;
  }
 
  mbedtls_sha256_context shaCtx;
  mbedtls_sha256_init(&shaCtx);
  mbedtls_sha256_starts(&shaCtx, 0); // 0 = SHA-256
 
  WiFiClient* stream = http.getStreamPtr();
  uint8_t buff[1024];
  size_t written = 0;
  while (http.connected()) {
    const size_t avail = stream->available();
    if (avail) {
      const size_t toRead = avail > sizeof(buff) ? sizeof(buff) : avail;
      const size_t read = stream->readBytes(buff, toRead);
      if (read == 0) {
        break;
      }
      if (Update.write(buff, read) != read) {
        sendTelemetry("update_write_fail", Update.errorString());
        recordTelemetry("failed", "update_write_fail", fw_version, targetVersion, "Write failed");
        Update.abort();
        mbedtls_sha256_free(&shaCtx);
        http.end();
        return false;
      }
      mbedtls_sha256_update(&shaCtx, buff, read);
      written += read;
    } else {
      delay(1);
      if (contentLength > 0 && written >= static_cast<size_t>(contentLength)) {
        break;
      }
    }
  }
 
  uint8_t hashResult[32];
  mbedtls_sha256_finish(&shaCtx, hashResult);
  mbedtls_sha256_free(&shaCtx);
  http.end();
 
  const String computedHash = toHex(hashResult, sizeof(hashResult));
  if (!computedHash.equalsIgnoreCase(expectedSha256)) {
    sendTelemetry("hash_mismatch", computedHash);
    recordTelemetry("failed", "sha_mismatch", fw_version, targetVersion, "SHA mismatch");
    Update.abort();
    return false;
  }
 
  if (!Update.end(true)) { // true = even if size is smaller
    sendTelemetry("update_end_fail", Update.errorString());
    recordTelemetry("failed", "update_end_fail", fw_version, targetVersion, "Update end failed");
    return false;
  }
 
  if (!Update.isFinished()) {
    sendTelemetry("update_not_finished", "not finished");
    recordTelemetry("failed", "update_not_finished", fw_version, targetVersion, "Update not finished");
    return false;
  }
 
  sendTelemetry("update_ok", "restarting");
   recordTelemetry("success", "", fw_version, targetVersion, "Update applied; restarting");
  delay(100);
  ESP.restart();
  return true; // not reached
}
 
void checkForOtaUpdate() {
  OtaManifest manifest;
  if (!fetchManifest(manifest)) {
    return;
  }
 
  if (manifest.version == fw_version) {
    recordTelemetry("noop", "", fw_version, manifest.version, "No update needed.");
    return; // already up to date
  }
 
  sendTelemetry("update_found", "v" + manifest.version);
  downloadAndUpdate(manifest.url, manifest.sha256, manifest.version);
}
 
void setup() {
  Serial.begin(115200);
  delay(200);
 
  loadTelemetryFromNvs();
 
  WiFi.mode(WIFI_STA);
  WiFi.begin(ssid, password);
 
  Serial.print("Connecting to WiFi");
  while (WiFi.status() != WL_CONNECTED) {
    delay(300);
    Serial.print(".");
  }
  Serial.println();
 
  Serial.print("Connected. IP: ");
  Serial.println(WiFi.localIP());
 
  // Simple landing page
  server.on("/", HTTP_GET, [](AsyncWebServerRequest *request) {
    request->send(200, "text/plain",
      "ESP32-C6 Sparrow OTA ready.\n"
      "Updates pull from manifest server.\n"
      "Go to /status for telemetry.");
  });
 
  server.on("/status", HTTP_GET, [](AsyncWebServerRequest *request) {
    String payload;
    payload.reserve(256);
    payload += "{";
    payload += "\"device_version\":\"" + String(fw_version) + "\",";
    payload += "\"last_result\":\"" + telemetry.lastResult + "\",";
    payload += "\"last_error\":\"" + telemetry.lastError + "\",";
    payload += "\"from_version\":\"" + telemetry.fromVersion + "\",";
    payload += "\"to_version\":\"" + telemetry.toVersion + "\",";
    payload += "\"message\":\"" + telemetry.message + "\"";
    payload += "}";
    request->send(200, "application/json", payload);
  });
 
  statusPixel.begin();
  statusPixel.setBrightness(40);
  statusPixel.show(); // initialize to off
 
  server.begin();
  Serial.println("HTTP server started.");
  Serial.println("Open http://<device-ip>/status");
 
  // First pull-based OTA check after boot
  checkForOtaUpdate();
}
 
void loop() {
  static bool ledOn = false;
  static unsigned long lastToggle = 0;
  const unsigned long blinkIntervalMs = 500;
 
  const unsigned long now = millis();
  if (now - lastToggle >= blinkIntervalMs) {
    lastToggle = now;
    ledOn = !ledOn;
    const uint32_t color = ledOn ? statusPixel.Color(255, 0, 0) : 0;
    statusPixel.setPixelColor(0, color);
    statusPixel.show();
  }
 
  if (WiFi.isConnected() && now - lastOtaCheckMs >= OTA_CHECK_INTERVAL_MS) {
    lastOtaCheckMs = now;
    checkForOtaUpdate();
  }
 
  // Yield to keep the watchdog happy
  delay(1);
}