#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>
#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;
    prefs.end();

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

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

    uint64_t sleepMs = FETCH_INTERVAL_MS;
    String intervalHdr = http.header("X-Interval-Ms");
    if (intervalHdr.length() > 0) {
        uint64_t v = strtoull(intervalHdr.c_str(), nullptr, 10);
        if (v > 0) sleepMs = std::min<uint64_t>(v, (uint64_t)FETCH_INTERVAL_MS);
    }

    bool displayInitialized = false;

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

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

        // 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 the pending flag.
        prefs.begin(NVS_NAMESPACE, false);
        prefs.putInt(NVS_KEY_DRAW_NEEDED, 0);
        prefs.end();

    } else if (code == 304) {
        http.end();
        // If a previous draw was interrupted (power loss mid-refresh), the image
        // file is in LittleFS and the ID is correct in NVS — just re-draw it.
        if (drawNeeded) {
            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.end();
            }
        }
    } 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 {
        http.end();
        displayInitialized = true;
        epd_init();
        epd_fill(COLOR_YELLOW);
    }

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