/* Lector y escritor de PNG (usa zlib solo para DEFLATE/CRC). */
#include <zlib.h>
#include "io_internal.h"

static uint32_t get_u32(const uint8_t *p) { return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) | ((uint32_t)p[2] << 8) | p[3]; }
void png_put_u32(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; }

static const uint8_t SIG[8] = {0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A};

void png_put_signature(tgb_sb *sb) { sb_append(sb, SIG, 8); }

void png_put_chunk(tgb_sb *sb, const char type[4], const uint8_t *data, size_t len) {
    uint8_t hdr[8];
    png_put_u32(hdr, (uint32_t)len);
    memcpy(hdr + 4, type, 4);
    sb_append(sb, hdr, 8);
    if (len) sb_append(sb, data, len);
    uLong crc = crc32(0, (const Bytef *)type, 4);
    if (len) crc = crc32(crc, data, (uInt)len);
    uint8_t c[4];
    png_put_u32(c, (uint32_t)crc);
    sb_append(sb, c, 4);
}

/* ================= decodificación ================= */
static int paeth(int a, int b, int c) {
    int p = a + b - c, pa = abs(p - a), pb = abs(p - b), pc = abs(p - c);
    return (pa <= pb && pa <= pc) ? a : (pb <= pc ? b : c);
}

static tgb_status bad(const char *msg) { return tgb_fail(TGB_ERR_CORRUPT, "PNG dañado: %s", msg); }

