#pragma once
#include <Arduino.h>
#include <Preferences.h>
#include <LittleFS.h>
#include "config.h"

#ifdef UNIT_TEST
// In unit tests, use mock stubs for hardware-dependent headers.
// The test build adds test/mocks to the include path via -iquote.
#include "epd_mock.h"
#include "esp_sleep.h"
#include "driver/gpio.h"
#else
#include "epd.h"
#include <esp_sleep.h>
#include <driver/gpio.h>
#include <mbedtls/sha256.h>
#endif

// ── SHA-256 of a LittleFS file ────────────────────────────────────────────────
// Used to verify the server-supplied X-Image-Sha256 against what we actually
// wrote to flash. ESP32-S3 has a hardware SHA accelerator; mbedtls picks it up
// automatically. Returns lowercase hex (matches PHP's hash_file('sha256') and
// the X-Image-Sha256 header style). Empty string on failure or in unit tests
// (where mbedtls isn't linked) — callers treat that as "skip verification".
inline String sha256_of_file(const char* path) {
#ifdef UNIT_TEST
    (void)path;
    return String("");
#else
    File f = LittleFS.open(path, "r");
    if (!f) return String("");

    mbedtls_sha256_context ctx;
    mbedtls_sha256_init(&ctx);
    mbedtls_sha256_starts(&ctx, 0);  // 0 = SHA-256 (1 = SHA-224)

    uint8_t buf[1024];
    while (f.available()) {
        int n = f.read(buf, sizeof(buf));
        if (n <= 0) break;
        mbedtls_sha256_update(&ctx, buf, (size_t)n);
    }

    uint8_t digest[32];
    mbedtls_sha256_finish(&ctx, digest);
    mbedtls_sha256_free(&ctx);
    f.close();

    char hex[65];
    for (int i = 0; i < 32; i++) sprintf(&hex[i * 2], "%02x", digest[i]);
    hex[64] = '\0';
    return String(hex);
#endif
}

#ifndef UNIT_TEST
// Defined in main.cpp
static void show_setup_qr(const String& mac);
#else
// Stub for native tests — tracks call count
extern int g_show_setup_qr_count;
inline void show_setup_qr(const String& mac) { g_show_setup_qr_count++; }
#endif

// ── Utility: derive AP SSID from MAC ─────────────────────────────────────────
// Strips colons, uppercases, takes the last 4 chars.
// Builds via std::string so single-char append is unambiguous on all targets.
inline String ap_ssid_from_mac(const String& mac) {
    std::string cleaned;
    const char* p = mac.c_str();
    while (*p) {
        if (*p != ':') cleaned += (char)toupper((unsigned char)*p);
        ++p;
    }
    std::string suffix = cleaned.substr(cleaned.size() - 4);
    return String(("WeVisto-" + suffix).c_str());
}

// ── WiFi connection attempt ───────────────────────────────────────────────────
inline bool attempt_wifi(const char* ssid, const char* pass) {
    WiFi.mode(WIFI_STA);
    WiFi.begin(ssid, pass);
    uint32_t start = millis();
    while (true) {
        int s = WiFi.status();
        if (s == WL_CONNECTED) return true;
        // Bail the moment the radio reports a terminal failure — bad PSK
        // surfaces as WL_CONNECT_FAILED and missing SSID as WL_NO_SSID_AVAIL
        // within a few seconds. Without this the user stares at the yellow
        // Step 1/2 for the full WIFI_TIMEOUT_MS before the red retry repaints.
        if (s == WL_CONNECT_FAILED || s == WL_NO_SSID_AVAIL) return false;
        if (millis() - start > WIFI_TIMEOUT_MS) return false;
        delay(200);
    }
}

// ── Reset button hold detection ───────────────────────────────────────────────
inline bool check_reset_button() {
    uint32_t hold_start = millis();
    while (digitalRead(PIN_BTN_RESET) == LOW) {
        if (millis() - hold_start >= RESET_HOLD_MS) {
            return true;
        }
        delay(50);
    }
    return false;
}

