// Part of JoeyLib, Copyright (c) 2026 Scott Duensing, under the MIT License.
// See LICENSE.md for license information.
// SPDX-License-Identifier: MIT
//
// toolPng.c - PNG for the JoeyLib host tools (toolPng.h).
//
// Decode: chunks (CRCs checked), the IDATs inflated as one zlib stream, then
// each Adam7 pass (one pass when not interlaced) unfiltered a row at a time
// and its pixels unpacked into the image. Encode: one IDAT, each row's
// filter the one whose bytes, as signed, sum smallest (the usual
// heuristic).

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "toolFile.h"
#include "toolPng.h"
#include "toolZlib.h"

#define PNG_SIGNATURE_BYTES 8u
#define CHUNK_OVERHEAD      12u
#define IHDR_BYTES          13u
#define COLOR_GRAY          0u
#define COLOR_RGB           2u
#define COLOR_PALETTE       3u
#define COLOR_GRAY_ALPHA    4u
#define COLOR_RGBA          6u
#define FILTER_NONE         0u
#define FILTER_SUB          1u
#define FILTER_UP           2u
#define FILTER_AVERAGE      3u
#define FILTER_PAETH        4u
#define FILTER_COUNT        5u
#define ADAM7_PASSES        7u
#define OPAQUE              255u
#define MAX_DIMENSION       0x8000u

typedef struct DecodeT {
    uint32_t       width;
    uint32_t       height;
    uint8_t        depth;
    uint8_t        color;
    uint8_t        channels;
    bool           hasKey;
    uint16_t       key[3];
    ToolImageT    *img;
} DecodeT;

static const uint8_t kSignature[PNG_SIGNATURE_BYTES] = { 0x89, 'P', 'N', 'G', '\r', '\n', 0x1A, '\n' };
static const uint8_t kAdam7X0[ADAM7_PASSES]          = { 0, 4, 0, 2, 0, 1, 0 };
static const uint8_t kAdam7Y0[ADAM7_PASSES]          = { 0, 0, 4, 0, 2, 0, 1 };
static const uint8_t kAdam7Dx[ADAM7_PASSES]          = { 8, 8, 4, 4, 2, 2, 1 };
static const uint8_t kAdam7Dy[ADAM7_PASSES]          = { 8, 8, 8, 4, 4, 2, 2 };

static bool     appendBytes(uint8_t **buf, size_t *len, size_t *cap, const uint8_t *data, size_t count);
static bool     appendChunk(uint8_t **buf, size_t *len, size_t *cap, const char *type, const uint8_t *data, size_t count);
static void     decodePixel(const DecodeT *d, const uint8_t *row, uint32_t i, uint32_t x, uint32_t y);
static uint8_t  paeth(uint8_t a, uint8_t b, uint8_t c);
static uint32_t readBe32(const uint8_t *p);
static uint16_t sample(const uint8_t *row, uint32_t index, uint8_t depth);
static bool     unfilterRow(uint8_t *row, const uint8_t *prior, size_t rowBytes, size_t bpp, uint8_t filter);
static void     writeBe32(uint8_t *p, uint32_t v);


static bool appendBytes(uint8_t **buf, size_t *len, size_t *cap, const uint8_t *data, size_t count) {
    uint8_t *grown;
    size_t   want;

    if (*len + count > *cap) {
        want = *cap == 0u ? 4096u : *cap;
        while (want < *len + count) {
            want *= 2u;
        }
        grown = (uint8_t *)realloc(*buf, want);
        if (grown == NULL) {
            return false;
        }
        *buf = grown;
        *cap = want;
    }
    memcpy(*buf + *len, data, count);
    *len += count;
    return true;
}


