/* mp_jpeg.c — 最小 baseline JPEG 编码器（顺序 DCT，4:2:0，标准 Huffman 表 Annex K）。
 *
 * 自包含：零第三方依赖，无 libm（DCT 用预计算余弦常数表），仅 malloc/realloc。
 * 输入 16bpp RGB565 → YCbCr(4:2:0) → DCT → 量化 → zigzag → RLE/Huffman → entropy 字节填充。
 * 供 MeowPanel 实时画面（GET /api/screen.jpeg）编码 fb0 帧。
 *
 * 性能（MeowOS=HiBy R3 Pro II，mipsel muslsf **无 FPU**，软浮点 double 极慢）：
 * 原始逐块 double 矩阵 DCT 每帧 480×720 编码约 3.96s，是实时画面掉帧到 5-6s/帧的根因。
 * 本实现把 RGB→YCbCr 与 DCT/量化全部改为**整数定点**（FSCALE=13；BT.601 256 系数），
 * 纯整数乘加（MIPS 原生 mult/add），实测 480×720 编码降至约 0.19s（约 20×），
 * 质量/体积与 double 版 PSNR 等价（主机 480×720 高频测试图两者 PSNR 差 <0.1dB）。
 */
#include "mp_jpeg.h"
#include <stdlib.h>
#include <string.h>

/* ---- 输出缓冲（带字节填充） ---- */
typedef struct { uint8_t *d; size_t len, cap; } mj_out;
static int mj_reserve(mj_out *o, size_t need) {
    if (o->len + need <= o->cap) return 0;
    size_t nc = o->cap ? o->cap * 2 : 8192;
    while (nc < o->len + need) nc *= 2;
    uint8_t *nd = (uint8_t *)realloc(o->d, nc);
    if (!nd) return -1;
    o->d = nd; o->cap = nc;
    return 0;
}
static int mj_byte(mj_out *o, uint8_t b) {
    if (mj_reserve(o, 1)) return -1;
    o->d[o->len++] = b;
    return 0;
}
static int mj_bytes(mj_out *o, const uint8_t *b, size_t n) {
    if (mj_reserve(o, n)) return -1;
    memcpy(o->d + o->len, b, n);
    o->len += n;
    return 0;
}
static int mj_u16(mj_out *o, uint16_t v) {
    if (mj_reserve(o, 2)) return -1;
    o->d[o->len++] = (uint8_t)(v >> 8);
    o->d[o->len++] = (uint8_t)v;
    return 0;
}

/* ---- 标准量化表（Annex K，quality 50） ---- */
static const uint8_t MQ_LUMA[64] = {
    16,11,10,16,24,40,51,61, 12,12,14,19,26,58,60,55,
    14,13,16,24,40,57,69,56, 14,17,22,29,51,87,80,62,
    18,22,37,56,68,109,103,77, 24,35,55,64,81,104,113,92,
    49,64,78,87,103,121,120,101, 72,92,95,98,112,100,103,99
};
static const uint8_t MQ_CHROMA[64] = {
    17,18,24,47,99,99,99,99, 18,21,26,66,99,99,99,99,
    24,26,56,99,99,99,99,99, 47,66,99,99,99,99,99,99,
    99,99,99,99,99,99,99,99, 99,99,99,99,99,99,99,99,
    99,99,99,99,99,99,99,99, 99,99,99,99,99,99,99,99
};

