#include "ota.h"
#include "http.h"
#include "json.h"
#include "log.h"
#include "printf.h"
#include "tls.h"
#include "url.h"
#include "util.h"

#ifndef MG_OTA_MAX_URL_LEN
#define MG_OTA_MAX_URL_LEN 256
#endif

// Scannable version tag embedded in every firmware binary, for server-side
// version extraction. The version string starts after the "MG_VERSION:" prefix.
static const char s_fw_version[] = "MG_VERSION:" MG_FIRMWARE_VERSION;

static bool s_autocommit_ok;  // True after OTA server confirms "same version"

static struct mg_ota_state {
  char json_url[MG_OTA_MAX_URL_LEN];
  char url[MG_OTA_MAX_URL_LEN];
  size_t size;
  uint8_t sha256[32];
  void (*fn)(const char *error_message);
  size_t received;  // Firmware bytes written so far
  bool open;        // mg_ota_begin() succeeded, mg_ota_end() is pending
  const char *err;  // First error during download, NULL if none
} *s_ota;

static void s_firmware_fn(struct mg_connection *c, int ev, void *ev_data);

static bool s_tls_init(struct mg_connection *c, const char *url) {
  struct mg_tls_opts opts;
  memset(&opts, 0, sizeof(opts));
  opts.ca = mg_str(MG_OTA_TLS_CA);
  opts.name = mg_url_host(url);
  if (opts.ca.len == 0) {
    mg_error(c, "MG_OTA_TLS_CA is not configured");
  } else {
    mg_tls_init(c, &opts);
  }
  return c->is_tls_hs != 0 && !c->is_closing;
}

#if MG_ENABLE_CUSTOM_DEVICE_ID
#else
void mg_ota_device_id(char *buf, size_t len) {
#if defined(UID_BASE) &&                                               \
    (defined(__SYSTEM_STM32F4XX_H) || defined(__SYSTEM_STM32F7XX_H) || \
     defined(SYSTEM_STM32H5XX_H) || defined(SYSTEM_STM32H7XX_H) ||     \
     defined(SYSTEM_STM32N6XX_H) || defined(SYSTEM_STM32U5XX_H))
  uint32_t *p = (uint32_t *) UID_BASE;
  mg_snprintf(buf, len, "stm32_%08x%08x%08x", p[0], p[1], p[2]);
#elif MG_ARCH == MG_ARCH_PICOSDK
  pico_unique_board_id_t x = {0};
  pico_get_unique_board_id(&x);
  mg_snprintf(buf, len, "rp_%02x%02x%02x%02x%02x%02x%02x%02x", x.id[0], x.id[1],
              x.id[2], x.id[3], x.id[4], x.id[5], x.id[6], x.id[7]);
#elif defined(OCOTP_BASE) && OCOTP_BASE == 0x401F4000u
  mg_snprintf(buf, len, "rt10xx_%08x%08x", OCOTP->CFG0, OCOTP->CFG1);
#elif defined(OCOTP_BASE) && OCOTP_BASE == 0x40CAC000u
  mg_snprintf(buf, len, "rt11xx_%08x%08x", OCOTP->FUSEN[16].FUSE,
              OCOTP->FUSEN[17].FUSE);
#elif MG_ARCH == MG_ARCH_ESP32
  uint8_t mac[6] = {0};
  esp_efuse_mac_get_default(mac);
  mg_snprintf(buf, len, "esp32_%02x%02x%02x%02x%02x%02x", mac[0], mac[1],
              mac[2], mac[3], mac[4], mac[5]);
#else
  mg_snprintf(buf, len, "%d", 0);
#endif
}
#endif