tgb_status tgb_png_decode_memory(const uint8_t *data, size_t len, tgb_canvas **out) {
    if (!data || !out) return tgb_fail(TGB_ERR_INVALID_ARG, "tgb_png_decode_memory: argumento nulo");
    *out = NULL;
    if (len < 8 || memcmp(data, SIG, 8) != 0) return tgb_fail(TGB_ERR_FORMAT, "No es una imagen PNG (firma incorrecta)");
    size_t pos = 8;
    uint32_t w = 0, h = 0;
    int depth = 0, ctype = 0;
    bool have_ihdr = false, have_iend = false;
    uint8_t palette[256 * 3];
    uint8_t pal_alpha[256];
    size_t pal_n = 0;
    bool have_trns = false;
    uint16_t key[3] = {0, 0, 0};
    uint8_t *idat = NULL;
    size_t idat_len = 0, idat_cap = 0;
    memset(pal_alpha, 255, sizeof pal_alpha);
    tgb_status s = TGB_OK;

    while (pos < len && !have_iend) {
        if (len - pos < 12) { s = bad("chunk incompleto al final del archivo"); goto done; }
        uint32_t clen = get_u32(data + pos);
        const uint8_t *type = data + pos + 4;
        if (clen > 0x7FFFFFFFu || (size_t)clen > len - pos - 12) { s = bad("longitud de chunk mayor que el archivo (¿truncado?)"); goto done; }
        const uint8_t *cdata = data + pos + 8;
        uLong crc = crc32(crc32(0, type, 4), cdata, clen);
        if ((uint32_t)crc != get_u32(cdata + clen)) { s = bad("CRC incorrecto"); goto done; }
        if (!have_ihdr && memcmp(type, "IHDR", 4) != 0) { s = bad("el primer chunk no es IHDR"); goto done; }
        if (!memcmp(type, "IHDR", 4)) {
            if (have_ihdr || clen != 13) { s = bad("IHDR no válido"); goto done; }
            w = get_u32(cdata); h = get_u32(cdata + 4);
            depth = cdata[8]; ctype = cdata[9];
            if (cdata[10] != 0 || cdata[11] != 0) { s = tgb_fail(TGB_ERR_UNSUPPORTED, "PNG con método de compresión o filtro desconocido"); goto done; }
            if (cdata[12] != 0) { s = tgb_fail(TGB_ERR_UNSUPPORTED, "PNG entrelazado (Adam7) no soportado todavía"); goto done; }
            if (w == 0 || h == 0 || w > TGB_MAX_DIMENSION || h > TGB_MAX_DIMENSION) {
                s = tgb_fail(TGB_ERR_LIMIT, "Imagen PNG de %ux%u fuera de los límites (máximo %d)", (unsigned)w, (unsigned)h, TGB_MAX_DIMENSION);
                goto done;
            }
            bool ok = (ctype == 0 && (depth == 1 || depth == 2 || depth == 4 || depth == 8 || depth == 16)) ||
                      (ctype == 2 && (depth == 8 || depth == 16)) || (ctype == 3 && (depth == 1 || depth == 2 || depth == 4 || depth == 8)) ||
                      (ctype == 4 && (depth == 8 || depth == 16)) || (ctype == 6 && (depth == 8 || depth == 16));
            if (!ok) { s = bad("combinación de tipo de color y profundidad no válida"); goto done; }
            have_ihdr = true;
        } else if (!memcmp(type, "PLTE", 4)) {
            if (clen % 3 != 0 || clen == 0 || clen > 768) { s = bad("paleta no válida"); goto done; }
            pal_n = clen / 3;
            memcpy(palette, cdata, clen);
        } else if (!memcmp(type, "tRNS", 4)) {
            have_trns = true;
            if (ctype == 3) { if (clen > 256) { s = bad("tRNS demasiado largo"); goto done; } memcpy(pal_alpha, cdata, clen); }
            else if (ctype == 0 && clen == 2) key[0] = (uint16_t)((cdata[0] << 8) | cdata[1]);
            else if (ctype == 2 && clen == 6) for (int i = 0; i < 3; i++) key[i] = (uint16_t)((cdata[i * 2] << 8) | cdata[i * 2 + 1]);
            else if (ctype == 0 || ctype == 2) { s = bad("tRNS no válido"); goto done; }
        } else if (!memcmp(type, "IDAT", 4)) {
            size_t need = idat_len + clen;
            if (need > TGB_MAX_FILE_BYTES) { s = tgb_fail(TGB_ERR_LIMIT, "Datos de imagen PNG demasiado grandes"); goto done; }
            if (need > idat_cap) {
                size_t nc = idat_cap ? idat_cap * 2 : 65536;
                while (nc < need) nc *= 2;
                uint8_t *ni = realloc(idat, nc);
                if (!ni) { s = tgb_fail(TGB_ERR_NOMEM, "Sin memoria para decodificar PNG"); goto done; }
                idat = ni; idat_cap = nc;
            }
            memcpy(idat + idat_len, cdata, clen);
            idat_len += clen;
        } else if (!memcmp(type, "IEND", 4)) {
            have_iend = true;
        } else if ((type[0] & 0x20) == 0) {   /* chunk crítico desconocido */
            s = tgb_fail(TGB_ERR_UNSUPPORTED, "PNG con un chunk crítico desconocido '%.4s'", (const char *)type);
            goto done;
        }
        pos += 12 + clen;
    }
    if (!have_ihdr) { s = bad("falta IHDR"); goto done; }
    if (!have_iend) { s = bad("falta IEND: el archivo está truncado"); goto done; }
    if (idat_len == 0) { s = bad("no contiene datos de imagen"); goto done; }
    if (ctype == 3 && pal_n == 0) { s = bad("falta la paleta (PLTE)"); goto done; }

    {
        int samples = ctype == 0 ? 1 : ctype == 2 ? 3 : ctype == 3 ? 1 : ctype == 4 ? 2 : 4;
        size_t bits_pp = (size_t)samples * (size_t)depth;
        size_t rowb = ((size_t)w * bits_pp + 7) / 8;
        size_t bpp = bits_pp < 8 ? 1 : bits_pp / 8;
        size_t rawsz;
        if (!tgb_mul_ok(rowb + 1, h, &rawsz)) { s = tgb_fail(TGB_ERR_LIMIT, "Imagen PNG demasiado grande"); goto done; }
        uint8_t *raw = malloc(rawsz);
        if (!raw) { s = tgb_fail(TGB_ERR_NOMEM, "Sin memoria para decodificar PNG"); goto done; }
        uLongf dest = (uLongf)rawsz;
        int zr = uncompress(raw, &dest, idat, (uLong)idat_len);
        if (zr == Z_BUF_ERROR) { free(raw); s = bad("datos comprimidos incompletos o con más contenido del esperado"); goto done; }
        if (zr != Z_OK) { free(raw); s = bad("datos comprimidos (zlib) corruptos"); goto done; }
        if (dest != rawsz) { free(raw); s = bad("los datos de imagen son más cortos de lo esperado"); goto done; }
        for (uint32_t y = 0; y < h; y++) {
            uint8_t *row = raw + (size_t)y * (rowb + 1) + 1;
            uint8_t ft = row[-1];
            const uint8_t *up = y ? row - (rowb + 1) : NULL;
            if (ft > 4) { free(raw); s = bad("tipo de filtro de fila desconocido"); goto done; }
            for (size_t i = 0; i < rowb; i++) {
                int a = i >= bpp ? row[i - bpp] : 0, b = up ? up[i] : 0, c = (up && i >= bpp) ? up[i - bpp] : 0;
                switch (ft) {
                case 1: row[i] = (uint8_t)(row[i] + a); break;
                case 2: row[i] = (uint8_t)(row[i] + b); break;
                case 3: row[i] = (uint8_t)(row[i] + (a + b) / 2); break;
                case 4: row[i] = (uint8_t)(row[i] + paeth(a, b, c)); break;
                default: break;
                }
            }
        }
        tgb_canvas *cv = tgb_canvas_create((int)w, (int)h);
        if (!cv) { free(raw); s = TGB_ERR_NOMEM; goto done; }
        uint8_t *px = tgb_canvas_pixels_mut(cv);
        unsigned maxv = (1u << (depth > 8 ? 8 : depth)) - 1u;
        for (uint32_t y = 0; y < h; y++) {
            const uint8_t *row = raw + (size_t)y * (rowb + 1) + 1;
            uint8_t *o = px + (size_t)y * w * 4;
            for (uint32_t x = 0; x < w; x++, o += 4) {
                uint8_t r, g, b, a = 255;
                if (ctype == 6) {
                    size_t k = (size_t)x * 4 * (depth / 8) ;
                    if (depth == 8) { r = row[k]; g = row[k + 1]; b = row[k + 2]; a = row[k + 3]; }
                    else { r = row[k]; g = row[k + 2]; b = row[k + 4]; a = row[k + 6]; }
                } else if (ctype == 2) {
                    size_t k = (size_t)x * 3 * (depth / 8);
                    uint16_t vr, vg, vb;
                    if (depth == 8) { vr = row[k]; vg = row[k + 1]; vb = row[k + 2]; r = row[k]; g = row[k + 1]; b = row[k + 2]; }
                    else { vr = (uint16_t)((row[k] << 8) | row[k + 1]); vg = (uint16_t)((row[k + 2] << 8) | row[k + 3]); vb = (uint16_t)((row[k + 4] << 8) | row[k + 5]); r = row[k]; g = row[k + 2]; b = row[k + 4]; }
                    if (have_trns && vr == key[0] && vg == key[1] && vb == key[2]) a = 0;
                } else if (ctype == 4) {
                    size_t k = (size_t)x * 2 * (depth / 8);
                    if (depth == 8) { r = g = b = row[k]; a = row[k + 1]; }
                    else { r = g = b = row[k]; a = row[k + 2]; }
                } else {   /* tipos 0 (gris) y 3 (paleta), cualquier profundidad */
                    unsigned v;
                    if (depth == 16) v = (unsigned)((row[x * 2] << 8) | row[x * 2 + 1]);
                    else if (depth == 8) v = row[x];
                    else { size_t bit = (size_t)x * (size_t)depth; v = (row[bit / 8] >> (8 - depth - (int)(bit % 8))) & maxv; }
                    if (ctype == 3) {
                        if (v >= pal_n) { tgb_canvas_release(cv); free(raw); s = bad("índice de paleta fuera de rango"); goto done; }
                        r = palette[v * 3]; g = palette[v * 3 + 1]; b = palette[v * 3 + 2]; a = pal_alpha[v];
                    } else {
                        unsigned g8 = depth == 16 ? (v >> 8) : (v * 255u / maxv);
                        r = g = b = (uint8_t)g8;
                        if (have_trns && v == key[0]) a = 0;
                    }
                }
                if (a == 0) { o[0] = o[1] = o[2] = o[3] = 0; }     /* forma canónica */
                else { o[0] = r; o[1] = g; o[2] = b; o[3] = a; }
            }
        }
        free(raw);
        *out = cv;
    }
done:
    free(idat);
    return s;
}