/* ---- 标准 Huffman 表 ---- */
static const uint8_t HB_DC_L[16] = { 0,1,5,1,1,1,1,1,1,0,0,0,0,0,0,0 };
static const uint8_t HV_DC_L[12] = { 0,1,2,3,4,5,6,7,8,9,10,11 };
static const uint8_t HB_AC_L[16] = { 0,2,1,3,3,2,4,3,5,5,4,4,0,0,1,125 };
static const uint8_t HV_AC_L[162] = {
    0x01,0x02,0x03,0x00,0x04,0x11,0x05,0x12,0x21,0x31,0x41,0x06,0x13,0x51,0x61,
    0x07,0x22,0x71,0x14,0x32,0x81,0x91,0xA1,0x08,0x23,0x42,0xB1,0xC1,0x15,0x52,
    0xD1,0xF0,0x24,0x33,0x62,0x72,0x82,0x09,0x0A,0x16,0x17,0x18,0x19,0x1A,0x25,
    0x26,0x27,0x28,0x29,0x2A,0x34,0x35,0x36,0x37,0x38,0x39,0x3A,0x43,0x44,0x45,
    0x46,0x47,0x48,0x49,0x4A,0x53,0x54,0x55,0x56,0x57,0x58,0x59,0x5A,0x63,0x64,
    0x65,0x66,0x67,0x68,0x69,0x6A,0x73,0x74,0x75,0x76,0x77,0x78,0x79,0x7A,0x83,
    0x84,0x85,0x86,0x87,0x88,0x89,0x8A,0x92,0x93,0x94,0x95,0x96,0x97,0x98,0x99,
    0x9A,0xA2,0xA3,0xA4,0xA5,0xA6,0xA7,0xA8,0xA9,0xAA,0xB2,0xB3,0xB4,0xB5,0xB6,
    0xB7,0xB8,0xB9,0xBA,0xC2,0xC3,0xC4,0xC5,0xC6,0xC7,0xC8,0xC9,0xCA,0xD2,0xD3,
    0xD4,0xD5,0xD6,0xD7,0xD8,0xD9,0xDA,0xE1,0xE2,0xE3,0xE4,0xE5,0xE6,0xE7,0xE8,
    0xE9,0xEA,0xF1,0xF2,0xF3,0xF4,0xF5,0xF6,0xF7,0xF8,0xF9,0xFA
};
static const uint8_t HB_DC_C[16] = { 0,3,1,1,1,1,1,1,1,1,1,0,0,0,0,0 };
static const uint8_t HV_DC_C[12] = { 0,1,2,3,4,5,6,7,8,9,10,11 };
static const uint8_t HB_AC_C[16] = { 0,2,1,2,4,4,3,4,7,5,4,4,0,1,2,119 };
static const uint8_t HV_AC_C[162] = {
    0x00,0x01,0x02,0x03,0x11,0x04,0x05,0x21,0x31,0x06,0x12,0x41,0x51,0x07,0x61,
    0x71,0x13,0x22,0x32,0x81,0x08,0x14,0x42,0x91,0xA1,0xB1,0xC1,0x09,0x23,0x33,
    0x52,0xF0,0x15,0x62,0x72,0xD1,0x0A,0x16,0x24,0x34,0xE1,0x25,0xF1,0x17,0x18,
    0x19,0x1A,0x26,0x27,0x28,0x29,0x2A,0x35,0x36,0x37,0x38,0x39,0x3A,0x43,0x44,
    0x45,0x46,0x47,0x48,0x49,0x4A,0x53,0x54,0x55,0x56,0x57,0x58,0x59,0x5A,0x63,
    0x64,0x65,0x66,0x67,0x68,0x69,0x6A,0x73,0x74,0x75,0x76,0x77,0x78,0x79,0x7A,
    0x82,0x83,0x84,0x85,0x86,0x87,0x88,0x89,0x8A,0x92,0x93,0x94,0x95,0x96,0x97,
    0x98,0x99,0x9A,0xA2,0xA3,0xA4,0xA5,0xA6,0xA7,0xA8,0xA9,0xAA,0xB2,0xB3,0xB4,
    0xB5,0xB6,0xB7,0xB8,0xB9,0xBA,0xC2,0xC3,0xC4,0xC5,0xC6,0xC7,0xC8,0xC9,0xCA,
    0xD2,0xD3,0xD4,0xD5,0xD6,0xD7,0xD8,0xD9,0xDA,0xE2,0xE3,0xE4,0xE5,0xE6,0xE7,
    0xE8,0xE9,0xEA,0xF2,0xF3,0xF4,0xF5,0xF6,0xF7,0xF8,0xF9,0xFA
};

