#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"
#else
#include "epd.h"
#include <esp_sleep.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(("PictureFrame-" + 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 (WiFi.status() != WL_CONNECTED) {
        if (millis() - start > WIFI_TIMEOUT_MS) return false;
        delay(200);
    }
    return true;
}

// ── 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) {
    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);
    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));
    }

    // 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");

    const char* collectHeaders[] = { "X-Interval-Ms", "X-Image-Id", "X-Image-Sha256" };
    http.collectHeaders(collectHeaders, 3);
    int code = http.GET();

    // 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;
            epd_init();
            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.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) {
        http.end();
        displayInitialized = true;
        epd_init();
        show_setup_qr(mac);
    } else if (code == 404) {
        http.end();
        displayInitialized = true;
        epd_init();
        show_setup_qr(mac);
    } else {
        // Sync failed (5xx, timeout, malformed). Per FR38, the last-good image
        // must persist; only the border indicates the error. epd_draw_image_with_border
        // falls back to a full fill if the cached file is missing or wrong size,
        // so first-boot error still gets a visible signal.
        http.end();
        displayInitialized = true;
        epd_init();
        File r = LittleFS.open(IMAGE_PATH, "r");
        if (r) {
            Serial.println(String("[op] sync fail code=") + String(code) +
                           " -> drawing image with yellow border");
            epd_draw_image_with_border(r, COLOR_YELLOW, BORDER_THICKNESS_PX);
            r.close();
        } else {
            Serial.println(String("[op] sync fail code=") + String(code) +
                           " -> no cached image, falling back to full yellow fill");
            epd_fill(COLOR_YELLOW);
        }
        prefs.begin(NVS_NAMESPACE, false);
        prefs.putInt(NVS_KEY_ERR_BORDER, 1);
        prefs.end();
    }

    // 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();

    esp_sleep_enable_timer_wakeup(sleepMs * 1000ULL);
    esp_deep_sleep_start();
}