static void s_version_fn(struct mg_connection *c, int ev, void *ev_data) {
  uint64_t expiration = *(uint64_t *) c->data;
  if (s_ota == NULL) {
    c->fn = NULL;
    c->is_closing = 1;
    return;
  }
  if (ev == MG_EV_POLL) {
    if (mg_millis() > expiration) mg_error(c, "Metadata timeout");
  } else if (ev == MG_EV_CONNECT) {
    char id[40];
    struct mg_str host = mg_url_host(s_ota->json_url);
    const char *uri = mg_url_uri(s_ota->json_url);
    const char *sep = strchr(uri, '?') == NULL ? "?" : "&";
    if (mg_url_is_ssl(s_ota->json_url)) {
      if (!s_tls_init(c, s_ota->json_url)) return;
    }
    mg_ota_device_id(id, sizeof(id));
    id[sizeof(id) - 1] = '\0';
    mg_printf(
        c,
        "GET %s%sarch=%d&version=%s&id=%s&interval=%d&boot=%d HTTP/1.1\r\n"
        "Host: %.*s\r\n"
        "Connection: close\r\n\r\n",
        uri, sep, MG_ARCH, s_fw_version + 11, id, MG_OTA_PULL_INTERVAL_SECONDS,
        MG_OTA_STATE_GET(), host.len, host.buf);
  } else if (ev == MG_EV_HTTP_MSG) {
    struct mg_http_message *hm = (struct mg_http_message *) ev_data;
    char version[MG_OTA_MAX_VERSION_LEN];
    double result;
    MG_DEBUG(("Got metadata: %.*s", hm->body.len, hm->body.buf));
    if (mg_http_status(hm) != 200 || mg_json_get(hm->body, "$", NULL) != 0 ||
        !mg_json_unescape(hm->body, "$.version", version, sizeof(version)) ||
        !mg_json_unescape(hm->body, "$.url", s_ota->url, sizeof(s_ota->url)) ||
        !mg_json_get_num(hm->body, "$.size", &result) || result <= 0) {
      char buf[100];
      mg_snprintf(buf, sizeof(buf), "Bad metadata: %.*s", hm->body.len,
                  hm->body.buf);
      s_ota->fn(buf);
      mg_free(s_ota);
      s_ota = NULL;
    } else if (strcmp(version, s_fw_version + 11) == 0) {
      s_autocommit_ok = true;
      s_ota->fn("Same version");
      mg_free(s_ota);
      s_ota = NULL;
    } else {
      struct mg_connection *fc;
      s_ota->size = (size_t) result;
      // TODO (robertc2000): parse and validate sha256
      MG_DEBUG(("Firmware version: %s, url: %s, size: %ld", version, s_ota->url,
                s_ota->size));
      fc = mg_http_connect(c->mgr, s_ota->url, s_firmware_fn, NULL);
      if (fc == NULL) {
        s_ota->fn("Failed to connect");
        mg_free(s_ota);
        s_ota = NULL;
      } else {
        *(uint64_t *) fc->data = mg_millis() + 5 * 1000;  // Set expiration
      }
    }
    c->is_closing = 1;
  } else if (ev == MG_EV_ERROR) {
    s_ota->fn((char *) ev_data);
    mg_free(s_ota);
    s_ota = NULL;
  }
}

static void status_fn(const char *errmsg) {
  if (errmsg) MG_ERROR(("OTA failed: %s", errmsg));
}

// Firmware download callback for mg_http_stream_body(). Writes chunks to the
// OTA target and reports the result via s_ota->fn(). The end call also happens
// when the connection drops before the whole body arrived.
static bool s_firmware_cb(struct mg_http_message *hm, struct mg_str *data,
                          void **p) {
  bool finalized = false, ok;
  (void) p;
  if (s_ota == NULL) return false;
  if (hm != NULL) {  // Start
    s_ota->received = 0;
    s_ota->open = mg_ota_begin(hm->body.len);
    if (!s_ota->open) s_ota->err = "OTA begin failed";
    return s_ota->open;
  } else if (data != NULL) {  // Next chunk
    if (!mg_ota_write(data->buf, data->len)) {
      s_ota->err = "OTA write failed";
      return false;
    }
    s_ota->received += data->len;
    return true;
  }
  // End
  if (s_ota->open) finalized = mg_ota_end();
  if (s_ota->err == NULL && s_ota->received != s_ota->size) {
    s_ota->err = "Connection closed";
  }
  if (s_ota->err == NULL && !finalized) s_ota->err = "OTA finalize failed";
  ok = s_ota->err == NULL;
  s_ota->fn(s_ota->err);
  mg_free(s_ota);
  s_ota = NULL;
  return ok;
}

