#include "jpegdec.h"
#include "display.h"
#include "esp_jpeg_dec.h"
#include "esp_timer.h"
#include "esp_heap_caps.h"
#include "esp_log.h"
#include <string.h>

static const char *TAG = "jpeg";
#define MAX_INTERMEDIATE (2 * 1024 * 1024)

static uint16_t *mid;                    // decoded image before resampling (PSRAM)
static size_t mid_cap;

esp_err_t jpegdec_init(void)
{
    mid_cap = 640 * 360 * 2;
    mid = heap_caps_aligned_alloc(16, mid_cap, MALLOC_CAP_SPIRAM);   // esp_new_jpeg wants 16-byte alignment
    return mid ? ESP_OK : ESP_ERR_NO_MEM;
}

esp_err_t jpegdec_dims(const uint8_t *d, size_t len, int *w, int *h)
{
    if (len < 4 || d[0] != 0xFF || d[1] != 0xD8) return ESP_ERR_INVALID_ARG;
    size_t i = 2;
    while (i + 9 < len) {
        if (d[i] != 0xFF) { i++; continue; }
        uint8_t m = d[i + 1];
        if (m == 0xFF) { i++; continue; }
        if (m == 0xD8 || m == 0x01 || (m >= 0xD0 && m <= 0xD7)) { i += 2; continue; }
        if (m == 0xDA || m == 0xD9) return ESP_ERR_NOT_FOUND;              // hit scan data without a SOF
        if (m == 0xC0 || m == 0xC1) {                                        // baseline / extended sequential
            *h = (d[i + 5] << 8) | d[i + 6];
            *w = (d[i + 7] << 8) | d[i + 8];
            return (*w && *h) ? ESP_OK : ESP_ERR_INVALID_SIZE;
        }
        if (m == 0xC2 || m == 0xC6 || m == 0xCA) return ESP_ERR_NOT_SUPPORTED;   // progressive
        i += 2 + ((d[i + 2] << 8) | d[i + 3]);
    }
    return ESP_ERR_NOT_FOUND;
}

// Bilinear "cover" resample: scale so the image fills the screen, crop the overflow, centre it.
static void resample_cover(const uint16_t *src, int sw, int sh, uint16_t *dst)
{
    const int dw = LCD_H_RES, dh = LCD_V_RES;
    int64_t step_x = ((int64_t)sw << 16) / dw, step_y = ((int64_t)sh << 16) / dh;
    int64_t step = step_x < step_y ? step_x : step_y;                 // same step both axes
    int64_t ox = (((int64_t)sw << 16) - step * dw) / 2 + step / 2 - 32768;
    int64_t oy = (((int64_t)sh << 16) - step * dh) / 2 + step / 2 - 32768;
    const int64_t max_x = (int64_t)(sw - 1) << 16, max_y = (int64_t)(sh - 1) << 16;

    for (int y = 0; y < dh; y++) {
        int64_t sy = oy + y * step;
        if (sy < 0) sy = 0;
        if (sy > max_y) sy = max_y;
        int y0 = sy >> 16, fy = (sy >> 8) & 0xff;
        int y1 = y0 + 1 < sh ? y0 + 1 : y0;
        const uint16_t *r0 = src + y0 * sw, *r1 = src + y1 * sw;
        uint16_t *o = dst + y * dw;
        for (int x = 0; x < dw; x++) {
            int64_t sx = ox + x * step;
            if (sx < 0) sx = 0;
            if (sx > max_x) sx = max_x;
            int x0 = sx >> 16, fx = (sx >> 8) & 0xff;
            int x1 = x0 + 1 < sw ? x0 + 1 : x0;
            uint16_t p00 = r0[x0], p01 = r0[x1], p10 = r1[x0], p11 = r1[x1];
            int w00 = (256 - fx) * (256 - fy), w01 = fx * (256 - fy), w10 = (256 - fx) * fy, w11 = fx * fy;
            int r = ((p00 >> 11) * w00 + (p01 >> 11) * w01 + (p10 >> 11) * w10 + (p11 >> 11) * w11) >> 16;
            int g = (((p00 >> 5) & 63) * w00 + ((p01 >> 5) & 63) * w01 + ((p10 >> 5) & 63) * w10 + ((p11 >> 5) & 63) * w11) >> 16;
            int b = ((p00 & 31) * w00 + (p01 & 31) * w01 + (p10 & 31) * w10 + (p11 & 31) * w11) >> 16;
            o[x] = (r << 11) | (g << 5) | b;
        }
    }
}

