#include <Arduino.h>
#include <WiFi.h>
#include <HTTPClient.h>
#include <JPEGDEC.h>
#include <GxEPD2_EPD.h>
#include <epd7c/GxEPD2_730c_GDEP073E01.h>
#include <SPI.h>

#include "config.h"
#include "dither.h"

// Display instance
SPIClass displaySPI(SPI2_HOST);
GxEPD2_730c_GDEP073E01 epd(PIN_EPD_CS, PIN_EPD_DC, PIN_EPD_RST, PIN_EPD_BUSY);

// Buffers (allocated in PSRAM)
static uint8_t* jpegBuf    = nullptr;  // raw JPEG data from server
static uint8_t* rgbBuf     = nullptr;  // decoded RGB888 (800*480*3 = 1,152,000 bytes)
static uint8_t* displayBuf = nullptr;  // dithered 4bpp output (800*480/2 = 192,000 bytes)
static size_t   jpegLen    = 0;
static String   photoName  = "";  // name of the photo currently being displayed

// JPEG decoder callback — writes decoded pixels into rgbBuf
static int jpegDrawCallback(JPEGDRAW* pDraw) {
    for (int y = 0; y < pDraw->iHeight; y++) {
        int destY = pDraw->y + y;
        if (destY >= DISPLAY_HEIGHT) continue;

        for (int x = 0; x < pDraw->iWidth; x++) {
            int destX = pDraw->x + x;
            if (destX >= DISPLAY_WIDTH) continue;

            // JPEGDEC outputs RGB565 by default
            uint16_t pixel = pDraw->pPixels[y * pDraw->iWidth + x];
            uint8_t r = ((pixel >> 11) & 0x1F) << 3;
            uint8_t g = ((pixel >> 5)  & 0x3F) << 2;
            uint8_t b = ( pixel        & 0x1F) << 3;

            int idx = (destY * DISPLAY_WIDTH + destX) * 3;
            rgbBuf[idx + 0] = r;
            rgbBuf[idx + 1] = g;
            rgbBuf[idx + 2] = b;
        }
    }
    return 1;  // continue decoding
}

bool connectWiFi() {
    Serial.printf("Connecting to WiFi '%s'...\n", WIFI_SSID);
    WiFi.begin(WIFI_SSID, WIFI_PASSWORD);

    int attempts = 0;
    while (WiFi.status() != WL_CONNECTED && attempts < 60) {
        delay(500);
        Serial.print(".");
        attempts++;
    }

    if (WiFi.status() == WL_CONNECTED) {
        Serial.printf("\nConnected! IP: %s\n", WiFi.localIP().toString().c_str());
        return true;
    }

    Serial.println("\nWiFi connection failed.");
    return false;
}

bool fetchPhoto() {
    HTTPClient http;
    http.begin(PHOTO_SERVER_URL);
    http.setTimeout(30000);
    const char* headerKeys[] = {"X-Photo-Name"};
    http.collectHeaders(headerKeys, 1);

    Serial.printf("Fetching photo from %s\n", PHOTO_SERVER_URL);
    int httpCode = http.GET();

    if (httpCode != HTTP_CODE_OK) {
        Serial.printf("HTTP error: %d\n", httpCode);
        http.end();
        return false;
    }

    // Capture photo name from response header
    photoName = http.header("X-Photo-Name");
    Serial.printf("Photo: %s\n", photoName.c_str());

    jpegLen = http.getSize();
    if (jpegLen <= 0 || jpegLen > 2 * 1024 * 1024) {
        // Unknown size — read into buffer with a max cap
        Serial.println("Reading response (unknown or oversized content-length)...");
        WiFiClient* stream = http.getStreamPtr();
        jpegLen = 0;
        size_t maxSize = 2 * 1024 * 1024;
        while (stream->connected() || stream->available()) {
            size_t avail = stream->available();
            if (avail == 0) { delay(10); continue; }
            size_t toRead = min(avail, maxSize - jpegLen);
            if (toRead == 0) break;
            size_t read = stream->readBytes(jpegBuf + jpegLen, toRead);
            jpegLen += read;
        }
    } else {
        Serial.printf("Content-Length: %u bytes\n", jpegLen);
        WiFiClient* stream = http.getStreamPtr();
        size_t received = 0;
        while (received < jpegLen) {
            size_t avail = stream->available();
            if (avail == 0) { delay(10); continue; }
            size_t read = stream->readBytes(jpegBuf + received, min(avail, jpegLen - received));
            received += read;
        }
    }

    http.end();
    Serial.printf("Received %u bytes of JPEG data\n", jpegLen);
    return jpegLen > 0;
}

