// Fast exact PNG decoding with Wuffs for the common Factorio sprite formats.
// Produces what AppKit's NSBitmapImageRep gives on macOS: tightly packed, non-premultiplied
// 8-bit samples with the source channel count (gray 1, gray+alpha 2, RGB 3, RGBA 4).
// Returns 0 for anything else (palette, tRNS, 16-bit, malformed) so the caller can fall back.
#define WUFFS_IMPLEMENTATION
#define WUFFS_CONFIG__MODULES
#define WUFFS_CONFIG__MODULE__BASE
#define WUFFS_CONFIG__MODULE__ADLER32
#define WUFFS_CONFIG__MODULE__CRC32
#define WUFFS_CONFIG__MODULE__DEFLATE
#define WUFFS_CONFIG__MODULE__ZLIB
#define WUFFS_CONFIG__MODULE__PNG
#include "../../../third_party/wuffs/wuffs-v0.3.c"

#include <stdlib.h>
#include <string.h>
#include "pngdecode.h"

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

// Channel count for supported PNGs (from IHDR, no tRNS before IDAT), else 0.
static int fx_png_channels(const uint8_t *d, size_t n) {
    static const uint8_t sig[8] = {0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'};
    if (n < 33 || memcmp(d, sig, 8) || memcmp(d + 12, "IHDR", 4)) return 0;
    uint8_t depth = d[24], type = d[25];
    int channels = type == 0 ? 1 : type == 2 ? 3 : type == 4 ? 2 : type == 6 ? 4 : 0;
    if (depth != 8 || !channels) return 0;
    for (size_t at = 8; at + 12 <= n;) {
        uint32_t len = be32(d + at);
        const uint8_t *kind = d + at + 4;
        if (!memcmp(kind, "IDAT", 4)) break;
        if (!memcmp(kind, "tRNS", 4)) return 0;
        if ((uint64_t)at + 12 + len > n) return 0;
        at += 12 + (size_t)len;
    }
    return channels;
}

int fx_png_decode(const uint8_t *data, size_t size, uint8_t **out, int *width, int *height, int *channels) {
    int ch = fx_png_channels(data, size);
    if (!ch) return 0;
    wuffs_png__decoder *dec = wuffs_png__decoder__alloc();
    if (!dec) return 0;
    // Skipping zlib/CRC verification does not change decoded pixels (verified byte-identical with
    // tests/native/png_parity.m); a corrupt file decodes to garbage instead of failing.
    wuffs_png__decoder__set_quirk_enabled(dec, WUFFS_BASE__QUIRK_IGNORE_CHECKSUM, true);
    int ok = 0;
    uint8_t *rgba = NULL, *work = NULL;
    wuffs_base__io_buffer src = wuffs_base__ptr_u8__reader((uint8_t *)data, size, true);
    wuffs_base__image_config ic = {0};
    if (wuffs_png__decoder__decode_image_config(dec, &ic, &src).repr) goto done;
    uint32_t w = wuffs_base__pixel_config__width(&ic.pixcfg), h = wuffs_base__pixel_config__height(&ic.pixcfg);
    if (!w || !h || (uint64_t)w * h > (1u << 28)) goto done;
    wuffs_base__pixel_config__set(&ic.pixcfg, WUFFS_BASE__PIXEL_FORMAT__RGBA_NONPREMUL,
                                  WUFFS_BASE__PIXEL_SUBSAMPLING__NONE, w, h);
    size_t pixels = (size_t)w * h;
    rgba = malloc(pixels * 4);
    uint64_t work_len = wuffs_png__decoder__workbuf_len(dec).max_incl;
    work = work_len ? malloc(work_len) : NULL;
    if (!rgba || (work_len && !work)) goto done;
    wuffs_base__pixel_buffer pb;
    if (wuffs_base__pixel_buffer__set_from_slice(&pb, &ic.pixcfg, wuffs_base__make_slice_u8(rgba, pixels * 4)).repr) goto done;
    if (wuffs_png__decoder__decode_frame(dec, &pb, &src, WUFFS_BASE__PIXEL_BLEND__SRC,
                                         wuffs_base__make_slice_u8(work, work_len), NULL).repr) goto done;
    // Pack to the source channel count (gray sources decode with R == G == B).
    if (ch != 4) {
        uint8_t *dst = rgba;
        for (size_t i = 0; i < pixels; i++) {
            const uint8_t *p = rgba + i * 4;
            if (ch == 1) { dst[i] = p[0]; }
            else if (ch == 2) { dst[2 * i] = p[0]; dst[2 * i + 1] = p[3]; }
            else { uint8_t r = p[0], g = p[1], b = p[2]; dst[3 * i] = r; dst[3 * i + 1] = g; dst[3 * i + 2] = b; }
        }
    }
    *out = rgba; rgba = NULL;
    *width = (int)w; *height = (int)h; *channels = ch;
    ok = 1;
done:
    free(rgba);
    free(work);
    free(dec);
    return ok;
}