#include "image_pipeline.h"
#include <M5GFX.h>
#include <esp_heap_caps.h>
#include <cstring>
#include <cmath>

// Spectra 6 calibrated palette (measured display appearance from epdoptimize)
// These represent what the display ACTUALLY shows, used for dithering decisions
static const uint8_t PALETTE_CALIBRATED[6][3] = {
    {0x1F, 0x22, 0x26},  // Black  -> appears as dark gray
    {0xB9, 0xC7, 0xC9},  // White  -> appears as light gray-blue
    {0x62, 0x20, 0x1E},  // Red    -> appears as dark red/brown
    {0x35, 0x56, 0x3A},  // Green  -> appears as dark forest green
    {0x23, 0x3F, 0x8E},  // Blue   -> appears as dark navy
    {0xC1, 0xBB, 0x1E}   // Yellow -> appears as olive/mustard
};

// JPEGDEC draw callback: receives decoded MCU blocks and writes RGB888 to buffer
static int jpegDrawCallback(JPEGDRAW* pDraw) {
    auto* ctx = static_cast<DecodeContext*>(pDraw->pUser);
    if (ctx == nullptr || ctx->rgbBuffer == nullptr) return 0;

    for (int y = 0; y < pDraw->iHeight; y++) {
        int dstRow = pDraw->y + y;
        if (dstRow < 0 || dstRow >= ctx->bufferHeight) continue;

        for (int x = 0; x < pDraw->iWidthUsed; x++) {
            int dstCol = pDraw->x + x;
            if (dstCol < 0 || dstCol >= ctx->bufferWidth) continue;

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

            size_t dstIdx = ((size_t)dstRow * ctx->bufferWidth + dstCol) * 3;
            ctx->rgbBuffer[dstIdx]     = r;
            ctx->rgbBuffer[dstIdx + 1] = g;
            ctx->rgbBuffer[dstIdx + 2] = b;
        }
    }
    return 1;
}

// Contrast/saturation enhancement applied before dithering
static void enhanceContrast(uint8_t* rgb, size_t pixelCount) {
    static constexpr float CONTRAST = 1.25f;
    static constexpr float SATURATION = 1.15f;
    static constexpr float MID = 128.0f;

    for (size_t i = 0; i < pixelCount; i++) {
        size_t idx = i * 3;
        float r = rgb[idx];
        float g = rgb[idx + 1];
        float b = rgb[idx + 2];

        r = (r - MID) * CONTRAST + MID;
        g = (g - MID) * CONTRAST + MID;
        b = (b - MID) * CONTRAST + MID;

        float lum = 0.299f * r + 0.587f * g + 0.114f * b;
        r = lum + (r - lum) * SATURATION;
        g = lum + (g - lum) * SATURATION;
        b = lum + (b - lum) * SATURATION;

        rgb[idx]     = (uint8_t)fminf(fmaxf(r, 0.0f), 255.0f);
        rgb[idx + 1] = (uint8_t)fminf(fmaxf(g, 0.0f), 255.0f);
        rgb[idx + 2] = (uint8_t)fminf(fmaxf(b, 0.0f), 255.0f);
    }
}

// Average a single edge of the fitted image (4 rows or columns deep)
static constexpr int EDGE_DEPTH = 4;

struct EdgeColor { uint8_t r, g, b; };

static EdgeColor averageEdge(const uint8_t* rgb, uint16_t w, uint16_t h,
                             int side) {
    // side: 0=top, 1=bottom, 2=left, 3=right
    uint32_t sumR = 0, sumG = 0, sumB = 0, count = 0;

    switch (side) {
        case 0: { // top rows
            int rows = min((int)h, EDGE_DEPTH);
            for (int y = 0; y < rows; y++)
                for (int x = 0; x < w; x++) {
                    size_t idx = ((size_t)y * w + x) * 3;
                    sumR += rgb[idx]; sumG += rgb[idx+1]; sumB += rgb[idx+2]; count++;
                }
            break;
        }
        case 1: { // bottom rows
            int rows = min((int)h, EDGE_DEPTH);
            for (int y = h - rows; y < h; y++)
                for (int x = 0; x < w; x++) {
                    size_t idx = ((size_t)y * w + x) * 3;
                    sumR += rgb[idx]; sumG += rgb[idx+1]; sumB += rgb[idx+2]; count++;
                }
            break;
        }
        case 2: { // left columns
            int cols = min((int)w, EDGE_DEPTH);
            for (int y = 0; y < h; y++)
                for (int x = 0; x < cols; x++) {
                    size_t idx = ((size_t)y * w + x) * 3;
                    sumR += rgb[idx]; sumG += rgb[idx+1]; sumB += rgb[idx+2]; count++;
                }
            break;
        }
        case 3: { // right columns
            int cols = min((int)w, EDGE_DEPTH);
            for (int y = 0; y < h; y++)
                for (int x = w - cols; x < w; x++) {
                    size_t idx = ((size_t)y * w + x) * 3;
                    sumR += rgb[idx]; sumG += rgb[idx+1]; sumB += rgb[idx+2]; count++;
                }
            break;
        }
    }

    if (count == 0) return {0, 0, 0};
    return {(uint8_t)(sumR / count), (uint8_t)(sumG / count), (uint8_t)(sumB / count)};
}