static bool appendChunk(uint8_t **buf, size_t *len, size_t *cap, const char *type, const uint8_t *data, size_t count) {
    uint8_t  head[8];
    uint8_t  tail[4];
    uint32_t crc;

    writeBe32(head, (uint32_t)count);
    memcpy(head + 4, type, 4);
    crc = toolCrc32(0u, head + 4, 4u);
    crc = toolCrc32(crc, data, count);
    writeBe32(tail, crc);
    return appendBytes(buf, len, cap, head, sizeof(head)) && appendBytes(buf, len, cap, data, count) && appendBytes(buf, len, cap, tail, sizeof(tail));
}


// Pixel i of an unfiltered row, stored at (x, y) of the image.
static void decodePixel(const DecodeT *d, const uint8_t *row, uint32_t i, uint32_t x, uint32_t y) {
    uint8_t *out;
    uint16_t v[4];
    uint16_t c;
    uint16_t maxValue;

    if (d->color == COLOR_PALETTE) {
        d->img->pixels[(size_t)y * d->width + x] = (uint8_t)sample(row, i, d->depth);
        return;
    }
    for (c = 0; c < d->channels; c++) {
        v[c] = sample(row, i * d->channels + c, d->depth);
    }
    out      = d->img->pixels + ((size_t)y * d->width + x) * 4u;
    maxValue = (uint16_t)((1u << d->depth) - 1u);
    if (d->color == COLOR_GRAY || d->color == COLOR_GRAY_ALPHA) {
        out[0] = (uint8_t)(d->depth == 16u ? v[0] >> 8 : (uint32_t)v[0] * 255u / maxValue);
        out[1] = out[0];
        out[2] = out[0];
        if (d->color == COLOR_GRAY_ALPHA) {
            out[3] = (uint8_t)(d->depth == 16u ? v[1] >> 8 : v[1]);
        } else {
            out[3] = (d->hasKey && v[0] == d->key[0]) ? 0u : OPAQUE;
        }
        return;
    }
    for (c = 0; c < 3u; c++) {
        out[c] = (uint8_t)(d->depth == 16u ? v[c] >> 8 : v[c]);
    }
    if (d->color == COLOR_RGBA) {
        out[3] = (uint8_t)(d->depth == 16u ? v[3] >> 8 : v[3]);
    } else {
        out[3] = (d->hasKey && v[0] == d->key[0] && v[1] == d->key[1] && v[2] == d->key[2]) ? 0u : OPAQUE;
    }
}


static uint8_t paeth(uint8_t a, uint8_t b, uint8_t c) {
    int32_t p;
    int32_t pa;
    int32_t pb;
    int32_t pc;

    p  = (int32_t)a + (int32_t)b - (int32_t)c;
    pa = abs(p - (int32_t)a);
    pb = abs(p - (int32_t)b);
    pc = abs(p - (int32_t)c);
    if (pa <= pb && pa <= pc) {
        return a;
    }
    return pb <= pc ? b : c;
}


static uint32_t readBe32(const uint8_t *p) {
    return (uint32_t)p[0] << 24 | (uint32_t)p[1] << 16 | (uint32_t)p[2] << 8 | (uint32_t)p[3];
}


// Sample `index` of a row of `depth`-bit samples, packed high bits first.
static uint16_t sample(const uint8_t *row, uint32_t index, uint8_t depth) {
    uint32_t bit;

    if (depth == 16u) {
        return (uint16_t)((uint16_t)row[index * 2u] << 8 | row[index * 2u + 1u]);
    }
    if (depth == 8u) {
        return row[index];
    }
    bit = index * depth;
    return (uint16_t)((row[bit >> 3] >> (8u - depth - (bit & 7u))) & ((1u << depth) - 1u));
}


bool toolImageCreate(ToolImageT *img, uint32_t width, uint32_t height, bool indexed) {
    memset(img, 0, sizeof(*img));
    img->width   = width;
    img->height  = height;
    img->indexed = indexed;
    img->pixels  = (uint8_t *)calloc((size_t)width * height, indexed ? 1u : 4u);
    return img->pixels != NULL || (size_t)width * height == 0u;
}


