#include "png_rgb332.h"

namespace roro::png {

uint32_t crc32(uint32_t crc, const uint8_t* data, size_t len) {
    crc = ~crc;
    for (size_t i = 0; i < len; i++) {
        crc ^= data[i];
        for (int bit = 0; bit < 8; bit++) crc = (crc >> 1) ^ (0xEDB88320u & (0u - (crc & 1)));
    }
    return ~crc;
}

uint32_t adler32(uint32_t adler, const uint8_t* data, size_t len) {
    uint32_t a = adler & 0xFFFF, b = adler >> 16;
    for (size_t i = 0; i < len; i++) {
        a = (a + data[i]) % 65521;
        b = (b + a) % 65521;
    }
    return (b << 16) | a;
}

namespace {

void be32(uint8_t* out, uint32_t v) {
    out[0] = static_cast<uint8_t>(v >> 24);
    out[1] = static_cast<uint8_t>(v >> 16);
    out[2] = static_cast<uint8_t>(v >> 8);
    out[3] = static_cast<uint8_t>(v);
}

size_t rawSize(int w, int h) { return static_cast<size_t>(w + 1) * h; }  // a filter byte before each row
size_t idatSize(int w, int h) { return 2 + 5 + rawSize(w, h) + 4; }       // zlib header, block header, data, adler

}  // namespace

size_t Rgb332Writer::fileSize(int w, int h) {
    return 8 + (12 + 13) + (12 + 768) + (12 + idatSize(w, h)) + 12;  // signature, IHDR, PLTE, IDAT, IEND
}

bool Rgb332Writer::put(const uint8_t* data, size_t len, bool inIdat) {
    if (inIdat) crc_ = crc32(crc_, data, len);
    return sink_(data, len);
}

bool Rgb332Writer::put32(uint32_t value, bool inIdat) {
    uint8_t b[4];
    be32(b, value);
    return put(b, 4, inIdat);
}

bool Rgb332Writer::begin() {
    if (w_ <= 0 || h_ <= 0 || rawSize(w_, h_) > 65535) return false;
    static const uint8_t signature[] = {0x89, 'P', 'N', 'G', '\r', '\n', 0x1A, '\n'};
    if (!sink_(signature, sizeof signature)) return false;

    uint8_t ihdr[4 + 13] = {'I', 'H', 'D', 'R'};
    be32(ihdr + 4, static_cast<uint32_t>(w_));
    be32(ihdr + 8, static_cast<uint32_t>(h_));
    ihdr[12] = 8;  // bits a pixel
    ihdr[13] = 3;  // indexed colour
    ihdr[14] = ihdr[15] = ihdr[16] = 0;
    uint8_t word[4];
    be32(word, 13);
    if (!sink_(word, 4) || !sink_(ihdr, sizeof ihdr)) return false;
    be32(word, crc32(0, ihdr, sizeof ihdr));
    if (!sink_(word, 4)) return false;

    // The palette: every RGB332 value is its own index, as scripts/rdbg.py expands them. Sixteen
    // colours at a time: this runs on a task with a small stack.
    be32(word, 768);
    const uint8_t plteKind[] = {'P', 'L', 'T', 'E'};
    if (!sink_(word, 4) || !sink_(plteKind, 4)) return false;
    uint32_t plteCrc = crc32(0, plteKind, 4);
    for (int first = 0; first < 256; first += 16) {
        uint8_t piece[48];
        for (int i = 0; i < 16; i++) {
            int v = first + i;
            piece[i * 3] = static_cast<uint8_t>((v >> 5) * 255 / 7);
            piece[i * 3 + 1] = static_cast<uint8_t>(((v >> 2) & 7) * 255 / 7);
            piece[i * 3 + 2] = static_cast<uint8_t>((v & 3) * 255 / 3);
        }
        plteCrc = crc32(plteCrc, piece, sizeof piece);
        if (!sink_(piece, sizeof piece)) return false;
    }
    be32(word, plteCrc);
    if (!sink_(word, 4)) return false;

    // IDAT: a zlib stream of one stored block. Its length is known, so it can be written first.
    size_t raw = rawSize(w_, h_);
    be32(word, static_cast<uint32_t>(idatSize(w_, h_)));
    if (!sink_(word, 4)) return false;
    crc_ = 0;
    const uint8_t head[] = {'I', 'D', 'A', 'T', 0x78, 0x01, 0x01, static_cast<uint8_t>(raw), static_cast<uint8_t>(raw >> 8),
                            static_cast<uint8_t>(~raw), static_cast<uint8_t>(~raw >> 8)};
    return put(head, sizeof head, true);
}

bool Rgb332Writer::row(const uint8_t* pixels) {
    if (rows_ >= h_) return false;
    rows_++;
    const uint8_t filter = 0;  // none
    adler_ = adler32(adler_, &filter, 1);
    adler_ = adler32(adler_, pixels, static_cast<size_t>(w_));
    return put(&filter, 1, true) && put(pixels, static_cast<size_t>(w_), true);
}

bool Rgb332Writer::end() {
    if (rows_ != h_) return false;
    if (!put32(adler_, true)) return false;
    uint8_t word[4];
    be32(word, crc_);
    if (!sink_(word, 4)) return false;
    static const uint8_t iend[] = {0, 0, 0, 0, 'I', 'E', 'N', 'D', 0xAE, 0x42, 0x60, 0x82};
    return sink_(iend, sizeof iend);
}

}  // namespace roro::png