// In-place horizontal box blur on a rectangular sub-region of a row-major RGB buffer.
// Operates on rows from y0..y1-1, columns x0..x1-1, within a buffer of stride `stride` pixels.
static void boxBlurH(uint8_t* buf, uint16_t stride,
                     uint16_t x0, uint16_t y0, uint16_t x1, uint16_t y1, int radius) {
    int width = x1 - x0;
    if (width <= 0) return;

    // Temp row buffer for one scanline (3 bytes per pixel)
    uint8_t* tmp = (uint8_t*)malloc(width * 3);
    if (tmp == nullptr) return;

    int diam = radius * 2 + 1;

    for (uint16_t y = y0; y < y1; y++) {
        uint8_t* row = buf + ((size_t)y * stride + x0) * 3;

        // Running sum initialization
        int sumR = 0, sumG = 0, sumB = 0;
        for (int i = -radius; i <= radius; i++) {
            int xi = constrain(i, 0, width - 1);
            sumR += row[xi * 3];
            sumG += row[xi * 3 + 1];
            sumB += row[xi * 3 + 2];
        }
        tmp[0] = sumR / diam;
        tmp[1] = sumG / diam;
        tmp[2] = sumB / diam;

        for (int x = 1; x < width; x++) {
            int addIdx = constrain(x + radius, 0, width - 1);
            int remIdx = constrain(x - radius - 1, 0, width - 1);
            sumR += row[addIdx * 3]     - row[remIdx * 3];
            sumG += row[addIdx * 3 + 1] - row[remIdx * 3 + 1];
            sumB += row[addIdx * 3 + 2] - row[remIdx * 3 + 2];
            tmp[x * 3]     = sumR / diam;
            tmp[x * 3 + 1] = sumG / diam;
            tmp[x * 3 + 2] = sumB / diam;
        }

        memcpy(row, tmp, width * 3);
    }
    free(tmp);
}

// In-place vertical box blur on a rectangular sub-region.
static void boxBlurV(uint8_t* buf, uint16_t stride,
                     uint16_t x0, uint16_t y0, uint16_t x1, uint16_t y1, int radius) {
    int height = y1 - y0;
    if (height <= 0) return;

    uint8_t* tmp = (uint8_t*)malloc(height * 3);
    if (tmp == nullptr) return;

    int diam = radius * 2 + 1;

    for (uint16_t x = x0; x < x1; x++) {
        // Running sum initialization
        int sumR = 0, sumG = 0, sumB = 0;
        for (int i = -radius; i <= radius; i++) {
            int yi = y0 + constrain(i, 0, height - 1);
            size_t idx = ((size_t)yi * stride + x) * 3;
            sumR += buf[idx];
            sumG += buf[idx + 1];
            sumB += buf[idx + 2];
        }
        tmp[0] = sumR / diam;
        tmp[1] = sumG / diam;
        tmp[2] = sumB / diam;

        for (int y = 1; y < height; y++) {
            int addY = y0 + constrain(y + radius, 0, height - 1);
            int remY = y0 + constrain(y - radius - 1, 0, height - 1);
            size_t addIdx = ((size_t)addY * stride + x) * 3;
            size_t remIdx = ((size_t)remY * stride + x) * 3;
            sumR += buf[addIdx]     - buf[remIdx];
            sumG += buf[addIdx + 1] - buf[remIdx + 1];
            sumB += buf[addIdx + 2] - buf[remIdx + 2];
            tmp[y * 3]     = sumR / diam;
            tmp[y * 3 + 1] = sumG / diam;
            tmp[y * 3 + 2] = sumB / diam;
        }

        // Write back
        for (int y = 0; y < height; y++) {
            size_t idx = ((size_t)(y0 + y) * stride + x) * 3;
            buf[idx]     = tmp[y * 3];
            buf[idx + 1] = tmp[y * 3 + 1];
            buf[idx + 2] = tmp[y * 3 + 2];
        }
    }
    free(tmp);
}

// 3-pass box blur (approximates Gaussian) on a sub-region
static void gaussianBlurRegion(uint8_t* buf, uint16_t stride,
                               uint16_t x0, uint16_t y0, uint16_t x1, uint16_t y1,
                               int radius) {
    for (int pass = 0; pass < 3; pass++) {
        boxBlurH(buf, stride, x0, y0, x1, y1, radius);
        boxBlurV(buf, stride, x0, y0, x1, y1, radius);
    }
}