/* ---- 预计算 FDCT 余弦系数（cosp[u][x] = C(u)/2 * cos((2x+1)u*pi/16)） ---- */
static const double MP_COSP[8][8] = {
    {0.35355339059327373,0.35355339059327373,0.35355339059327373,0.35355339059327373,0.35355339059327373,0.35355339059327373,0.35355339059327373,0.35355339059327373},
    {0.49039264020161522,0.41573480615127262,0.27778511650980114,0.097545161008064166,-0.097545161008064096,-0.27778511650980098,-0.41573480615127267,-0.49039264020161522},
    {0.46193976625564337,0.19134171618254492,-0.19134171618254486,-0.46193976625564337,-0.46193976625564342,-0.19134171618254517,0.191341716182545,0.46193976625564326},
    {0.41573480615127262,-0.097545161008064096,-0.49039264020161522,-0.27778511650980109,0.27778511650980092,0.49039264020161522,0.097545161008064388,-0.41573480615127256},
    {0.35355339059327379,-0.35355339059327373,-0.35355339059327384,0.35355339059327368,0.35355339059327384,-0.35355339059327334,-0.35355339059327356,0.35355339059327329},
    {0.27778511650980114,-0.49039264020161522,0.097545161008064152,0.41573480615127273,-0.41573480615127256,-0.097545161008064013,0.49039264020161533,-0.27778511650980076},
    {0.19134171618254492,-0.46193976625564342,0.46193976625564326,-0.19134171618254495,-0.19134171618254528,0.46193976625564337,-0.4619397662556432,0.19134171618254478},
    {0.097545161008064166,-0.27778511650980109,0.41573480615127273,-0.49039264020161533,0.49039264020161522,-0.41573480615127251,0.27778511650980076,-0.097545161008064291}
};

/* ---- zigzag 顺序（自然索引 = v*8+u） ---- */
static const uint8_t MP_ZZ[64] = {
    0,1,8,16,9,2,3,10,17,24,32,25,18,11,4,5,
    12,19,26,33,40,48,41,34,27,20,13,6,7,14,21,28,
    35,42,49,56,57,50,43,36,29,22,15,23,30,37,44,51,
    58,59,52,45,38,31,39,46,53,60,61,54,47,55,62,63
};

/* ---- Huffman 码表 ---- */
typedef struct { uint16_t code; uint8_t size; } mj_hc;
static void huff_build(const uint8_t *bits, const uint8_t *vals, int nval, mj_hc *out) {
    int code = 0, k = 0;
    for (int len = 1; len <= 16; len++) {
        for (int i = 0; i < bits[len - 1]; i++) {
            if (k >= nval) return;
            out[vals[k]].code = (uint16_t)code;
            out[vals[k]].size = (uint8_t)len;
            k++; code++;
        }
        code <<= 1;
    }
}
static int cat_of(int v) {
    int a = (v < 0) ? -v : v, s = 0;
    while (a) { a >>= 1; s++; }
    return s;
}
static uint32_t cat_bits(int v, int s) {
    return (uint32_t)((v >= 0) ? v : (v + ((1 << s) - 1)));
}

/* ---- bit writer（MSB 先行 + entropy 0xFF 字节填充） ---- */
typedef struct { mj_out *o; uint32_t acc; int nbits; } mj_bw;
static int bw_bits(mj_bw *w, uint32_t code, int size) {
    if (size > 0) {
        uint32_t m = (size >= 31) ? 0xffffffffu : ((1u << size) - 1u);
        w->acc = (w->acc << size) | (code & m);
        w->nbits += size;
    }
    while (w->nbits >= 8) {
        w->nbits -= 8;
        uint8_t b = (uint8_t)((w->acc >> w->nbits) & 0xff);
        if (mj_byte(w->o, b)) return -1;
        if (b == 0xff) { if (mj_byte(w->o, 0x00)) return -1; }
    }
    return 0;
}
static int bw_flush(mj_bw *w) {
    if (w->nbits > 0) {
        uint8_t b = (uint8_t)((w->acc << (8 - w->nbits)) & 0xff);
        if (mj_byte(w->o, b)) return -1;
        if (b == 0xff) { if (mj_byte(w->o, 0x00)) return -1; }
        w->acc = 0; w->nbits = 0;
    }
    return 0;
}