/* ================= codificación ================= */
static unsigned filter_cost(const uint8_t *row, size_t n) {
    unsigned sum = 0;
    for (size_t i = 0; i < n; i++) { int v = (int8_t)row[i]; sum += (unsigned)(v < 0 ? -v : v); }
    return sum;
}

tgb_status png_compress_image(const tgb_canvas *c, uint8_t **out, size_t *len) {
    size_t w = (size_t)tgb_canvas_width(c), h = (size_t)tgb_canvas_height(c), rowb = w * 4;
    const uint8_t *px = tgb_canvas_pixels(c);
    uint8_t *raw = malloc((rowb + 1) * h);
    uint8_t *cand = malloc(rowb * 5);
    if (!raw || !cand) { free(raw); free(cand); return tgb_fail(TGB_ERR_NOMEM, "Sin memoria para codificar PNG"); }
    for (size_t y = 0; y < h; y++) {
        const uint8_t *cur = px + y * rowb, *up = y ? px + (y - 1) * rowb : NULL;
        for (size_t i = 0; i < rowb; i++) {
            int a = i >= 4 ? cur[i - 4] : 0, b = up ? up[i] : 0, cc = (up && i >= 4) ? up[i - 4] : 0;
            cand[i] = cur[i];
            cand[rowb + i] = (uint8_t)(cur[i] - a);
            cand[2 * rowb + i] = (uint8_t)(cur[i] - b);
            cand[3 * rowb + i] = (uint8_t)(cur[i] - (a + b) / 2);
            cand[4 * rowb + i] = (uint8_t)(cur[i] - paeth(a, b, cc));
        }
        int best = 0;
        unsigned best_cost = filter_cost(cand, rowb);
        for (int f = 1; f < 5; f++) { unsigned cost = filter_cost(cand + (size_t)f * rowb, rowb); if (cost < best_cost) { best_cost = cost; best = f; } }
        raw[y * (rowb + 1)] = (uint8_t)best;
        memcpy(raw + y * (rowb + 1) + 1, cand + (size_t)best * rowb, rowb);
    }
    free(cand);
    uLongf zlen = compressBound((uLong)((rowb + 1) * h));
    uint8_t *z = malloc(zlen);
    if (!z) { free(raw); return tgb_fail(TGB_ERR_NOMEM, "Sin memoria para codificar PNG"); }
    int r = compress2(z, &zlen, raw, (uLong)((rowb + 1) * h), 6);
    free(raw);
    if (r != Z_OK) { free(z); return tgb_fail(TGB_ERR_IO, "zlib no pudo comprimir la imagen (código %d)", r); }
    *out = z;
    *len = zlen;
    return TGB_OK;
}