// Fill letterbox bars with mirrored + blurred + faded content from the photo edge.
// The fitted image must already be placed in final_buf at (offsetX, offsetY).
static void fillLetterbox(uint8_t* final_buf, uint16_t targetW, uint16_t targetH,
                          const uint8_t* fitted, uint16_t fitW, uint16_t fitH,
                          uint16_t offsetX, uint16_t offsetY) {
    static constexpr int BLUR_RADIUS = 6;

    if (offsetY > 0) {
        // Horizontal bars (top and bottom)
        EdgeColor topEdge = averageEdge(fitted, fitW, fitH, 0);
        EdgeColor botEdge = averageEdge(fitted, fitW, fitH, 1);

        // --- Extend top edge into top bar ---
        // Every row in the bar copies from the photo's topmost row (row 0)
        for (uint16_t y = 0; y < offsetY; y++) {
            for (uint16_t x = 0; x < targetW; x++) {
                int srcX = (int)x - (int)offsetX;
                srcX = constrain(srcX, 0, (int)fitW - 1);
                size_t srcIdx = ((size_t)0 * fitW + srcX) * 3;  // always row 0
                size_t dstIdx = ((size_t)y * targetW + x) * 3;
                final_buf[dstIdx]     = fitted[srcIdx];
                final_buf[dstIdx + 1] = fitted[srcIdx + 1];
                final_buf[dstIdx + 2] = fitted[srcIdx + 2];
            }
        }

        // --- Extend bottom edge into bottom bar ---
        // Every row copies from the photo's bottommost row (fitH - 1)
        uint16_t botStart = offsetY + fitH;
        uint16_t lastRow = fitH - 1;
        for (uint16_t y = botStart; y < targetH; y++) {
            for (uint16_t x = 0; x < targetW; x++) {
                int srcX = (int)x - (int)offsetX;
                srcX = constrain(srcX, 0, (int)fitW - 1);
                size_t srcIdx = ((size_t)lastRow * fitW + srcX) * 3;
                size_t dstIdx = ((size_t)y * targetW + x) * 3;
                final_buf[dstIdx]     = fitted[srcIdx];
                final_buf[dstIdx + 1] = fitted[srcIdx + 1];
                final_buf[dstIdx + 2] = fitted[srcIdx + 2];
            }
        }

        // --- Blur both bars ---
        gaussianBlurRegion(final_buf, targetW, 0, 0, targetW, offsetY, BLUR_RADIUS);
        gaussianBlurRegion(final_buf, targetW, 0, botStart, targetW, targetH, BLUR_RADIUS);

        // --- Fade toward solid edge color ---
        for (uint16_t y = 0; y < offsetY; y++) {
            float alpha = (float)(offsetY - 1 - y) / (float)offsetY;  // 0 at photo, 1 at screen edge
            for (uint16_t x = 0; x < targetW; x++) {
                size_t idx = ((size_t)y * targetW + x) * 3;
                final_buf[idx]     = (uint8_t)(final_buf[idx]     * (1.0f - alpha) + topEdge.r * alpha);
                final_buf[idx + 1] = (uint8_t)(final_buf[idx + 1] * (1.0f - alpha) + topEdge.g * alpha);
                final_buf[idx + 2] = (uint8_t)(final_buf[idx + 2] * (1.0f - alpha) + topEdge.b * alpha);
            }
        }
        for (uint16_t y = botStart; y < targetH; y++) {
            float alpha = (float)(y - botStart) / (float)(targetH - botStart);
            for (uint16_t x = 0; x < targetW; x++) {
                size_t idx = ((size_t)y * targetW + x) * 3;
                final_buf[idx]     = (uint8_t)(final_buf[idx]     * (1.0f - alpha) + botEdge.r * alpha);
                final_buf[idx + 1] = (uint8_t)(final_buf[idx + 1] * (1.0f - alpha) + botEdge.g * alpha);
                final_buf[idx + 2] = (uint8_t)(final_buf[idx + 2] * (1.0f - alpha) + botEdge.b * alpha);
            }
        }

        Serial.printf("[pipeline] Letterbox TB: bars %dpx, top=#%02X%02X%02X bot=#%02X%02X%02X\n",
                      offsetY, topEdge.r, topEdge.g, topEdge.b, botEdge.r, botEdge.g, botEdge.b);
    }

    if (offsetX > 0) {
        // Vertical bars (left and right)
        EdgeColor leftEdge = averageEdge(fitted, fitW, fitH, 2);
        EdgeColor rightEdge = averageEdge(fitted, fitW, fitH, 3);

        // --- Extend left edge into left bar ---
        // Every column in the bar copies from the photo's leftmost column (col 0)
        for (uint16_t y = 0; y < targetH; y++) {
            int srcY = (int)y - (int)offsetY;
            srcY = constrain(srcY, 0, (int)fitH - 1);
            size_t srcIdx = ((size_t)srcY * fitW + 0) * 3;  // always column 0
            for (uint16_t x = 0; x < offsetX; x++) {
                size_t dstIdx = ((size_t)y * targetW + x) * 3;
                final_buf[dstIdx]     = fitted[srcIdx];
                final_buf[dstIdx + 1] = fitted[srcIdx + 1];
                final_buf[dstIdx + 2] = fitted[srcIdx + 2];
            }
        }

        // --- Extend right edge into right bar ---
        // Every column copies from the photo's rightmost column (fitW - 1)
        uint16_t rightStart = offsetX + fitW;
        uint16_t lastCol = fitW - 1;
        for (uint16_t y = 0; y < targetH; y++) {
            int srcY = (int)y - (int)offsetY;
            srcY = constrain(srcY, 0, (int)fitH - 1);
            size_t srcIdx = ((size_t)srcY * fitW + lastCol) * 3;
            for (uint16_t x = rightStart; x < targetW; x++) {
                size_t dstIdx = ((size_t)y * targetW + x) * 3;
                final_buf[dstIdx]     = fitted[srcIdx];
                final_buf[dstIdx + 1] = fitted[srcIdx + 1];
                final_buf[dstIdx + 2] = fitted[srcIdx + 2];
            }
        }

        // --- Blur both bars ---
        gaussianBlurRegion(final_buf, targetW, 0, 0, offsetX, targetH, BLUR_RADIUS);
        gaussianBlurRegion(final_buf, targetW, rightStart, 0, targetW, targetH, BLUR_RADIUS);

        // --- Fade toward solid edge color ---
        for (uint16_t y = 0; y < targetH; y++) {
            for (uint16_t x = 0; x < offsetX; x++) {
                float alpha = (float)(offsetX - 1 - x) / (float)offsetX;
                size_t idx = ((size_t)y * targetW + x) * 3;
                final_buf[idx]     = (uint8_t)(final_buf[idx]     * (1.0f - alpha) + leftEdge.r * alpha);
                final_buf[idx + 1] = (uint8_t)(final_buf[idx + 1] * (1.0f - alpha) + leftEdge.g * alpha);
                final_buf[idx + 2] = (uint8_t)(final_buf[idx + 2] * (1.0f - alpha) + leftEdge.b * alpha);
            }
            for (uint16_t x = rightStart; x < targetW; x++) {
                float alpha = (float)(x - rightStart) / (float)(targetW - rightStart);
                size_t idx = ((size_t)y * targetW + x) * 3;
                final_buf[idx]     = (uint8_t)(final_buf[idx]     * (1.0f - alpha) + rightEdge.r * alpha);
                final_buf[idx + 1] = (uint8_t)(final_buf[idx + 1] * (1.0f - alpha) + rightEdge.g * alpha);
                final_buf[idx + 2] = (uint8_t)(final_buf[idx + 2] * (1.0f - alpha) + rightEdge.b * alpha);
            }
        }

        Serial.printf("[pipeline] Letterbox LR: bars %dpx, left=#%02X%02X%02X right=#%02X%02X%02X\n",
                      offsetX, leftEdge.r, leftEdge.g, leftEdge.b,
                      rightEdge.r, rightEdge.g, rightEdge.b);
    }
}