static void s_firmware_fn(struct mg_connection *c, int ev, void *ev_data) {
  uint64_t expiration = *(uint64_t *) c->data;
  if (s_ota == NULL) {
    c->fn = NULL;
    c->is_closing = 1;
    return;
  }
  if (ev == MG_EV_POLL) {
    if (mg_millis() > expiration) mg_error(c, "OTA timeout");
  } else if (ev == MG_EV_CONNECT) {
    struct mg_str host = mg_url_host(s_ota->url);
    if (mg_url_is_ssl(s_ota->url)) {
      if (!s_tls_init(c, s_ota->url)) return;
    }
    mg_printf(c,
              "GET %s HTTP/1.1\r\n"
              "Host: %.*s\r\n"
              "Connection: close\r\n\r\n",
              mg_url_uri(s_ota->url), (int) host.len, host.buf);
    *(uint64_t *) c->data = mg_millis() + 300 * 1000;  // Set expiration
  } else if (ev == MG_EV_HTTP_HDRS) {
    struct mg_http_message *hm = (struct mg_http_message *) ev_data;
    int status = mg_http_status(hm);
    if (status != 200 || hm->body.len != s_ota->size) {
      mg_error(c, "Bad HTTP response: status %d, size %lu vs %lu", status,
               (unsigned long) hm->body.len, (unsigned long) s_ota->size);
    } else {
      MG_DEBUG(("Beginning OTA (%lu bytes)", (unsigned long) s_ota->size));
      mg_http_stream_body(c, hm, s_firmware_cb, NULL);
    }
  } else if (ev == MG_EV_ERROR) {
    s_ota->fn((char *) ev_data);
    mg_free(s_ota);
    s_ota = NULL;
  }
}

void mg_ota_url_check(struct mg_mgr *mgr, const char *json_url,
                      void (*fn)(const char *error_message)) {
  if (fn == NULL) fn = status_fn;
  if (s_ota != NULL) {
    fn("OTA already in progress");
  } else if ((s_ota = (struct mg_ota_state *) mg_calloc(1, sizeof(*s_ota))) ==
             NULL) {
    fn("Out of memory");
  } else {
    struct mg_connection *c;
    mg_snprintf(s_ota->json_url, sizeof(s_ota->json_url), "%s", json_url);
    MG_DEBUG(("Connecting to %s", json_url));
    c = mg_http_connect(mgr, s_ota->json_url, s_version_fn, NULL);
    if (c == NULL) {
      mg_free(s_ota);
      s_ota = NULL;
      fn("Failed to connect");
    } else {
      s_ota->fn = fn;
      *(uint64_t *) c->data = mg_millis() + 5 * 1000;  // Set expiration
    }
  }
}

// Uses *p as a flag: non-NULL means an OTA session is open and must be
// closed by mg_ota_end()
bool mg_http_ota_cb(struct mg_http_message *hm, struct mg_str *data, void **p) {
  if (hm != NULL) {  // Start
    if (!mg_ota_begin(hm->body.len)) return false;
    *p = (void *) (ptrdiff_t) 1;
    return true;
  } else if (data != NULL) {  // Next chunk
    return mg_ota_write(data->buf, data->len);
  }
  if (*p == NULL) return false;  // End. Begin failed, nothing to finish
  *p = NULL;
  return mg_ota_end();
}

void mg_ota_poll(struct mg_mgr *mgr) {
  static uint64_t t = 5000;  // Fire first time 5 sec after boot

  static uint64_t feed_timer;  // Advances 500ms per tick; tracks elapsed time
  if (MG_OTA_STATE_GET() == MG_OTA_FAILED &&
      mg_timer_expired(&feed_timer, 500, mg_millis())) {
    if (feed_timer < (uint64_t) MG_OTA_ROLLBACK_TIMEOUT_SECONDS * 1000) {
      MG_OTA_ROLLBACK_TIMER_FEED();  // Feed watchdog every 500ms
    } else if (s_autocommit_ok) {
      MG_INFO(("Auto-committing firmware"));
      MG_OTA_STATE_SET(MG_OTA_CONFIRMED);  // Stop feeding → IWDG resets cleanly
    } else {
      MG_INFO(("No commit confirmation, rolling back"));
      // Stop feeding → IWDG fires → resets into MG_OTA_FAILED → rollback
    }
  }

  if (MG_OTA_URL != NULL &&
      mg_timer_expired(&t, MG_OTA_PULL_INTERVAL_SECONDS * 1000, mg_millis())) {
    mg_ota_url_check(mgr, MG_OTA_URL, MG_OTA_STATUS_FN);
  }
}