static int64_t ms(void) { return esp_timer_get_time() / 1000; }

// esp_new_jpeg path: scale in the decoder to the smallest multiple-of-8 size that still
// covers the screen (uniform factor, min 1/8), then bilinear-resample the last few percent.
static esp_err_t decode_new(const uint8_t *jpeg, size_t len, int w, int h, uint16_t *dst, int64_t *t_dec)
{
    jpeg_dec_config_t cfg = DEFAULT_JPEG_DEC_CONFIG();
    cfg.output_type = JPEG_PIXEL_FORMAT_RGB565_LE;
    float f = LCD_H_RES / (float)w;
    if (LCD_V_RES / (float)h > f) f = LCD_V_RES / (float)h;
    if (f < 0.125f) f = 0.125f;
    int sw = w, sh = h;
    if (f < 1.0f) {
        sw = ((int)(w * f + 7) / 8) * 8;
        sh = ((int)(h * f + 7) / 8) * 8;
        if (sw < 8) sw = 8;
        if (sh < 8) sh = 8;
        cfg.scale.width = sw;
        cfg.scale.height = sh;
    }
    size_t need = (size_t)(sw + 16) * (sh + 16) * 2;
    if (need > MAX_INTERMEDIATE) return ESP_ERR_INVALID_SIZE;
    if (need > mid_cap) {
        uint16_t *bigger = heap_caps_aligned_alloc(16, need, MALLOC_CAP_SPIRAM);
        if (!bigger) return ESP_ERR_NO_MEM;
        heap_caps_free(mid);
        mid = bigger;
        mid_cap = need;
    }
    jpeg_dec_handle_t dec = NULL;
    if (jpeg_dec_open(&cfg, &dec) != JPEG_ERR_OK) return ESP_FAIL;
    jpeg_dec_io_t io = {.inbuf = (uint8_t *)jpeg, .inbuf_len = len};
    jpeg_dec_header_info_t info = {0};
    jpeg_error_t err = jpeg_dec_parse_header(dec, &io, &info);
    int outlen = 0;
    if (err == JPEG_ERR_OK) err = jpeg_dec_get_outbuf_len(dec, &outlen);
    if (err == JPEG_ERR_OK && (size_t)outlen > mid_cap) err = JPEG_ERR_NO_MEM;
    int64_t t0 = ms();
    if (err == JPEG_ERR_OK) {
        io.outbuf = (uint8_t *)mid;
        err = jpeg_dec_process(dec, &io);
    }
    *t_dec = ms() - t0;
    jpeg_dec_close(dec);
    if (err != JPEG_ERR_OK) {
        ESP_LOGW(TAG, "esp_new_jpeg failed (%d) %dx%d -> %dx%d", err, w, h, sw, sh);
        return ESP_FAIL;
    }
    int ow = cfg.scale.width ? sw : w, oh = cfg.scale.height ? sh : h;
    resample_cover(mid, ow, oh, dst);
    return ESP_OK;
}

esp_err_t jpegdec_to_screen(const uint8_t *jpeg, size_t len, uint16_t *dst, int *src_w, int *src_h)
{
    int w, h;
    esp_err_t err = jpegdec_dims(jpeg, len, &w, &h);
    if (err != ESP_OK) return err;
    *src_w = w;
    *src_h = h;
    int64_t t_dec = 0, t0 = ms();
    err = decode_new(jpeg, len, w, h, dst, &t_dec);
    ESP_LOGD(TAG, "%dx%d %u bytes: decode %lld ms, total %lld ms", w, h, (unsigned)len, (long long)t_dec, (long long)(ms() - t0));
    return err;
}