// Bilinear resize from src (srcW x srcH) to dst (dstW x dstH)
static void bilinearResize(const uint8_t* src, uint16_t srcW, uint16_t srcH,
                           uint8_t* dst, uint16_t dstW, uint16_t dstH) {
    if (dstW <= 1 || dstH <= 1 || srcW <= 1 || srcH <= 1) return;

    float xRatio = (float)(srcW - 1) / (float)(dstW - 1);
    float yRatio = (float)(srcH - 1) / (float)(dstH - 1);

    for (uint16_t y = 0; y < dstH; y++) {
        float srcY = y * yRatio;
        uint16_t y0 = (uint16_t)srcY;
        uint16_t y1 = min((uint16_t)(y0 + 1), (uint16_t)(srcH - 1));
        float yFrac = srcY - y0;

        for (uint16_t x = 0; x < dstW; x++) {
            float srcX = x * xRatio;
            uint16_t x0 = (uint16_t)srcX;
            uint16_t x1 = min((uint16_t)(x0 + 1), (uint16_t)(srcW - 1));
            float xFrac = srcX - x0;

            for (int c = 0; c < 3; c++) {
                float top = src[(y0 * srcW + x0) * 3 + c] * (1 - xFrac) +
                           src[(y0 * srcW + x1) * 3 + c] * xFrac;
                float bot = src[(y1 * srcW + x0) * 3 + c] * (1 - xFrac) +
                           src[(y1 * srcW + x1) * 3 + c] * xFrac;
                float val = top * (1 - yFrac) + bot * yFrac;
                dst[(y * dstW + x) * 3 + c] = (uint8_t)(val + 0.5f);
            }
        }
    }
}

