#include "image_pipeline.h"
#include <M5GFX.h>
#include <lgfx/utility/lgfx_tjpgd.h>
#include <esp_heap_caps.h>
#include <cstring>

// Spectra 6 palette
static const uint8_t PALETTE_RGB[6][3] = {
    {PALETTE_BLACK_R, PALETTE_BLACK_G, PALETTE_BLACK_B},
    {PALETTE_WHITE_R, PALETTE_WHITE_G, PALETTE_WHITE_B},
    {PALETTE_RED_R, PALETTE_RED_G, PALETTE_RED_B},
    {PALETTE_GREEN_R, PALETTE_GREEN_G, PALETTE_GREEN_B},
    {PALETTE_BLUE_R, PALETTE_BLUE_G, PALETTE_BLUE_B},
    {PALETTE_YELLOW_R, PALETTE_YELLOW_G, PALETTE_YELLOW_B}
};

struct JpegMemReader {
    const uint8_t* data;
    size_t size;
    size_t pos;
};

static uint32_t jpegMemInput(void* device, uint8_t* buf, uint32_t len) {
    auto* ctx = static_cast<JpegMemReader*>(device);
    if (ctx->pos >= ctx->size) {
        return 0;
    }
    uint32_t remain = static_cast<uint32_t>(ctx->size - ctx->pos);
    uint32_t n = len < remain ? len : remain;
    memcpy(buf, ctx->data + ctx->pos, n);
    ctx->pos += n;
    return n;
}

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

    // Decode JPEG
    uint16_t srcW, srcH;
    uint8_t* rgb = decodeJpeg(jpegData, jpegSize, &srcW, &srcH);
    if (rgb == nullptr) {
        Serial.println("[pipeline] JPEG decode failed");
        return result;
    }

    Serial.printf("[pipeline] Decoded: %dx%d\n", srcW, srcH);

    // Calculate crop region (fill-crop to display aspect ratio)
    uint16_t cropX, cropY, cropW, cropH;
    centerCrop(rgb, srcW, srcH, DISPLAY_WIDTH, DISPLAY_HEIGHT,
               &cropX, &cropY, &cropW, &cropH);

    // Resize cropped region to display dimensions
    uint8_t* cropped = (uint8_t*)ps_malloc(cropW * cropH * 3);
    if (cropped == nullptr) {
        free(rgb);
        return result;
    }

    // Extract crop region
    for (uint16_t y = 0; y < cropH; y++) {
        memcpy(cropped + y * cropW * 3,
               rgb + ((cropY + y) * srcW + cropX) * 3,
               cropW * 3);
    }
    free(rgb);

    // Resize to display dimensions
    uint8_t* resized = resize(cropped, cropW, cropH, DISPLAY_WIDTH, DISPLAY_HEIGHT);
    free(cropped);

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

    // Dither to 6-color palette
    uint8_t* dithered = dither(resized, DISPLAY_WIDTH, DISPLAY_HEIGHT);
    free(resized);

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

    result.framebuffer = dithered;
    result.valid = true;
    Serial.println("[pipeline] Processing complete");
    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};

    // Each portrait gets half the width minus gap
    uint16_t portraitW = (DISPLAY_WIDTH - PORTRAIT_GAP_PX) / 2;
    uint16_t portraitH = DISPLAY_HEIGHT;

    // Allocate combined RGB buffer
    uint8_t* combined = (uint8_t*)ps_calloc(DISPLAY_WIDTH * DISPLAY_HEIGHT * 3, 1);
    if (combined == nullptr) return result;

    // Process first portrait
    uint16_t src1W, src1H;
    uint8_t* rgb1 = decodeJpeg(jpeg1Data, jpeg1Size, &src1W, &src1H);
    if (rgb1 != nullptr) {
        uint16_t cropX, cropY, cropW, cropH;
        centerCrop(rgb1, src1W, src1H, portraitW, portraitH,
                   &cropX, &cropY, &cropW, &cropH);

        uint8_t* cropped1 = (uint8_t*)ps_malloc(cropW * cropH * 3);
        if (cropped1) {
            for (uint16_t y = 0; y < cropH; y++) {
                memcpy(cropped1 + y * cropW * 3,
                       rgb1 + ((cropY + y) * src1W + cropX) * 3, cropW * 3);
            }
            uint8_t* resized1 = resize(cropped1, cropW, cropH, portraitW, portraitH);
            free(cropped1);
            if (resized1) {
                // Copy into left side of combined buffer
                for (uint16_t y = 0; y < portraitH; y++) {
                    memcpy(combined + y * DISPLAY_WIDTH * 3,
                           resized1 + y * portraitW * 3, portraitW * 3);
                }
                free(resized1);
            }
        }
        free(rgb1);
    }

    // Process second portrait
    uint16_t src2W, src2H;
    uint8_t* rgb2 = decodeJpeg(jpeg2Data, jpeg2Size, &src2W, &src2H);
    if (rgb2 != nullptr) {
        uint16_t cropX, cropY, cropW, cropH;
        centerCrop(rgb2, src2W, src2H, portraitW, portraitH,
                   &cropX, &cropY, &cropW, &cropH);

        uint8_t* cropped2 = (uint8_t*)ps_malloc(cropW * cropH * 3);
        if (cropped2) {
            for (uint16_t y = 0; y < cropH; y++) {
                memcpy(cropped2 + y * cropW * 3,
                       rgb2 + ((cropY + y) * src2W + cropX) * 3, cropW * 3);
            }
            uint8_t* resized2 = resize(cropped2, cropW, cropH, portraitW, portraitH);
            free(cropped2);
            if (resized2) {
                // Copy into right side of combined buffer
                uint16_t offsetX = portraitW + PORTRAIT_GAP_PX;
                for (uint16_t y = 0; y < portraitH; y++) {
                    memcpy(combined + (y * DISPLAY_WIDTH + offsetX) * 3,
                           resized2 + y * portraitW * 3, portraitW * 3);
                }
                free(resized2);
            }
        }
        free(rgb2);
    }

    // Dither combined buffer
    uint8_t* dithered = dither(combined, DISPLAY_WIDTH, DISPLAY_HEIGHT);
    free(combined);

    if (dithered == nullptr) return result;

    result.framebuffer = dithered;
    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::decodeJpeg(uint8_t* data, size_t size,
                                    uint16_t* outWidth, uint16_t* outHeight) {
    JpegMemReader reader = {data, size, 0};
    lgfxJdec jdec;
    static constexpr uint16_t poolSize = 3900;
    uint8_t* pool = (uint8_t*)malloc(poolSize);
    if (pool == nullptr) {
        return nullptr;
    }

    JRESULT prep = lgfx_jd_prepare(&jdec, jpegMemInput, pool, poolSize, &reader);
    if (prep != JDR_OK) {
        free(pool);
        return nullptr;
    }

    uint16_t w = jdec.width;
    uint16_t h = jdec.height;
    free(pool);

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

    if (!sprite.createSprite(w, h)) {
        return nullptr;
    }

    if (!sprite.drawJpg(data, size, 0, 0)) {
        sprite.deleteSprite();
        return nullptr;
    }

    size_t bufSize = static_cast<size_t>(w) * h * 3;
    uint8_t* rgb = (uint8_t*)ps_malloc(bufSize);
    if (rgb == nullptr) {
        sprite.deleteSprite();
        return nullptr;
    }

    for (uint16_t y = 0; y < h; y++) {
        for (uint16_t x = 0; x < w; x++) {
            lgfx::bgr888_t color = sprite.readPixelRGB(x, y);
            size_t idx = (static_cast<size_t>(y) * w + x) * 3;
            rgb[idx] = color.r;
            rgb[idx + 1] = color.g;
            rgb[idx + 2] = color.b;
        }
    }

    sprite.deleteSprite();
    *outWidth = w;
    *outHeight = h;
    return rgb;
}