void toolImageFree(ToolImageT *img) {
    free(img->pixels);
    img->pixels = NULL;
}


bool toolImageToRgba(const ToolImageT *src, ToolImageT *dst) {
    size_t n;
    size_t i;

    if (!toolImageCreate(dst, src->width, src->height, false)) {
        return false;
    }
    n = (size_t)src->width * src->height;
    if (!src->indexed) {
        memcpy(dst->pixels, src->pixels, n * 4u);
        return true;
    }
    for (i = 0; i < n; i++) {
        memcpy(dst->pixels + i * 4u, src->palette[src->pixels[i]], 4u);
    }
    return true;
}


bool toolPngDecode(const uint8_t *data, size_t len, ToolImageT *img, char err[TOOL_PNG_ERROR_BYTES]) {
    DecodeT  d;
    uint8_t *idat;
    uint8_t *raw;
    uint8_t *prior;
    uint8_t *row;
    size_t   idatLen;
    size_t   idatCap;
    size_t   rawLen;
    size_t   pos;
    size_t   bpp;
    size_t   rowBytes;
    size_t   at;
    uint32_t chunkLen;
    uint32_t passW;
    uint32_t passH;
    uint32_t pass;
    uint32_t passes;
    uint32_t px;
    uint32_t py;
    uint32_t i;
    uint8_t  palette[TOOL_PNG_PALETTE_MAX][4];
    uint16_t paletteCount;
    uint8_t  interlace;
    bool     seenEnd;

    memset(&d, 0, sizeof(d));
    memset(palette, OPAQUE, sizeof(palette));
    paletteCount = 0u;
    memset(img, 0, sizeof(*img));
    idat      = NULL;
    idatLen   = 0;
    idatCap   = 0;
    seenEnd   = false;
    interlace = 0u;
    if (len < PNG_SIGNATURE_BYTES || memcmp(data, kSignature, PNG_SIGNATURE_BYTES) != 0) {
        snprintf(err, TOOL_PNG_ERROR_BYTES, "not a PNG");
        return false;
    }
    pos = PNG_SIGNATURE_BYTES;
    while (!seenEnd) {
        if (pos + CHUNK_OVERHEAD > len) {
            snprintf(err, TOOL_PNG_ERROR_BYTES, "truncated PNG");
            free(idat);
            return false;
        }
        chunkLen = readBe32(data + pos);
        if (chunkLen > len - pos - CHUNK_OVERHEAD) {
            snprintf(err, TOOL_PNG_ERROR_BYTES, "truncated PNG chunk");
            free(idat);
            return false;
        }
        if (toolCrc32(0u, data + pos + 4u, chunkLen + 4u) != readBe32(data + pos + 8u + chunkLen)) {
            snprintf(err, TOOL_PNG_ERROR_BYTES, "bad CRC in a %.4s chunk", (const char *)(data + pos + 4u));
            free(idat);
            return false;
        }
        if (memcmp(data + pos + 4u, "IHDR", 4) == 0 && chunkLen == IHDR_BYTES) {
            at          = pos + 8u;
            d.width     = readBe32(data + at);
            d.height    = readBe32(data + at + 4u);
            d.depth     = data[at + 8u];
            d.color     = data[at + 9u];
            interlace   = data[at + 12u];
        } else if (memcmp(data + pos + 4u, "PLTE", 4) == 0) {
            paletteCount = (uint16_t)(chunkLen / 3u > TOOL_PNG_PALETTE_MAX ? TOOL_PNG_PALETTE_MAX : chunkLen / 3u);
            for (i = 0; i < paletteCount; i++) {
                palette[i][0] = data[pos + 8u + i * 3u];
                palette[i][1] = data[pos + 9u + i * 3u];
                palette[i][2] = data[pos + 10u + i * 3u];
            }
        } else if (memcmp(data + pos + 4u, "tRNS", 4) == 0) {
            if (d.color == COLOR_PALETTE) {
                for (i = 0; i < chunkLen && i < TOOL_PNG_PALETTE_MAX; i++) {
                    palette[i][3] = data[pos + 8u + i];
                }
            } else if (d.color == COLOR_GRAY && chunkLen >= 2u) {
                d.hasKey = true;
                d.key[0] = (uint16_t)((uint16_t)data[pos + 8u] << 8 | data[pos + 9u]);
            } else if (d.color == COLOR_RGB && chunkLen >= 6u) {
                d.hasKey = true;
                for (i = 0; i < 3u; i++) {
                    d.key[i] = (uint16_t)((uint16_t)data[pos + 8u + i * 2u] << 8 | data[pos + 9u + i * 2u]);
                }
            }
        } else if (memcmp(data + pos + 4u, "IDAT", 4) == 0) {
            if (!appendBytes(&idat, &idatLen, &idatCap, data + pos + 8u, chunkLen)) {
                snprintf(err, TOOL_PNG_ERROR_BYTES, "out of memory");
                free(idat);
                return false;
            }
        } else if (memcmp(data + pos + 4u, "IEND", 4) == 0) {
            seenEnd = true;
        }
        pos += CHUNK_OVERHEAD + chunkLen;
    }
    switch (d.color) {
        case COLOR_GRAY:
            d.channels = 1u;
            break;
        case COLOR_RGB:
            d.channels = 3u;
            break;
        case COLOR_PALETTE:
            d.channels = 1u;
            break;
        case COLOR_GRAY_ALPHA:
            d.channels = 2u;
            break;
        case COLOR_RGBA:
            d.channels = 4u;
            break;
        default:
            d.channels = 0u;
            break;
    }
    if (d.channels == 0u || d.width == 0u || d.height == 0u || d.width > MAX_DIMENSION || d.height > MAX_DIMENSION || (d.depth != 1u && d.depth != 2u && d.depth != 4u && d.depth != 8u && d.depth != 16u) || interlace > 1u) {
        snprintf(err, TOOL_PNG_ERROR_BYTES, "unsupported PNG header (%ux%u, depth %u, color type %u)", d.width, d.height, d.depth, d.color);
        free(idat);
        return false;
    }
    if (!toolInflate(idat, idatLen, TOOL_ZLIB_ZLIB, &raw, &rawLen)) {
        snprintf(err, TOOL_PNG_ERROR_BYTES, "bad image data");
        free(idat);
        return false;
    }
    free(idat);
    if (!toolImageCreate(img, d.width, d.height, d.color == COLOR_PALETTE)) {
        snprintf(err, TOOL_PNG_ERROR_BYTES, "out of memory");
        free(raw);
        return false;
    }
    img->paletteCount = paletteCount;
    memcpy(img->palette, palette, sizeof(palette));
    if (img->indexed && img->paletteCount == 0u) {
        snprintf(err, TOOL_PNG_ERROR_BYTES, "palette PNG without a PLTE");
        free(raw);
        toolImageFree(img);
        return false;
    }
    d.img  = img;
    bpp    = (size_t)d.channels * d.depth / 8u;
    if (bpp == 0u) {
        bpp = 1u;
    }
    passes = interlace != 0u ? ADAM7_PASSES : 1u;
    pos    = 0;
    for (pass = 0; pass < passes; pass++) {
        if (interlace != 0u) {
            passW = d.width > kAdam7X0[pass] ? (d.width - kAdam7X0[pass] + kAdam7Dx[pass] - 1u) / kAdam7Dx[pass] : 0u;
            passH = d.height > kAdam7Y0[pass] ? (d.height - kAdam7Y0[pass] + kAdam7Dy[pass] - 1u) / kAdam7Dy[pass] : 0u;
        } else {
            passW = d.width;
            passH = d.height;
        }
        if (passW == 0u || passH == 0u) {
            continue;
        }
        rowBytes = ((size_t)passW * d.channels * d.depth + 7u) / 8u;
        prior    = NULL;
        for (py = 0; py < passH; py++) {
            if (pos + 1u + rowBytes > rawLen) {
                snprintf(err, TOOL_PNG_ERROR_BYTES, "image data too short");
                free(raw);
                toolImageFree(img);
                return false;
            }
            row = raw + pos + 1u;
            if (!unfilterRow(row, prior, rowBytes, bpp, raw[pos])) {
                snprintf(err, TOOL_PNG_ERROR_BYTES, "bad row filter %u", raw[pos]);
                free(raw);
                toolImageFree(img);
                return false;
            }
            for (px = 0; px < passW; px++) {
                if (interlace != 0u) {
                    decodePixel(&d, row, px, kAdam7X0[pass] + px * kAdam7Dx[pass], kAdam7Y0[pass] + py * kAdam7Dy[pass]);
                } else {
                    decodePixel(&d, row, px, px, py);
                }
            }
            prior = row;
            pos  += 1u + rowBytes;
        }
    }
    free(raw);
    return true;
}