// Parse JPEG markers properly to find SOF and extract dimensions.
static bool parseJpegHeader(const uint8_t* data, size_t size,
                            uint16_t* width, uint16_t* height, bool* isProgressive) {
    *width = 0;
    *height = 0;
    *isProgressive = false;

    if (size < 2 || data[0] != 0xFF || data[1] != 0xD8) return false;

    size_t pos = 2;
    while (pos + 4 < size) {
        if (data[pos] != 0xFF) { pos++; continue; }
        while (pos < size && data[pos] == 0xFF) pos++;
        if (pos >= size) break;

        uint8_t marker = data[pos++];

        if (marker >= 0xC0 && marker <= 0xCF && marker != 0xC4 && marker != 0xCC) {
            if (pos + 7 > size) break;
            *height = (data[pos + 3] << 8) | data[pos + 4];
            *width = (data[pos + 5] << 8) | data[pos + 6];
            *isProgressive = (marker == 0xC2);
            return true;
        }

        if (marker == 0xDA) break;

        if (pos + 1 >= size) break;
        uint16_t segLen = (data[pos] << 8) | data[pos + 1];
        pos += segLen;
    }
    return false;
}

// Fallback decoder using LGFX's built-in TJPGD
static uint8_t* decodeFallbackLGFX(uint8_t* data, size_t size,
                                    uint16_t targetW, uint16_t targetH,
                                    uint16_t imgW, uint16_t imgH,
                                    uint16_t* outW, uint16_t* outH) {
    // Calculate decode dimensions (cap at 960px longest side)
    static constexpr uint16_t MAX_DECODE = 960;
    uint16_t decodeW = imgW;
    uint16_t decodeH = imgH;
    uint16_t maxSide = max(imgW, imgH);

    if (maxSide > MAX_DECODE) {
        float scale = (float)MAX_DECODE / (float)maxSide;
        decodeW = (uint16_t)(imgW * scale);
        decodeH = (uint16_t)(imgH * scale);
    }

    Serial.printf("[pipeline/lgfx] Decoding %dx%d -> %dx%d sprite\n",
                  imgW, imgH, decodeW, decodeH);

    lgfx::LGFX_Sprite sprite;
    sprite.setPsram(true);
    sprite.setColorDepth(lgfx::color_depth_t::rgb888_3Byte);

    if (!sprite.createSprite(decodeW, decodeH)) {
        Serial.println("[pipeline/lgfx] Sprite alloc failed");
        return nullptr;
    }

    sprite.fillSprite(0);
    float scaleX = (float)decodeW / (float)imgW;
    float scaleY = (float)decodeH / (float)imgH;

    if (!sprite.drawJpg(data, size, 0, 0, decodeW, decodeH, 0, 0, scaleX, scaleY)) {
        Serial.println("[pipeline/lgfx] drawJpg failed (possibly progressive)");
        sprite.deleteSprite();
        return nullptr;
    }

    // Fit-contain: scale to fit entirely within target
    float fitScaleW = (float)targetW / (float)decodeW;
    float fitScaleH = (float)targetH / (float)decodeH;
    float fitScale = fminf(fitScaleW, fitScaleH);
    uint16_t fitW = (uint16_t)(decodeW * fitScale);
    uint16_t fitH = (uint16_t)(decodeH * fitScale);

    // Extract full decoded image from sprite
    size_t decodeBufSize = (size_t)decodeW * decodeH * 3;
    uint8_t* decoded = (uint8_t*)ps_malloc(decodeBufSize);
    if (decoded == nullptr) {
        sprite.deleteSprite();
        return nullptr;
    }
    for (uint16_t y = 0; y < decodeH; y++) {
        for (uint16_t x = 0; x < decodeW; x++) {
            lgfx::bgr888_t color = sprite.readPixelRGB(x, y);
            size_t idx = ((size_t)y * decodeW + x) * 3;
            decoded[idx]     = color.r;
            decoded[idx + 1] = color.g;
            decoded[idx + 2] = color.b;
        }
    }
    sprite.deleteSprite();

    // Resize to fit dimensions
    uint8_t* fitted = (uint8_t*)ps_malloc((size_t)fitW * fitH * 3);
    if (fitted == nullptr) {
        free(decoded);
        return nullptr;
    }
    bilinearResize(decoded, decodeW, decodeH, fitted, fitW, fitH);
    free(decoded);

    // Allocate final target buffer (zeroed)
    size_t finalSize = (size_t)targetW * targetH * 3;
    uint8_t* final_buf = (uint8_t*)ps_calloc(targetW * targetH, 3);
    if (final_buf == nullptr) {
        free(fitted);
        return nullptr;
    }

    // Center the fitted image in the final buffer
    uint16_t offsetX = (targetW - fitW) / 2;
    uint16_t offsetY = (targetH - fitH) / 2;
    for (uint16_t y = 0; y < fitH; y++) {
        memcpy(final_buf + ((size_t)(offsetY + y) * targetW + offsetX) * 3,
               fitted + (size_t)y * fitW * 3,
               fitW * 3);
    }

    // Fill letterbox bars with per-edge average colors
    fillLetterbox(final_buf, targetW, targetH, fitted, fitW, fitH, offsetX, offsetY);
    free(fitted);

    Serial.printf("[pipeline/lgfx] Fit %dx%d into %dx%d (offset %d,%d)\n",
                  fitW, fitH, targetW, targetH, offsetX, offsetY);

    *outW = targetW;
    *outH = targetH;
    return final_buf;
}