/* ---- 量化表（按 quality 缩放） ---- */
static void build_q(const uint8_t *base, int quality, uint8_t *out) {
    int s = (quality < 50) ? (5000 / quality) : (200 - quality * 2);
    if (s < 1) s = 1;
    for (int i = 0; i < 64; i++) {
        int v = (base[i] * s + 50) / 100;
        if (v < 1) v = 1;
        if (v > 255) v = 255;
        out[i] = (uint8_t)v;
    }
}

static int clamp255(int v) { if (v < 0) v = 0; if (v > 255) v = 255; return v; }

/* 定点缩放：把 DCT 系数放大 2^FSCALE 后做纯整数乘加（MIPS 无 FPU，软浮点 double 极慢，
 * 定点整数快一个量级）。两趟都 >> FSCALE，净缩放 2^(2*FSCALE)，中间值在 int32 内（FSCALE=13 安全）。 */
#define FSCALE 13
static int mp_fix_round(int a) {   /* 有符号四舍五入右移 */
    return (a + (a >= 0 ? (1 << (FSCALE - 1)) : -(1 << (FSCALE - 1)))) >> FSCALE;
}

/* ---- 抓 8x8 样本（越界 clamp 到边缘；样本为中心化整数 -128..127） ---- */
static void grab8(const uint8_t *plane, int pstride, int pw, int ph,
                  int x0, int y0, int s[8][8]) {
    for (int y = 0; y < 8; y++) {
        int py = y0 + y; if (py >= ph) py = ph - 1;
        const uint8_t *row = plane + (size_t)py * (size_t)pstride;
        for (int x = 0; x < 8; x++) {
            int px = x0 + x; if (px >= pw) px = pw - 1;
            s[y][x] = (int)row[px] - 128;
        }
    }
}

/* ---- 编码一个 8x8 块（定点 DCT→量化→zigzag→RLE/Huffman） ---- */
static int encode_block(mj_bw *w, const int samp[8][8], const uint8_t *q,
                        const mj_hc *dc, const mj_hc *ac, int *prev_dc) {
    static int ci[8][8];
    static int ci_ready = 0;
    if (!ci_ready) {
        for (int u = 0; u < 8; u++)
            for (int x = 0; x < 8; x++)
                ci[u][x] = (int)(MP_COSP[u][x] * (1 << FSCALE) + 0.5);
        ci_ready = 1;
    }
    int coef[64];
    int R[8][8];
    for (int y = 0; y < 8; y++)
        for (int u = 0; u < 8; u++) {
            int acc = 0;
            for (int x = 0; x < 8; x++) acc += ci[u][x] * samp[y][x];
            R[u][y] = mp_fix_round(acc);
        }
    for (int u = 0; u < 8; u++)
        for (int v = 0; v < 8; v++) {
            int acc = 0;
            for (int y = 0; y < 8; y++) acc += ci[v][y] * R[u][y];
            int c = mp_fix_round(acc);
            int Q = q[v * 8 + u];
            coef[v * 8 + u] = (c + (c >= 0 ? Q / 2 : -Q / 2)) / Q;
        }

    int dcv = coef[0];
    int diff = dcv - *prev_dc;
    *prev_dc = dcv;
    int ds = cat_of(diff);
    if (bw_bits(w, dc[ds].code, dc[ds].size)) return -1;
    if (ds && bw_bits(w, cat_bits(diff, ds), ds)) return -1;

    int run = 0;
    for (int p = 1; p < 64; p++) {
        int val = coef[MP_ZZ[p]];
        if (!val) { run++; continue; }
        while (run >= 16) {
            if (bw_bits(w, ac[0xF0].code, ac[0xF0].size)) return -1;
            run -= 16;
        }
        int sc = cat_of(val);
        int sym = (run << 4) | sc;
        if (bw_bits(w, ac[sym].code, ac[sym].size)) return -1;
        if (bw_bits(w, cat_bits(val, sc), sc)) return -1;
        run = 0;
    }
    if (run > 0) {
        if (bw_bits(w, ac[0x00].code, ac[0x00].size)) return -1;
    }
    return 0;
}