bool toolPngEncode(const ToolImageT *img, uint8_t **out, size_t *outLen) {
    uint8_t       *raw;
    uint8_t       *z;
    uint8_t       *buf;
    uint8_t       *trial;
    uint8_t        ihdr[IHDR_BYTES];
    uint8_t        plte[TOOL_PNG_PALETTE_MAX * 3u];
    uint8_t        trns[TOOL_PNG_PALETTE_MAX];
    size_t         len;
    size_t         cap;
    size_t         zLen;
    size_t         bpp;
    size_t         rowBytes;
    size_t         i;
    uint32_t       y;
    uint32_t       cost;
    uint32_t       best;
    uint16_t       trnsCount;
    uint8_t        f;
    uint8_t        bestFilter;
    uint8_t        left;
    uint8_t        up;
    uint8_t        upLeft;
    uint8_t        v;
    bool           ok;
    const uint8_t *row;
    const uint8_t *prior;

    bpp      = img->indexed ? 1u : 4u;
    rowBytes = (size_t)img->width * bpp;
    raw      = (uint8_t *)malloc((rowBytes + 1u) * img->height);
    trial    = (uint8_t *)malloc(rowBytes * FILTER_COUNT);
    if (raw == NULL || trial == NULL) {
        free(raw);
        free(trial);
        return false;
    }
    for (y = 0; y < img->height; y++) {
        row        = img->pixels + (size_t)y * rowBytes;
        prior      = y > 0u ? row - rowBytes : NULL;
        best       = 0xFFFFFFFFu;
        bestFilter = FILTER_NONE;
        for (f = 0; f < FILTER_COUNT; f++) {
            cost = 0u;
            for (i = 0; i < rowBytes; i++) {
                left   = i >= bpp ? row[i - bpp] : 0u;
                up     = prior != NULL ? prior[i] : 0u;
                upLeft = (i >= bpp && prior != NULL) ? prior[i - bpp] : 0u;
                switch (f) {
                    case FILTER_SUB:
                        v = (uint8_t)(row[i] - left);
                        break;
                    case FILTER_UP:
                        v = (uint8_t)(row[i] - up);
                        break;
                    case FILTER_AVERAGE:
                        v = (uint8_t)(row[i] - ((uint32_t)left + up) / 2u);
                        break;
                    case FILTER_PAETH:
                        v = (uint8_t)(row[i] - paeth(left, up, upLeft));
                        break;
                    default:
                        v = row[i];
                        break;
                }
                trial[f * rowBytes + i] = v;
                cost += (uint32_t)(v < 128u ? v : 256u - v);
            }
            if (cost < best) {
                best       = cost;
                bestFilter = f;
            }
        }
        raw[(size_t)y * (rowBytes + 1u)] = bestFilter;
        memcpy(raw + (size_t)y * (rowBytes + 1u) + 1u, trial + bestFilter * rowBytes, rowBytes);
    }
    free(trial);
    ok = toolDeflateZlib(raw, (rowBytes + 1u) * img->height, &z, &zLen);
    free(raw);
    if (!ok) {
        return false;
    }
    writeBe32(ihdr, img->width);
    writeBe32(ihdr + 4, img->height);
    ihdr[8]  = 8u;
    ihdr[9]  = img->indexed ? COLOR_PALETTE : COLOR_RGBA;
    ihdr[10] = 0u;
    ihdr[11] = 0u;
    ihdr[12] = 0u;
    buf      = NULL;
    len      = 0;
    cap      = 0;
    ok       = appendBytes(&buf, &len, &cap, kSignature, PNG_SIGNATURE_BYTES) && appendChunk(&buf, &len, &cap, "IHDR", ihdr, IHDR_BYTES);
    if (ok && img->indexed) {
        trnsCount = 0u;
        for (i = 0; i < img->paletteCount; i++) {
            plte[i * 3u]      = img->palette[i][0];
            plte[i * 3u + 1u] = img->palette[i][1];
            plte[i * 3u + 2u] = img->palette[i][2];
            trns[i]           = img->palette[i][3];
            if (trns[i] != OPAQUE) {
                trnsCount = (uint16_t)(i + 1u);
            }
        }
        ok = appendChunk(&buf, &len, &cap, "PLTE", plte, (size_t)img->paletteCount * 3u);
        if (ok && trnsCount > 0u) {
            ok = appendChunk(&buf, &len, &cap, "tRNS", trns, trnsCount);
        }
    }
    ok = ok && appendChunk(&buf, &len, &cap, "IDAT", z, zLen) && appendChunk(&buf, &len, &cap, "IEND", NULL, 0u);
    free(z);
    if (!ok) {
        free(buf);
        return false;
    }
    *out    = buf;
    *outLen = len;
    return true;
}