uint8_t* ImagePipeline::decodeAndFit(uint8_t* data, size_t size,
                                      uint16_t targetW, uint16_t targetH,
                                      uint16_t* outW, uint16_t* outH) {
    // Parse JPEG header to get dimensions and detect progressive
    uint16_t headerW = 0, headerH = 0;
    bool progressive = false;
    if (parseJpegHeader(data, size, &headerW, &headerH, &progressive)) {
        Serial.printf("[pipeline] Header: %dx%d %s\n", headerW, headerH,
                      progressive ? "progressive" : "baseline");
    } else {
        Serial.println("[pipeline] Could not parse JPEG header");
    }

    if (progressive) {
        Serial.println("[pipeline] Progressive JPEG — skipping (server still reprocessing?)");
        return nullptr;
    }

    // Primary path: JPEGDEC
    JPEGDEC jpeg;

    if (!jpeg.openRAM(data, (int)size, jpegDrawCallback)) {
        int err = jpeg.getLastError();
        Serial.printf("[pipeline] JPEGDEC openRAM failed (err %d), trying LGFX fallback\n", err);
        if (headerW > 0 && headerH > 0) {
            return decodeFallbackLGFX(data, size, targetW, targetH,
                                      headerW, headerH, outW, outH);
        }
        return nullptr;
    }

    int imgW = jpeg.getWidth();
    int imgH = jpeg.getHeight();
    Serial.printf("[pipeline] JPEG: %dx%d, type=baseline\n", imgW, imgH);

    // Determine decode scale (reduce large images before fit)
    int decodeOptions = 0;
    int decodeW = imgW;
    int decodeH = imgH;

    if (imgW > targetW * 4 && imgH > targetH * 4) {
        decodeOptions = JPEG_SCALE_QUARTER;
        decodeW = imgW / 4;
        decodeH = imgH / 4;
        Serial.printf("[pipeline] Using 1/4 scale: %dx%d\n", decodeW, decodeH);
    } else if (imgW > targetW * 2 && imgH > targetH * 2) {
        decodeOptions = JPEG_SCALE_HALF;
        decodeW = imgW / 2;
        decodeH = imgH / 2;
        Serial.printf("[pipeline] Using 1/2 scale: %dx%d\n", decodeW, decodeH);
    }

    // Decode full image (no cropping — we want the whole photo)
    size_t decodeBufSize = (size_t)decodeW * decodeH * 3;
    uint8_t* decodedRgb = (uint8_t*)ps_malloc(decodeBufSize);
    if (decodedRgb == nullptr) {
        Serial.printf("[pipeline] Failed to allocate %u KB for decode buffer\n",
                      (unsigned)(decodeBufSize / 1024));
        jpeg.close();
        return nullptr;
    }
    memset(decodedRgb, 0, decodeBufSize);

    DecodeContext ctx;
    ctx.rgbBuffer = decodedRgb;
    ctx.bufferWidth = (uint16_t)decodeW;
    ctx.bufferHeight = (uint16_t)decodeH;

    jpeg.setUserPointer(&ctx);
    jpeg.setPixelType(RGB565_LITTLE_ENDIAN);

    if (!jpeg.decode(0, 0, decodeOptions)) {
        Serial.printf("[pipeline] JPEGDEC decode failed (err %d), trying LGFX fallback\n",
                      jpeg.getLastError());
        free(decodedRgb);
        jpeg.close();
        if (headerW > 0 && headerH > 0) {
            return decodeFallbackLGFX(data, size, targetW, targetH,
                                      headerW, headerH, outW, outH);
        }
        return nullptr;
    }
    jpeg.close();

    // Fit-contain: scale to fit entirely within target dimensions
    float fitScaleW = (float)targetW / (float)decodeW;
    float fitScaleH = (float)targetH / (float)decodeH;
    float fitScale = fminf(fitScaleW, fitScaleH);
    uint16_t fitW = (uint16_t)(decodeW * fitScale);
    uint16_t fitH = (uint16_t)(decodeH * fitScale);

    Serial.printf("[pipeline] Fit: %dx%d -> %dx%d (scale %.2f)\n",
                  decodeW, decodeH, fitW, fitH, fitScale);

    // Resize decoded image to fit dimensions
    uint8_t* fitted = (uint8_t*)ps_malloc((size_t)fitW * fitH * 3);
    if (fitted == nullptr) {
        Serial.println("[pipeline] Fit alloc failed");
        free(decodedRgb);
        return nullptr;
    }
    bilinearResize(decodedRgb, (uint16_t)decodeW, (uint16_t)decodeH, fitted, fitW, fitH);
    free(decodedRgb);

    // Allocate final target-sized buffer (zeroed)
    size_t finalSize = (size_t)targetW * targetH * 3;
    uint8_t* final_buf = (uint8_t*)ps_calloc(targetW * targetH, 3);
    if (final_buf == nullptr) {
        Serial.println("[pipeline] Final buffer alloc failed");
        free(fitted);
        return nullptr;
    }

    // Center the fitted image into the final buffer
    uint16_t offsetX = (targetW - fitW) / 2;
    uint16_t offsetY = (targetH - fitH) / 2;
    for (uint16_t y = 0; y < fitH; y++) {
        memcpy(final_buf + ((size_t)(offsetY + y) * targetW + offsetX) * 3,
               fitted + (size_t)y * fitW * 3,
               fitW * 3);
    }

    // Fill letterbox bars — each bar matches the adjacent photo edge
    fillLetterbox(final_buf, targetW, targetH, fitted, fitW, fitH, offsetX, offsetY);
    free(fitted);

    *outW = targetW;
    *outH = targetH;
    return final_buf;
}