/* ---- DHT 段 ---- */
static int emit_dht(mj_out *o, uint8_t spec, const uint8_t *bits,
                    const uint8_t *vals, int nval) {
    if (mj_byte(o, 0xFF) || mj_byte(o, 0xC4)) return -1;
    if (mj_u16(o, (uint16_t)(2 + 1 + 16 + nval))) return -1;
    if (mj_byte(o, spec)) return -1;
    if (mj_bytes(o, bits, 16)) return -1;
    if (mj_bytes(o, vals, (size_t)nval)) return -1;
    return 0;
}

int mp_jpeg_encode_rgb565(const uint8_t *rgb565, int w, int h, int quality,
                          uint8_t **out, size_t *outlen) {
    if (!rgb565 || w <= 0 || h <= 0 || !out || !outlen) return -1;
    if (quality < 1) quality = 1;
    if (quality > 100) quality = 100;

    int wc = (w + 1) / 2, hc = (h + 1) / 2;
    uint8_t *Y = (uint8_t *)malloc((size_t)w * (size_t)h);
    uint8_t *Cb = (uint8_t *)malloc((size_t)wc * (size_t)hc);
    uint8_t *Cr = (uint8_t *)malloc((size_t)wc * (size_t)hc);
    uint8_t ql[64], qc[64];
    if (!Y || !Cb || !Cr) { free(Y); free(Cb); free(Cr); return -1; }
    build_q(MQ_LUMA, quality, ql);
    build_q(MQ_CHROMA, quality, qc);

    /* RGB565 → Y 平面(全分辨率) + Cb/Cr 平面(2x2 平均)，全部整数定点（BT.601 近似）。
     * Y=(77R+150G+29B)>>8；Cb=128+((-43R-85G+128B)>>8)；Cr=128+((128R-107G-21B)>>8)。 */
    for (int cy = 0; cy < h; cy += 2)
        for (int cx = 0; cx < w; cx += 2) {
            int cb = 0, cr = 0, n = 0;
            for (int dy = 0; dy < 2; dy++)
                for (int dx = 0; dx < 2; dx++) {
                    int py = cy + dy, px = cx + dx;
                    if (py >= h || px >= w) continue;
                    uint16_t v = (uint16_t)(rgb565[((size_t)py * (size_t)w + (size_t)px) * 2]
                                            | (rgb565[((size_t)py * (size_t)w + (size_t)px) * 2 + 1] << 8));
                    int r5 = (v >> 11) & 0x1f, g6 = (v >> 5) & 0x3f, b5 = v & 0x1f;
                    int R = (r5 << 3) | (r5 >> 2);
                    int G = (g6 << 2) | (g6 >> 4);
                    int B = (b5 << 3) | (b5 >> 2);
                    Y[(size_t)py * (size_t)w + (size_t)px] = (uint8_t)((77 * R + 150 * G + 29 * B) >> 8);
                    cb += (-43) * R - 85 * G + 128 * B;
                    cr += 128 * R - 107 * G - 21 * B;
                    n++;
                }
            size_t k = (size_t)(cy / 2) * (size_t)wc + (size_t)(cx / 2);
            int dn = n ? n : 1;
            Cb[k] = (uint8_t)clamp255(128 + ((cb / dn) >> 8));
            Cr[k] = (uint8_t)clamp255(128 + ((cr / dn) >> 8));
        }

    mj_out o = {0};
    mj_hc hdc_l[16], hac_l[256], hdc_c[16], hac_c[256];
    memset(hdc_l, 0, sizeof hdc_l); memset(hac_l, 0, sizeof hac_l);
    memset(hdc_c, 0, sizeof hdc_c); memset(hac_c, 0, sizeof hac_c);
    huff_build(HB_DC_L, HV_DC_L, 12, hdc_l);
    huff_build(HB_AC_L, HV_AC_L, 162, hac_l);
    huff_build(HB_DC_C, HV_DC_C, 12, hdc_c);
    huff_build(HB_AC_C, HV_AC_C, 162, hac_c);

    int rc = -1;
    /* SOI */
    if (mj_byte(&o, 0xFF) || mj_byte(&o, 0xD8)) goto done;
    /* APP0 JIF */
    static const uint8_t jfif[14] = { 'J','F','I','F',0,1,1,0,0,1,0,1,0,0 };
    if (mj_byte(&o, 0xFF) || mj_byte(&o, 0xE0) || mj_u16(&o, 0x0010)) goto done;
    if (mj_bytes(&o, jfif, 14)) goto done;
    /* DQT（2 张表） */
    if (mj_byte(&o, 0xFF) || mj_byte(&o, 0xDB) || mj_u16(&o, 0x0084)) goto done;
    if (mj_byte(&o, 0x00) || mj_bytes(&o, ql, 64)) goto done;
    if (mj_byte(&o, 0x01) || mj_bytes(&o, qc, 64)) goto done;
    /* SOF0 baseline, 8-bit, 3 分量, 4:2:0 */
    if (mj_byte(&o, 0xFF) || mj_byte(&o, 0xC0) || mj_u16(&o, 0x0011)) goto done;
    if (mj_byte(&o, 0x08)) goto done;
    if (mj_u16(&o, (uint16_t)h) || mj_u16(&o, (uint16_t)w)) goto done;
    if (mj_byte(&o, 0x03)) goto done;
    if (mj_byte(&o, 0x01) || mj_byte(&o, 0x22) || mj_byte(&o, 0x00)) goto done;  /* Y */
    if (mj_byte(&o, 0x02) || mj_byte(&o, 0x11) || mj_byte(&o, 0x01)) goto done;  /* Cb */
    if (mj_byte(&o, 0x03) || mj_byte(&o, 0x11) || mj_byte(&o, 0x01)) goto done;  /* Cr */
    /* DHT（4 张表） */
    if (emit_dht(&o, 0x00, HB_DC_L, HV_DC_L, 12)) goto done;
    if (emit_dht(&o, 0x10, HB_AC_L, HV_AC_L, 162)) goto done;
    if (emit_dht(&o, 0x01, HB_DC_C, HV_DC_C, 12)) goto done;
    if (emit_dht(&o, 0x11, HB_AC_C, HV_AC_C, 162)) goto done;
    /* SOS */
    if (mj_byte(&o, 0xFF) || mj_byte(&o, 0xDA) || mj_u16(&o, 0x000C)) goto done;
    if (mj_byte(&o, 0x03)) goto done;
    if (mj_byte(&o, 0x01) || mj_byte(&o, 0x00)) goto done;  /* Y: DC0 AC0 */
    if (mj_byte(&o, 0x02) || mj_byte(&o, 0x11)) goto done;  /* Cb: DC1 AC1 */
    if (mj_byte(&o, 0x03) || mj_byte(&o, 0x11)) goto done;  /* Cr: DC1 AC1 */
    if (mj_byte(&o, 0x00) || mj_byte(&o, 0x3F) || mj_byte(&o, 0x00)) goto done;

    {
        mj_bw bb = { &o, 0, 0 };
        int pd_y = 0, pd_cb = 0, pd_cr = 0;   /* 每分量独立 DC 预测器（JPEG 规范） */
        int mx_count = (w + 15) / 16, my_count = (h + 15) / 16;
        for (int my = 0; my < my_count; my++)
            for (int mx = 0; mx < mx_count; mx++) {
                for (int by = 0; by < 2; by++)
                    for (int bx = 0; bx < 2; bx++) {
                        int s[8][8];
                        grab8(Y, w, w, h, mx * 16 + bx * 8, my * 16 + by * 8, s);
                        if (encode_block(&bb, s, ql, hdc_l, hac_l, &pd_y)) goto done;
                    }
                {
                    int s[8][8];
                    grab8(Cb, wc, wc, hc, mx * 8, my * 8, s);
                    if (encode_block(&bb, s, qc, hdc_c, hac_c, &pd_cb)) goto done;
                }
                {
                    int s[8][8];
                    grab8(Cr, wc, wc, hc, mx * 8, my * 8, s);
                    if (encode_block(&bb, s, qc, hdc_c, hac_c, &pd_cr)) goto done;
                }
            }
        if (bw_flush(&bb)) goto done;
    }
    if (mj_byte(&o, 0xFF) || mj_byte(&o, 0xD9)) goto done;  /* EOI */
    rc = 0;

done:
    free(Y); free(Cb); free(Cr);
    if (rc == 0) { *out = o.d; *outlen = o.len; }
    else { free(o.d); *out = NULL; *outlen = 0; }
    return rc;
}