bool decodeJpeg() {
    Serial.println("Decoding JPEG...");
    JPEGDEC jpeg;

    if (!jpeg.openRAM(jpegBuf, jpegLen, jpegDrawCallback)) {
        Serial.println("Failed to open JPEG");
        return false;
    }

    Serial.printf("JPEG: %dx%d\n", jpeg.getWidth(), jpeg.getHeight());

    // Use 1/1 scale if the server already resized to 800x480.
    // If the image is larger, we could use 1/2 or 1/4 scale, but
    // the server should handle resizing.
    jpeg.setPixelType(RGB565_BIG_ENDIAN);
    if (!jpeg.decode(0, 0, 0)) {
        Serial.println("JPEG decode failed");
        return false;
    }

    Serial.println("JPEG decoded successfully");
    return true;
}

void displayImage() {
    Serial.println("Dithering to 6-color palette...");
    memset(displayBuf, 0, DISPLAY_WIDTH * DISPLAY_HEIGHT / 2);
    ditherFloydSteinberg(rgbBuf, displayBuf, DISPLAY_WIDTH, DISPLAY_HEIGHT);
    Serial.println("Dithering complete");

    Serial.println("Writing to e-paper display...");
    epd.init(115200, true, 50, false);

    // Write raw 4bpp native-encoded pixel data directly to the panel.
    // invert=true tells GxEPD2 to skip its color remapping since our
    // dither output already uses native panel color indices.
    epd.writeNative(displayBuf, nullptr, 0, 0, DISPLAY_WIDTH, DISPLAY_HEIGHT,
                     true, false, false);
    epd.refresh(false);  // full refresh

    Serial.println("Waiting for display refresh (~15s)...");
    epd.hibernate();
    Serial.println("Display updated!");
}

void sendHeartbeat() {
    HTTPClient http;
    String url = PHOTO_SERVER_URL;
    // Replace /photo with /heartbeat
    url = url.substring(0, url.lastIndexOf('/')) + "/heartbeat";

    http.begin(url);
    http.addHeader("Content-Type", "application/json");

    String body = "{\"photo\":\"" + photoName + "\",\"free_heap\":" + String(ESP.getFreeHeap()) + "}";
    int code = http.POST(body);
    Serial.printf("Heartbeat sent (%d)\n", code);
    http.end();
}

void enterDeepSleep() {
    Serial.printf("Going to deep sleep for %d seconds...\n", SLEEP_DURATION_SEC);
    WiFi.disconnect(true);
    WiFi.mode(WIFI_OFF);

    esp_sleep_enable_timer_wakeup((uint64_t)SLEEP_DURATION_SEC * 1000000ULL);
    esp_deep_sleep_start();
}

void setup() {
    Serial.begin(115200);
    delay(1000);
    Serial.println("\n=== E-Ink Photo Frame ===");

    // Check PSRAM
    if (!psramFound()) {
        Serial.println("ERROR: PSRAM not found! Cannot continue.");
        return;
    }
    Serial.printf("PSRAM: %u bytes free\n", ESP.getFreePsram());

    // Allocate buffers in PSRAM
    jpegBuf    = (uint8_t*)ps_malloc(2 * 1024 * 1024);  // 2MB for JPEG
    rgbBuf     = (uint8_t*)ps_malloc(DISPLAY_WIDTH * DISPLAY_HEIGHT * 3);  // ~1.1MB
    displayBuf = (uint8_t*)ps_malloc(DISPLAY_WIDTH * DISPLAY_HEIGHT / 2);  // 192KB

    if (!jpegBuf || !rgbBuf || !displayBuf) {
        Serial.println("ERROR: Failed to allocate PSRAM buffers!");
        return;
    }

    // Initialize SPI for the display
    displaySPI.begin(PIN_SPI_SCLK, -1, PIN_SPI_MOSI, PIN_EPD_CS);
    epd.selectSPI(displaySPI, SPISettings(20000000, MSBFIRST, SPI_MODE0));

    if (!connectWiFi()) {
        Serial.println("No WiFi — going back to sleep");
        enterDeepSleep();
        return;
    }

    if (!fetchPhoto()) {
        Serial.println("Failed to fetch photo — going back to sleep");
        enterDeepSleep();
        return;
    }

    if (!decodeJpeg()) {
        Serial.println("Failed to decode JPEG — going back to sleep");
        enterDeepSleep();
        return;
    }

    displayImage();
    sendHeartbeat();
    enterDeepSleep();
}

void loop() {
    // Never reached — we deep sleep from setup()
}