ProcessedImage ImagePipeline::process(uint8_t* jpegData, size_t jpegSize) {
    ProcessedImage result = {nullptr, DISPLAY_WIDTH, DISPLAY_HEIGHT, false};

    uint16_t outW, outH;
    uint8_t* rgb = decodeAndFit(jpegData, jpegSize, DISPLAY_WIDTH, DISPLAY_HEIGHT, &outW, &outH);
    if (rgb == nullptr) {
        Serial.println("[pipeline] Decode/fit failed");
        return result;
    }

    Serial.printf("[pipeline] Got %dx%d fitted RGB\n", outW, outH);

    // Enhance contrast/saturation for e-ink readability
    enhanceContrast(rgb, (size_t)outW * outH);

    // Row-by-row Floyd-Steinberg dithering to 6-color palette
    uint8_t* dithered = ditherRowByRow(rgb, outW, outH);
    free(rgb);

    if (dithered == nullptr) {
        return result;
    }

    result.framebuffer = dithered;
    result.width = outW;
    result.height = outH;
    result.valid = true;
    Serial.printf("[pipeline] Processing complete (%dx%d)\n", outW, outH);
    return result;
}

ProcessedImage ImagePipeline::processPortraitPair(uint8_t* jpeg1Data, size_t jpeg1Size,
                                                    uint8_t* jpeg2Data, size_t jpeg2Size) {
    ProcessedImage result = {nullptr, DISPLAY_WIDTH, DISPLAY_HEIGHT, false};

    uint16_t portraitW = (DISPLAY_WIDTH - PORTRAIT_GAP_PX) / 2;
    uint16_t portraitH = DISPLAY_HEIGHT;

    uint16_t out1W, out1H;
    uint8_t* rgb1 = decodeAndFit(jpeg1Data, jpeg1Size, portraitW, portraitH, &out1W, &out1H);

    uint16_t out2W, out2H;
    uint8_t* rgb2 = decodeAndFit(jpeg2Data, jpeg2Size, portraitW, portraitH, &out2W, &out2H);

    // Allocate combined RGB buffer
    uint8_t* combined = (uint8_t*)ps_calloc(DISPLAY_WIDTH * DISPLAY_HEIGHT, 3);
    if (combined == nullptr) {
        if (rgb1) free(rgb1);
        if (rgb2) free(rgb2);
        return result;
    }

    if (rgb1 != nullptr) {
        for (uint16_t y = 0; y < out1H; y++) {
            memcpy(combined + y * DISPLAY_WIDTH * 3,
                   rgb1 + y * out1W * 3, out1W * 3);
        }
        free(rgb1);
    }

    if (rgb2 != nullptr) {
        uint16_t offsetX = portraitW + PORTRAIT_GAP_PX;
        for (uint16_t y = 0; y < out2H; y++) {
            memcpy(combined + (y * DISPLAY_WIDTH + offsetX) * 3,
                   rgb2 + y * out2W * 3, out2W * 3);
        }
        free(rgb2);
    }

    enhanceContrast(combined, (size_t)DISPLAY_WIDTH * DISPLAY_HEIGHT);

    uint8_t* dithered = ditherRowByRow(combined, DISPLAY_WIDTH, DISPLAY_HEIGHT);
    free(combined);

    if (dithered == nullptr) return result;

    result.framebuffer = dithered;
    result.width = DISPLAY_WIDTH;
    result.height = DISPLAY_HEIGHT;
    result.valid = true;
    return result;
}