uint8_t* ImagePipeline::resize(uint8_t* rgb, uint16_t srcW, uint16_t srcH,
                                uint16_t dstW, uint16_t dstH) {
    size_t bufSize = dstW * dstH * 3;
    uint8_t* dst = (uint8_t*)ps_malloc(bufSize);
    if (dst == nullptr) return nullptr;

    if (dstW <= 1 || dstH <= 1 || srcW == 0 || srcH == 0) {
        free(dst);
        return nullptr;
    }

    // Bilinear interpolation
    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 = rgb[(y0 * srcW + x0) * 3 + c] * (1 - xFrac) +
                           rgb[(y0 * srcW + x1) * 3 + c] * xFrac;
                float bot = rgb[(y1 * srcW + x0) * 3 + c] * (1 - xFrac) +
                           rgb[(y1 * srcW + x1) * 3 + c] * xFrac;
                float val = top * (1 - yFrac) + bot * yFrac;
                dst[(y * dstW + x) * 3 + c] = (uint8_t)(val + 0.5f);
            }
        }
    }

    return dst;
}

void ImagePipeline::centerCrop(uint8_t* rgb, uint16_t srcW, uint16_t srcH,
                                uint16_t targetW, uint16_t targetH,
                                uint16_t* cropX, uint16_t* cropY,
                                uint16_t* cropW, uint16_t* cropH) {
    (void)rgb;
    float targetAspect = (float)targetW / (float)targetH;
    float srcAspect = (float)srcW / (float)srcH;

    if (srcAspect > targetAspect) {
        // Source is wider — crop sides
        *cropH = srcH;
        *cropW = (uint16_t)(srcH * targetAspect);
        *cropX = (srcW - *cropW) / 2;
        *cropY = 0;
    } else {
        // Source is taller — crop top/bottom
        *cropW = srcW;
        *cropH = (uint16_t)(srcW / targetAspect);
        *cropX = 0;
        *cropY = (srcH - *cropH) / 2;
    }
}

uint8_t* ImagePipeline::dither(uint8_t* rgb, uint16_t width, uint16_t height) {
    size_t pixelCount = width * height;
    uint8_t* output = (uint8_t*)ps_malloc(pixelCount);
    if (output == nullptr) return nullptr;

    // Work buffer with int16 to handle error overflow
    int16_t* work = (int16_t*)ps_malloc(pixelCount * 3 * sizeof(int16_t));
    if (work == nullptr) {
        free(output);
        return nullptr;
    }

    // Copy to work buffer
    for (size_t i = 0; i < pixelCount * 3; i++) {
        work[i] = rgb[i];
    }

    // Floyd-Steinberg dithering
    for (uint16_t y = 0; y < height; y++) {
        for (uint16_t x = 0; x < width; x++) {
            size_t idx = (y * width + x) * 3;
            int r = constrain(work[idx], 0, 255);
            int g = constrain(work[idx + 1], 0, 255);
            int b = constrain(work[idx + 2], 0, 255);

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

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

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

    free(work);
    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_RGB[i][0];
        int dg = g - PALETTE_RGB[i][1];
        int db = b - PALETTE_RGB[i][2];
        int dist = dr * dr + dg * dg + db * db;
        if (dist < bestDist) {
            bestDist = dist;
            best = static_cast<uint8_t>(i);
        }
    }
    return best;
}