tgb_status tgb_png_encode_memory(const tgb_canvas *canvas, uint8_t **out_data, size_t *out_len) {
    if (!canvas || !out_data || !out_len) return tgb_fail(TGB_ERR_INVALID_ARG, "tgb_png_encode_memory: argumento nulo");
    uint8_t *z; size_t zl;
    tgb_status s = png_compress_image(canvas, &z, &zl);
    if (TGB_FAILED(s)) return s;
    tgb_sb sb; sb_init(&sb);
    uint8_t ihdr[13];
    png_put_u32(ihdr, (uint32_t)tgb_canvas_width(canvas));
    png_put_u32(ihdr + 4, (uint32_t)tgb_canvas_height(canvas));
    ihdr[8] = 8; ihdr[9] = 6; ihdr[10] = 0; ihdr[11] = 0; ihdr[12] = 0;
    png_put_signature(&sb);
    png_put_chunk(&sb, "IHDR", ihdr, 13);
    png_put_chunk(&sb, "IDAT", z, zl);
    png_put_chunk(&sb, "IEND", NULL, 0);
    free(z);
    if (sb.oom) { sb_free(&sb); return tgb_fail(TGB_ERR_NOMEM, "Sin memoria para codificar PNG"); }
    *out_data = (uint8_t *)sb_take(&sb, out_len);
    return TGB_OK;
}

void tgb_free(void *p) { free(p); }