void ImagePipeline::freeImage(ProcessedImage& img) {
    if (img.framebuffer) {
        free(img.framebuffer);
        img.framebuffer = nullptr;
    }
    img.valid = false;
}

uint8_t* ImagePipeline::ditherRowByRow(uint8_t* rgb, uint16_t width, uint16_t height) {
    size_t pixelCount = (size_t)width * height;
    uint8_t* output = (uint8_t*)ps_malloc(pixelCount);
    if (output == nullptr) {
        Serial.println("[pipeline] Dither output alloc failed");
        return nullptr;
    }

    size_t rowBytes = (size_t)width * 3 * sizeof(int16_t);
    int16_t* errCurrent = (int16_t*)ps_malloc(rowBytes);
    int16_t* errNext = (int16_t*)ps_malloc(rowBytes);
    if (errCurrent == nullptr || errNext == nullptr) {
        Serial.println("[pipeline] Dither error buffer alloc failed");
        free(output);
        if (errCurrent) free(errCurrent);
        if (errNext) free(errNext);
        return nullptr;
    }

    // Initialize first row
    for (uint16_t x = 0; x < width; x++) {
        size_t srcIdx = x * 3;
        size_t errIdx = x * 3;
        errCurrent[errIdx]     = rgb[srcIdx];
        errCurrent[errIdx + 1] = rgb[srcIdx + 1];
        errCurrent[errIdx + 2] = rgb[srcIdx + 2];
    }

    for (uint16_t y = 0; y < height; y++) {
        // Prepare next row from source
        if (y + 1 < height) {
            size_t nextRowOffset = (size_t)(y + 1) * width * 3;
            for (uint16_t x = 0; x < width; x++) {
                size_t srcIdx = nextRowOffset + x * 3;
                size_t errIdx = x * 3;
                errNext[errIdx]     = rgb[srcIdx];
                errNext[errIdx + 1] = rgb[srcIdx + 1];
                errNext[errIdx + 2] = rgb[srcIdx + 2];
            }
        } else {
            memset(errNext, 0, rowBytes);
        }

        for (uint16_t x = 0; x < width; x++) {
            size_t errIdx = x * 3;

            int r = constrain(errCurrent[errIdx], 0, 255);
            int g = constrain(errCurrent[errIdx + 1], 0, 255);
            int b = constrain(errCurrent[errIdx + 2], 0, 255);

            uint8_t nearest = findNearest(r, g, b);
            output[y * width + x] = nearest;

            int errR = r - PALETTE_CALIBRATED[nearest][0];
            int errG = g - PALETTE_CALIBRATED[nearest][1];
            int errB = b - PALETTE_CALIBRATED[nearest][2];

            if (x + 1 < width) {
                size_t ni = (x + 1) * 3;
                errCurrent[ni]     += errR * 7 / 16;
                errCurrent[ni + 1] += errG * 7 / 16;
                errCurrent[ni + 2] += errB * 7 / 16;
            }
            if (y + 1 < height && x > 0) {
                size_t ni = (x - 1) * 3;
                errNext[ni]     += errR * 3 / 16;
                errNext[ni + 1] += errG * 3 / 16;
                errNext[ni + 2] += errB * 3 / 16;
            }
            if (y + 1 < height) {
                size_t ni = x * 3;
                errNext[ni]     += errR * 5 / 16;
                errNext[ni + 1] += errG * 5 / 16;
                errNext[ni + 2] += errB * 5 / 16;
            }
            if (y + 1 < height && x + 1 < width) {
                size_t ni = (x + 1) * 3;
                errNext[ni]     += errR * 1 / 16;
                errNext[ni + 1] += errG * 1 / 16;
                errNext[ni + 2] += errB * 1 / 16;
            }
        }

        int16_t* tmp = errCurrent;
        errCurrent = errNext;
        errNext = tmp;
    }

    free(errCurrent);
    free(errNext);
    return output;
}

uint8_t ImagePipeline::findNearest(int r, int g, int b) {
    uint8_t best = 0;
    int bestDist = INT32_MAX;
    for (int i = 0; i < DISPLAY_COLORS; i++) {
        int dr = r - PALETTE_CALIBRATED[i][0];
        int dg = g - PALETTE_CALIBRATED[i][1];
        int db = b - PALETTE_CALIBRATED[i][2];
        int dist = dr * dr + dg * dg + db * db;
        if (dist < bestDist) {
            bestDist = dist;
            best = static_cast<uint8_t>(i);
        }
    }
    return best;
}