bool toolPngLoad(const char *path, ToolImageT *img, char err[TOOL_PNG_ERROR_BYTES]) {
    uint8_t *data;
    size_t   len;
    bool     ok;

    if (!toolReadFile(path, &data, &len)) {
        snprintf(err, TOOL_PNG_ERROR_BYTES, "cannot read it");
        return false;
    }
    ok = toolPngDecode(data, len, img, err);
    free(data);
    return ok;
}


bool toolPngSave(const char *path, const ToolImageT *img) {
    uint8_t *data;
    size_t   len;
    bool     ok;

    if (!toolPngEncode(img, &data, &len)) {
        return false;
    }
    ok = toolWriteFile(path, data, len);
    free(data);
    return ok;
}


static bool unfilterRow(uint8_t *row, const uint8_t *prior, size_t rowBytes, size_t bpp, uint8_t filter) {
    size_t  i;
    uint8_t left;
    uint8_t up;
    uint8_t upLeft;

    for (i = 0; i < rowBytes; i++) {
        left   = i >= bpp ? row[i - bpp] : 0u;
        up     = prior != NULL ? prior[i] : 0u;
        upLeft = (i >= bpp && prior != NULL) ? prior[i - bpp] : 0u;
        switch (filter) {
            case FILTER_NONE:
                break;
            case FILTER_SUB:
                row[i] = (uint8_t)(row[i] + left);
                break;
            case FILTER_UP:
                row[i] = (uint8_t)(row[i] + up);
                break;
            case FILTER_AVERAGE:
                row[i] = (uint8_t)(row[i] + ((uint32_t)left + up) / 2u);
                break;
            case FILTER_PAETH:
                row[i] = (uint8_t)(row[i] + paeth(left, up, upLeft));
                break;
            default:
                return false;
        }
    }
    return true;
}


static void writeBe32(uint8_t *p, uint32_t v) {
    p[0] = (uint8_t)(v >> 24);
    p[1] = (uint8_t)(v >> 16);
    p[2] = (uint8_t)(v >> 8);
    p[3] = (uint8_t)v;
}