template<typename HTTP>
void normal_operation_impl(const String& mac, HTTP& http, const String& url, Preferences& prefs) {
    Serial.println("[op] start GET " + url);
    prefs.begin(NVS_NAMESPACE, true);
    int32_t currentImgId = prefs.getInt(NVS_KEY_IMG_ID, -1);
    bool drawNeeded = prefs.getInt(NVS_KEY_DRAW_NEEDED, 0) != 0;
    bool errBorder  = prefs.getInt(NVS_KEY_ERR_BORDER, 0) != 0;
    int  schemaV    = prefs.getInt(NVS_KEY_SCHEMA_V, 0);
    bool justProvisioned = prefs.getInt(NVS_KEY_JUST_PROVISIONED, 0) != 0;
    // TEMP: power-monitor — previous cycle's awake + panel-init durations.
    // Reported as X-Prev-Awake-Ms / X-Prev-Panel-Init-Ms below. Remove
    // along with the corresponding writes near deep_sleep_start + the
    // epd_init() timing block once the PIN_PWR cut is validated.
    uint32_t prevAwakeMs     = prefs.getUInt(NVS_KEY_PREV_AWAKE_MS, 0);
    uint32_t prevPanelInitMs = prefs.getUInt(NVS_KEY_PREV_PANEL_INIT_MS, 0);
    prefs.end();

    // Schema migration: on first boot under err-border-aware firmware, the
    // display may be holding a stale full-screen yellow from the old buggy
    // epd_fill(YELLOW) path. The old firmware never wrote NVS_KEY_ERR_BORDER,
    // so we'd have no other signal that a clean repaint is needed. Force one
    // by treating this boot as if errBorder were set, then bump schema_v so
    // it doesn't fire again.
    if (schemaV < NVS_SCHEMA_VERSION) {
        Serial.println(String("[op] schema migration v") + String(schemaV) + " -> v" +
                       String(NVS_SCHEMA_VERSION) + ", forcing one-shot recovery redraw");
        errBorder = true;
        prefs.begin(NVS_NAMESPACE, false);
        prefs.putInt(NVS_KEY_SCHEMA_V, NVS_SCHEMA_VERSION);
        prefs.end();
    }

    if (currentImgId >= 0) {
        http.addHeader("X-Current-Image-Id", String(currentImgId));
    }

    // Report the panel hardware so the server can route image-rendering
    // dimensions to the right DeviceModel. Comes from -DPANEL_ID in the
    // active env's build_flags (see config.h). Sent on every poll so the
    // server can correct a mis-set Device.model lazily without needing
    // a separate registration handshake.
    http.addHeader("X-Panel-Id", PANEL_ID);

    // Tell the server how we got here. The server uses this to honor a
    // power-cycle as a deliberate "force resync" — a poll that arrives with
    // X-Boot-Reason: cold gets a fresh rotation even outside configured wake
    // times, so unplugging and replugging the frame works as a manual refresh.
    // Timer wakes (the normal case) keep their schedule-gated semantics.
    const esp_sleep_wakeup_cause_t cause = esp_sleep_get_wakeup_cause();
    http.addHeader("X-Boot-Reason",
                   cause == ESP_SLEEP_WAKEUP_TIMER ? "timer" : "cold");

    // Tell the server we're freshly provisioned and haven't been
    // acknowledged yet. The server uses this to refuse to serve images
    // for a stale binding (e.g. sold device whose old owner forgot to
    // "Remove this frame") and to send X-Claimed: 1 once the binding is
    // current — at which point we clear the NVS flag below.
    if (justProvisioned) {
        http.addHeader("X-Just-Provisioned", "1");
    }

    // TEMP: power-monitor — last cycle's awake + panel-init times.
    // Lets us see, server-side, whether the PIN_PWR rail cut affects
    // either. Send only if non-zero (skips the first boot after a
    // firmware that didn't store them).
    if (prevAwakeMs > 0) {
        http.addHeader("X-Prev-Awake-Ms", String((unsigned long)prevAwakeMs));
    }
    if (prevPanelInitMs > 0) {
        http.addHeader("X-Prev-Panel-Init-Ms", String((unsigned long)prevPanelInitMs));
    }

    const char* collectHeaders[] = { "X-Interval-Ms", "X-Image-Id", "X-Image-Sha256", "X-Claimed" };
    http.collectHeaders(collectHeaders, 4);
    int code = http.GET();
    Serial.println("[op] GET code=" + String(code));

    // Server confirmed we're claimed → flag clears, regardless of what
    // happened to the response body. Without this, every poll forever
    // would carry X-Just-Provisioned and force the awaiting-claim gate.
    if (justProvisioned) {
        String claimed = http.header("X-Claimed");
        if (claimed == "1") {
            prefs.begin(NVS_NAMESPACE, false);
            prefs.putInt(NVS_KEY_JUST_PROVISIONED, 0);
            prefs.end();
            justProvisioned = false;
        }
    }

    // Honor the server's X-Interval-Ms — that's the user's configured
    // rotationIntervalMinutes / wakeTimes schedule, computed in
    // DeviceImageController::computeIntervalMs. Clamp to sane physical
    // limits so a malformed 0/garbage value doesn't burn the battery
    // (CLAMP_MIN) and a misconfigured "every 999 days" doesn't strand the
    // device for a week (CLAMP_MAX). When no header is present (server
    // bug, mid-deploy), fall back to FETCH_INTERVAL_MS_FALLBACK.
    uint64_t sleepMs = FETCH_INTERVAL_MS_FALLBACK;
    String intervalHdr = http.header("X-Interval-Ms");
    if (intervalHdr.length() > 0) {
        uint64_t v = strtoull(intervalHdr.c_str(), nullptr, 10);
        if (v > 0) {
            sleepMs = v;
            if (sleepMs < SLEEP_CLAMP_MIN_MS) sleepMs = SLEEP_CLAMP_MIN_MS;
            if (sleepMs > SLEEP_CLAMP_MAX_MS) sleepMs = SLEEP_CLAMP_MAX_MS;
        }
    }

    bool displayInitialized = false;

    if (code == 200) {
        String newId  = http.header("X-Image-Id");
        String newSha = http.header("X-Image-Sha256");

        File f = LittleFS.open(IMAGE_PATH, "w", true);
        if (f) { http.writeToStream(&f); f.close(); }
        http.end();

        // Verify integrity. If the server sent a SHA-256 and the bytes we
        // just wrote don't hash to it, the transfer corrupted somewhere
        // between Imagick and our flash — skip the panel update and leave
        // NVS_KEY_IMG_ID alone, so the next poll re-fetches from scratch
        // (the device will keep claiming the old image-id, server will see
        // the mismatch and send 200 again with fresh bytes). Set the err
        // border so the user sees a sync issue instead of garbage.
        bool integrityOk = true;
        if (newSha.length() > 0) {
            String actualSha = sha256_of_file(IMAGE_PATH);
            if (!newSha.equalsIgnoreCase(actualSha)) {
                integrityOk = false;
                Serial.println(String("[op] sha256 mismatch: expected=") + newSha +
                               " actual=" + actualSha + " — discarding image, retry next poll");
            }
        }

        if (!integrityOk) {
            prefs.begin(NVS_NAMESPACE, false);
            prefs.putInt(NVS_KEY_ERR_BORDER, 1);
            prefs.end();
            // Don't touch the panel — it keeps whatever was up before.
            // displayInitialized stays false; epd_sleep() at the bottom is
            // skipped, so no spurious wait_busy timeout on the next cycle.
        } else {
            // Persist ID and set draw_needed before touching the display.
            // If the device loses power during the ~20 s refresh, the flag survives
            // in NVS so the next boot re-draws from LittleFS instead of looping on 200.
            if (newId.length() > 0) {
                prefs.begin(NVS_NAMESPACE, false);
                prefs.putInt(NVS_KEY_DRAW_NEEDED, 1);
                prefs.putInt(NVS_KEY_IMG_ID, newId.toInt());
                prefs.end();
            }

            displayInitialized = true;
            // TEMP: power-monitor — time the panel init, stored in NVS at
            // end of cycle for next boot to report. Remove the timing
            // wrapper when the PIN_PWR cut is validated.
            uint32_t panelInitStart = millis();
            epd_init();
            uint32_t panelInitMs = millis() - panelInitStart;
            File r = LittleFS.open(IMAGE_PATH, "r");
            if (r) { epd_draw_image_from_file(r); r.close(); }

            // Draw complete — clear pending and error-border flags. The fresh
            // image fully overwrites any prior border, so error state is gone.
            prefs.begin(NVS_NAMESPACE, false);
            prefs.putInt(NVS_KEY_DRAW_NEEDED, 0);
            prefs.putInt(NVS_KEY_ERR_BORDER, 0);
            prefs.putUInt(NVS_KEY_PREV_PANEL_INIT_MS, panelInitMs);  // TEMP
            prefs.end();
        }

    } else if (code == 304) {
        http.end();
        // Redraw from LittleFS if either: a previous draw was interrupted
        // (drawNeeded), or a sync-fail border is currently on screen and the
        // server is healthy again (errBorder) — repaint clean to clear it.
        if (drawNeeded || errBorder) {
            Serial.println(String("[op] 304 with recovery flags (drawNeeded=") +
                           String((int)drawNeeded) + " errBorder=" +
                           String((int)errBorder) + ") -> repainting clean from /img.bin");
            File r = LittleFS.open(IMAGE_PATH, "r");
            if (r) {
                displayInitialized = true;
                epd_init();
                epd_draw_image_from_file(r);
                r.close();
                prefs.begin(NVS_NAMESPACE, false);
                prefs.putInt(NVS_KEY_DRAW_NEEDED, 0);
                prefs.putInt(NVS_KEY_ERR_BORDER, 0);
                prefs.end();
                Serial.println("[op] recovery redraw complete; flags cleared");
            } else {
                Serial.println("[op] recovery aborted: /img.bin not in LittleFS");
            }
        }
    } else if (code == 204 || code == 404) {
        // No image to serve. Don't touch the panel — whatever's already
        // displayed is the right thing:
        //   • currentImgId == -1 → the setup QR is up (painted by
        //     enter_provisioning after WiFi save). The 15s bootstrap poll
        //     hits this branch every cycle until the user claims via
        //     /setup/{mac}; redrawing the QR each time would put the panel
        //     in a perpetual ~20s e-ink redraw loop and risk ghosting.
        //   • currentImgId >= 0 → a real photo is up (server hiccup, asset
        //     missing, image deleted). Don't paint the setup QR over the
        //     user's photo; leave the last-good image alone.
        // displayInitialized stays false → epd_sleep() at the bottom is
        // also skipped, since the display was already in sleep from the
        // previous cycle.
        http.end();
    } else {
        // Sync failed (5xx, timeout, TLS handshake failure → code=-1, etc.).
        // Criterion is "does /img.bin exist?", not "is currentImgId >= 0?":
        //  • file exists → a real photo is on the panel. Draw it with a
        //    yellow border so the user sees a sync signal without losing it.
        //    Per FR38, the last-good image must persist.
        //  • no file    → bootstrap window. Panel is showing the setup-claim
        //    QR which the user still needs to scan. A transient TLS hiccup
        //    must NOT wipe that to a yellow fill — it would leave the
        //    recipient with no way to claim the frame until the next 200.
        http.end();
        File r = LittleFS.open(IMAGE_PATH, "r");
        if (r) {
            Serial.println(String("[op] sync fail code=") + String(code) +
                           " -> drawing image with yellow border");
            displayInitialized = true;
            epd_init();
            epd_draw_image_with_border(r, COLOR_YELLOW, BORDER_THICKNESS_PX);
            r.close();
            prefs.begin(NVS_NAMESPACE, false);
            prefs.putInt(NVS_KEY_ERR_BORDER, 1);
            prefs.end();
        } else {
            Serial.println(String("[op] sync fail code=") + String(code) +
                           " with no cached image; preserving setup screen");
        }
    }

    // Only power off the display if it was initialized this cycle. Calling
    // epd_sleep() when the display is already in hardware deep sleep (from the
    // previous cycle) causes wait_busy() to time out at 60 s, wasting the
    // entire poll interval on every 304 response.
    if (displayInitialized) epd_sleep();

    // Bootstrap window: the device has never received an image — the user is
    // still mid-setup, looking at the green setup QR, waiting for the first
    // photo to land. Return WITHOUT arming deep sleep. The caller is expected
    // to re-invoke us on a short timer until imgIdAfter goes >= 0, keeping
    // WiFi up and the device responsive. The previous "15 s deep-sleep
    // between bootstrap polls" wasted ~5 s per cycle on flash boot + WiFi
    // reconnect, so end-user wait until first image was ~20 s × N retries.
    prefs.begin(NVS_NAMESPACE, true);
    int32_t imgIdAfter = prefs.getInt(NVS_KEY_IMG_ID, -1);
    prefs.end();
    if (imgIdAfter < 0) {
        Serial.println("[op] bootstrap: no image yet, caller will retry");
        return;
    }

    esp_sleep_enable_timer_wakeup(sleepMs * 1000ULL);
    // Wake on the BOOT button so the user-facing "hold until the screen
    // starts to flash" reset works even during deep sleep. Without this,
    // the button only did anything during the brief awake window during
    // a poll, and a full sleep cycle of holding did nothing.
    // Pin is GPIO 0 (PIN_BTN_RESET); active-low because BOOT is pulled-up
    // and shorts to ground on press. After EXT0 wakes the chip, setup()
    // runs and check_reset_button() handles the remainder of the hold.
    // A too-short press wakes the device but check_reset_button returns
    // false → normal_operation_impl runs → the next poll fetches a fresh
    // image, which doubles as a "force refresh" gesture.
    esp_sleep_enable_ext0_wakeup((gpio_num_t)PIN_BTN_RESET, 0);
    // Latch any gpio_hold_en pins through the deep sleep period.
    // epd_sleep() cuts PIN_PWR LOW + holds it; without this call the
    // hold releases on the RTC transition and the panel rail comes back
    // up, losing the saving. Released per-pin in epd_setup_pins() on
    // wake via gpio_hold_dis().
    gpio_deep_sleep_hold_en();

    // TEMP: power-monitor — millis() at this point is the total awake
    // duration since boot. Next boot reads it back and reports as
    // X-Prev-Awake-Ms. Remove when PIN_PWR cut is validated.
    prefs.begin(NVS_NAMESPACE, false);
    prefs.putUInt(NVS_KEY_PREV_AWAKE_MS, millis());
    prefs.end();

    esp_deep_sleep_start();
}

// ── Bootstrap-stay-awake loop ────────────────────────────────────────────────
// Wraps normal_operation_impl in a retry loop for the pre-first-image window.
// Once an image arrives, impl calls esp_deep_sleep_start() and never returns —
// in production that's literal silicon-level halt; in unit tests the mock just
// sets g_deep_sleep_started and returns, so we check that flag to break out.
// Caller passes a callable that runs one poll iteration (instantiates a fresh
// HTTPClient bound to the WiFiClient, calls normal_operation_impl, etc.).
template<typename PollOnce>
inline void bootstrap_loop(PollOnce poll_once) {
    while (true) {
        poll_once();
#ifdef UNIT_TEST
        // Production: esp_deep_sleep_start never returns. Tests: it sets the
        // flag and returns, so without this guard the loop spins forever.
        if (g_deep_sleep_started) return;
#endif
        delay(BOOTSTRAP_RETRY_INTERVAL_MS);
    }
}