#include "png_printer.h"

#include "../../include/debug.h"


// rewrite of TinyPngOut https://www.nayuki.io/page/tiny-png-output

#define DEFLATE_MAX_BLOCK_SIZE 0xFFFF

void pngPrinter::uint32_to_array(uint32_t src, uint8_t dest[4])
{
    dest[0] = (uint8_t)((src >> 24) & 0xff);
    dest[1] = (uint8_t)((src >> 16) & 0xff);
    dest[2] = (uint8_t)((src >> 8) & 0xff);
    dest[3] = (uint8_t)(src & 0xff);
}

uint32_t pngPrinter::update_adler32(uint32_t adler, uint8_t data)
{
    // https://gist.github.com/kornelski/710db9d30a64db0807c5bfbdbdecf85e
    unsigned s1 = adler & 0xffff;
    unsigned s2 = (adler >> 16) & 0xffff;

    s1 += data;
    s1 %= 65521;

    s2 += s1;
    s2 %= 65521;

    return (s2 << 16) | s1;
}

uint32_t pngPrinter::rc_crc32(uint32_t crc, const uint8_t *buf, size_t len)
// https://rosettacode.org/wiki/CRC-32#Implementation_2
{
    static uint32_t table[256];
    static int have_table = 0;
    uint32_t rem;
    uint8_t octet;
    int i, j;
    const uint8_t *p, *q;

    /* This check is not thread safe; there is no mutex. */
    if (have_table == 0)
    {
        /* Calculate CRC table. */
        for (i = 0; i < 256; i++)
        {
            rem = i; /* remainder from polynomial division */
            for (j = 0; j < 8; j++)
            {
                if (rem & 1)
                {
                    rem >>= 1;
                    rem ^= 0xedb88320;
                }
                else
                    rem >>= 1;
            }
            table[i] = rem;
        }
        have_table = 1;
    }

    crc = ~crc;
    q = buf + len;
    for (p = buf; p < q; p++)
    {
        octet = *p; /* Cast to unsigned octet. */
        crc = (crc >> 8) ^ table[(crc & 0xff) ^ octet];
    }
    return ~crc;
}

// uint32_t pngPrinter::rc_crc32(uint32_t crc, uint8_t c)
// // pass a single character
// {
//     rc_crc32(crc, &c, 1);
// }

void pngPrinter::png_signature()
{
    Debug_println("Writing PNG Signature.");
    uint8_t sig[] = {0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A};
    fwrite(sig, 1, 8, _file);
}

void pngPrinter::png_header()
{
    /*
        https://www.w3.org/TR/REC-png.pdf

        4.1.1 IHDR Image header
        The IHDR chunk must appear FIRST. It contains:
        Width: 4 bytes
        Height: 4 bytes
        Bit depth: 1 byte
        Color type: 1 byte
        Compression method: 1 byte
        Filter method: 1 byte
        Interlace method: 1 byte
    */
    Debug_println("Writing PNG Header.");

    uint8_t header[] = {
        // IHDR chunk
        0x00, 0x00, 0x00, 0x0D, // 0-3      13 byte data length (indexes 8-20)
        'I', 'H', 'D', 'R',     // 4-7      IHDR
        0, 0, 0, 0,             // 8-11     'width' placeholder
        0, 0, 0, 0,             // 12-15    'height' placeholder
        0x08,                   // 16       1 byte depth
        0x03,                   // 17       0x03 color with palette
        0x00,                   // 18       compression method always 0
        0x00,                   // 19       no filter
        0x00,                   // 20       no interlace
        0, 0, 0, 0,             // 21-24    IHDR CRC-32 placeholder
    };

    uint32_to_array(width, &header[8]);
    uint32_to_array(height, &header[12]);
    /* 
        https://www.w3.org/TR/REC-png.pdf
        3.2 Chunk layout
        A 4-byte CRC (Cyclic Redundancy Check) calculated 
        on the preceding bytes in the chunk, including the 
        chunk type code and chunk data fields, but 
        not including the length field.
    */
    crc_value = rc_crc32(0, &header[4], 17);
    uint32_to_array(crc_value, &header[21]);
    fwrite(header, 1, 25, _file);
}

void pngPrinter::png_palette()
{
    Debug_println("Writing PNG Palette.");
    uint8_t len[] = {0x00, 0x00, 0x00, 0x00}; // 0-3      size placeholder
    const uint8_t data[] = {
        // IDAT chunk data
        'P', 'L', 'T', 'E', // 4-7      PLTE
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x13, 0x13, 0x13, 0x25, 0x25, 0x25, 0x37, 0x37, 0x37, 0x49,
        0x49, 0x49, 0x5F, 0x5F, 0x5F, 0x71, 0x71, 0x71, 0x7A, 0x7A, 0x7A, 0x8C, 0x8C, 0x8C, 0xA1, 0xA1,
        0xA1, 0xB3, 0xB3, 0xB3, 0xC5, 0xC5, 0xC5, 0xD7, 0xD7, 0xD7, 0xED, 0xED, 0xED, 0xFF, 0xFF, 0xFF,
        0x0A, 0x00, 0x00, 0x1C, 0x0A, 0x00, 0x32, 0x1F, 0x00, 0x44, 0x31, 0x00, 0x56, 0x43, 0x00, 0x68,
        0x55, 0x00, 0x7D, 0x6B, 0x00, 0x90, 0x7D, 0x00, 0x98, 0x86, 0x00, 0xAA, 0x98, 0x00, 0xC0, 0xAD,
        0x13, 0xD2, 0xBF, 0x25, 0xE4, 0xD1, 0x37, 0xF6, 0xE3, 0x49, 0xFF, 0xF9, 0x5F, 0xFF, 0xFF, 0x71,
        0x2A, 0x00, 0x00, 0x3C, 0x00, 0x00, 0x51, 0x08, 0x00, 0x63, 0x1A, 0x00, 0x75, 0x2C, 0x00, 0x87,
        0x3E, 0x00, 0x9D, 0x54, 0x00, 0xAF, 0x66, 0x08, 0xB8, 0x6E, 0x11, 0xCA, 0x81, 0x23, 0xDF, 0x96,
        0x38, 0xF1, 0xA8, 0x4A, 0xFF, 0xBA, 0x5C, 0xFF, 0xCC, 0x6F, 0xFF, 0xE2, 0x84, 0xFF, 0xF4, 0x96,
        0x3D, 0x00, 0x00, 0x4F, 0x00, 0x00, 0x64, 0x00, 0x00, 0x77, 0x05, 0x00, 0x89, 0x17, 0x08, 0x9B,
        0x29, 0x1A, 0xB0, 0x3F, 0x30, 0xC2, 0x51, 0x42, 0xCB, 0x59, 0x4A, 0xDD, 0x6B, 0x5D, 0xF2, 0x81,
        0x72, 0xFF, 0x93, 0x84, 0xFF, 0xA5, 0x96, 0xFF, 0xB7, 0xA8, 0xFF, 0xCD, 0xBE, 0xFF, 0xDF, 0xD0,
        0x40, 0x00, 0x00, 0x52, 0x00, 0x12, 0x68, 0x00, 0x27, 0x7A, 0x00, 0x39, 0x8C, 0x08, 0x4B, 0x9E,
        0x1A, 0x5D, 0xB4, 0x30, 0x73, 0xC6, 0x42, 0x85, 0xCE, 0x4B, 0x8D, 0xE1, 0x5D, 0xA0, 0xF6, 0x72,
        0xB5, 0xFF, 0x84, 0xC7, 0xFF, 0x96, 0xD9, 0xFF, 0xA8, 0xEB, 0xFF, 0xBE, 0xFF, 0xFF, 0xD0, 0xFF,
        0x33, 0x00, 0x3F, 0x45, 0x00, 0x51, 0x5B, 0x00, 0x66, 0x6D, 0x00, 0x78, 0x7F, 0x03, 0x8A, 0x91,
        0x15, 0x9C, 0xA7, 0x2A, 0xB2, 0xB9, 0x3C, 0xC4, 0xC1, 0x45, 0xCD, 0xD3, 0x57, 0xDF, 0xE9, 0x6D,
        0xF4, 0xFB, 0x7F, 0xFF, 0xFF, 0x91, 0xFF, 0xFF, 0xA3, 0xFF, 0xFF, 0xB8, 0xFF, 0xFF, 0xCA, 0xFF,
        0x18, 0x00, 0x6E, 0x2A, 0x00, 0x80, 0x40, 0x00, 0x95, 0x52, 0x00, 0xA7, 0x64, 0x07, 0xB9, 0x76,
        0x19, 0xCB, 0x8C, 0x2F, 0xE1, 0x9E, 0x41, 0xF3, 0xA6, 0x4A, 0xFC, 0xB8, 0x5C, 0xFF, 0xCE, 0x71,
        0xFF, 0xE0, 0x83, 0xFF, 0xF2, 0x95, 0xFF, 0xFF, 0xA8, 0xFF, 0xFF, 0xBD, 0xFF, 0xFF, 0xCF, 0xFF,
        0x00, 0x00, 0x83, 0x07, 0x00, 0x95, 0x1C, 0x00, 0xAB, 0x2E, 0x03, 0xBD, 0x40, 0x15, 0xCF, 0x52,
        0x27, 0xE1, 0x68, 0x3D, 0xF6, 0x7A, 0x4F, 0xFF, 0x83, 0x58, 0xFF, 0x95, 0x6A, 0xFF, 0xAA, 0x7F,
        0xFF, 0xBC, 0x91, 0xFF, 0xCE, 0xA3, 0xFF, 0xE0, 0xB6, 0xFF, 0xF6, 0xCB, 0xFF, 0xFF, 0xDD, 0xFF,
        0x00, 0x00, 0x7B, 0x00, 0x00, 0x8D, 0x00, 0x06, 0xA3, 0x09, 0x18, 0xB5, 0x1B, 0x2A, 0xC7, 0x2D,
        0x3C, 0xD9, 0x42, 0x52, 0xEF, 0x54, 0x64, 0xFF, 0x5D, 0x6C, 0xFF, 0x6F, 0x7E, 0xFF, 0x85, 0x94,
        0xFF, 0x97, 0xA6, 0xFF, 0xA9, 0xB8, 0xFF, 0xBB, 0xCA, 0xFF, 0xD0, 0xE0, 0xFF, 0xE2, 0xF2, 0xFF,
        0x00, 0x00, 0x57, 0x00, 0x08, 0x6A, 0x00, 0x1D, 0x7F, 0x00, 0x2F, 0x91, 0x00, 0x41, 0xA3, 0x0D,
        0x53, 0xB5, 0x22, 0x69, 0xCB, 0x35, 0x7B, 0xDD, 0x3D, 0x83, 0xE5, 0x4F, 0x96, 0xF8, 0x65, 0xAB,
        0xFF, 0x77, 0xBD, 0xFF, 0x89, 0xCF, 0xFF, 0x9B, 0xE1, 0xFF, 0xB0, 0xF7, 0xFF, 0xC3, 0xFF, 0xFF,
        0x00, 0x0B, 0x1E, 0x00, 0x1D, 0x31, 0x00, 0x32, 0x46, 0x00, 0x44, 0x58, 0x00, 0x57, 0x6A, 0x00,
        0x69, 0x7C, 0x0E, 0x7E, 0x92, 0x20, 0x90, 0xA4, 0x29, 0x99, 0xAD, 0x3B, 0xAB, 0xBF, 0x51, 0xC0,
        0xD4, 0x63, 0xD2, 0xE6, 0x75, 0xE5, 0xF8, 0x87, 0xF7, 0xFF, 0x9C, 0xFF, 0xFF, 0xAE, 0xFF, 0xFF,
        0x00, 0x1A, 0x00, 0x00, 0x2C, 0x00, 0x00, 0x42, 0x03, 0x00, 0x54, 0x15, 0x00, 0x66, 0x27, 0x00,
        0x78, 0x3A, 0x0A, 0x8D, 0x4F, 0x1C, 0x9F, 0x61, 0x25, 0xA8, 0x6A, 0x37, 0xBA, 0x7C, 0x4C, 0xD0,
        0x91, 0x5E, 0xE2, 0xA3, 0x70, 0xF4, 0xB5, 0x82, 0xFF, 0xC8, 0x98, 0xFF, 0xDD, 0xAA, 0xFF, 0xEF,
        0x00, 0x20, 0x00, 0x00, 0x32, 0x00, 0x00, 0x48, 0x00, 0x00, 0x5A, 0x00, 0x00, 0x6C, 0x00, 0x01,
        0x7E, 0x00, 0x16, 0x93, 0x0F, 0x28, 0xA6, 0x21, 0x31, 0xAE, 0x2A, 0x43, 0xC0, 0x3C, 0x58, 0xD6,
        0x52, 0x6A, 0xE8, 0x64, 0x7C, 0xFA, 0x76, 0x8F, 0xFF, 0x88, 0xA4, 0xFF, 0x9D, 0xB6, 0xFF, 0xAF,
        0x00, 0x1C, 0x00, 0x00, 0x2E, 0x00, 0x00, 0x44, 0x00, 0x00, 0x56, 0x00, 0x09, 0x68, 0x00, 0x1B,
        0x7A, 0x00, 0x30, 0x8F, 0x00, 0x42, 0xA2, 0x00, 0x4B, 0xAA, 0x00, 0x5D, 0xBC, 0x0C, 0x73, 0xD2,
        0x21, 0x85, 0xE4, 0x33, 0x97, 0xF6, 0x46, 0xA9, 0xFF, 0x58, 0xBE, 0xFF, 0x6D, 0xD0, 0xFF, 0x7F,
        0x00, 0x0F, 0x00, 0x00, 0x21, 0x00, 0x08, 0x36, 0x00, 0x1A, 0x48, 0x00, 0x2C, 0x5A, 0x00, 0x3E,
        0x6C, 0x00, 0x54, 0x82, 0x00, 0x66, 0x94, 0x00, 0x6E, 0x9D, 0x00, 0x81, 0xAF, 0x00, 0x96, 0xC4,
        0x0A, 0xA8, 0xD6, 0x1C, 0xBA, 0xE8, 0x2E, 0xCC, 0xFA, 0x40, 0xE2, 0xFF, 0x56, 0xF4, 0xFF, 0x68,
        0x06, 0x00, 0x00, 0x18, 0x0C, 0x00, 0x2E, 0x22, 0x00, 0x40, 0x34, 0x00, 0x52, 0x46, 0x00, 0x64,
        0x58, 0x00, 0x79, 0x6E, 0x00, 0x8B, 0x80, 0x00, 0x94, 0x88, 0x00, 0xA6, 0x9A, 0x00, 0xBC, 0xB0,
        0x10, 0xCE, 0xC2, 0x22, 0xE0, 0xD4, 0x34, 0xF2, 0xE6, 0x47, 0xFF, 0xFC, 0x5C, 0xFF, 0xFF, 0x6E};
    uint8_t ccc[] = {0, 0, 0, 0}; // crc placeholder

    uint32_to_array(768, &len[0]);
    crc_value = rc_crc32(0, &data[0], 4 + 768);
    uint32_to_array(crc_value, &ccc[0]);

    fwrite(len, 1, 4, _file);
    fwrite(data, 1, 4 + 768, _file);
    fwrite(ccc, 1, 4, _file);
}

void pngPrinter::png_data()
{
    /*  
    https://www.w3.org/TR/REC-png.pdf

    4.1.3 IDAT Image data
    The IDAT chunk contains the actual image data. To create this data:
    1. Begin with image scanlines represented as described in Image layout (Section 2.3); the layout and total
    size of this raw data are determined by the fields of IHDR.
    2. Filter the image data according to the filtering method specified by the IHDR chunk. (Note that with
    filter method 0, the only one currently defined, this implies prepending a filter type byte to each scanline.)
    3. Compress the filtered data using the compression method specified by the IHDR chunk.
    The IDAT chunk contains the output datastream of the compression algorithm.
    To read the image data, reverse this process.
    There can be multiple IDAT chunks; if so, they must appear consecutively with no other intervening chunks.
    The compressed datastream is then the concatenation of the contents of all the IDAT chunks. The encoder
    can divide the compressed datastream into IDAT chunks however it wishes. (Multiple IDAT chunks are
    allowed so that encoders can work in a fixed amount of memory; typically the chunk size will correspond
    to the encoder’s buffer size.) It is important to emphasize that IDAT chunk boundaries have no semantic
    significance and can occur at any point in the compressed datastream
*/
    Debug_println("Starting PNG Image Data...");
    uint8_t data[] = {
        // IDAT chunk
        0x00, 0x00, 0x00, 0x00, // 0-3      size placeholder
        'I', 'D', 'A', 'T',     // 4-7      IDAT
    };
    crc_value = rc_crc32(0, &data[4], 4); // begin CRC calculation

    // Compute data size
    // imgSize = (width + 1) * height; // +1 per line for filter 0's
    uint32_t numBlocks = imgSize / DEFLATE_MAX_BLOCK_SIZE;
    if (imgSize % DEFLATE_MAX_BLOCK_SIZE != 0)
        numBlocks++; // Round up

    dataSize += numBlocks * 5; // 5 bytes per DEFLATE uncompressed block header
    dataSize += 2;             // 2 bytes for zlib header
    dataSize += 4;             // 4 bytes for zlib Adler-32 checksum
    dataSize += imgSize;       // plus the image data

    uint32_to_array(dataSize, &data[0]); // store the computed size

    fwrite(data, 1, 8, _file); // write out the IDAT header
}

void pngPrinter::png_add_data(uint8_t *buf, uint32_t n)
{
    uint8_t c = 0;
    // Deflate-compressed datastreams within PNG are stored in the “zlib” format
    // https://tools.ietf.org/html/rfc1950#page-4

    if (img_pos == 0)
    {
        Debug_println("Writing ZLIB header.");
        // write out a ZLIB header
        // Compression method/flags code: 1 byte (For PNG compression method 0, the zlib compression method/flags code must specify method code 8 (“deflate” compression))
        c = 0x08; // ZLIB "Deflate" compression scheme
        crc_value = rc_crc32(crc_value, c);
        fputc(c, _file);
        //  Additional flags/check bits: 1 byte (must be such that method + flags, when viewed as a 16-bit unsigned integer stored in MSB order (CMF*256 + FLG), is a multiple of 31.)
        c = 0x1D; // precompute so that 0x081D is divisible by 31 [ (0x800 / 31 + 1) * 31 - 0x800 ]
        crc_value = rc_crc32(crc_value, c);
        fputc(c, _file);

        //printf("new image\r\n");
    }

    uint32_t idx = 0;
    while (idx < n)
    {
        // at beginning of block?
        if (blk_pos == 0)
        {
            blkSize = DEFLATE_MAX_BLOCK_SIZE;
            uint8_t c = 0;
            if (imgSize - img_pos <= (uint32_t)blkSize)
            {
                c = 1; // final block
                blkSize = (uint16_t)(imgSize - img_pos);
            }

            Debug_println("Writing ZLIB block header.");
            // write out block header
            crc_value = rc_crc32(crc_value, c);
            fputc(c, _file);

            // write out block size
            c = (uint8_t)(blkSize >> 0);
            crc_value = rc_crc32(crc_value, c);
            fputc(c, _file);
            c = (uint8_t)(blkSize >> 8);
            crc_value = rc_crc32(crc_value, c);
            fputc(c, _file);
            c = (uint8_t)((blkSize >> 0) ^ 0xFF);
            crc_value = rc_crc32(crc_value, c);
            fputc(c, _file);
            c = (uint8_t)((blkSize >> 8) ^ 0xFF);
            crc_value = rc_crc32(crc_value, c);
            fputc(c, _file);

            //printf("new block\r\n");
        }

        //at beginning of a line?
        if (Xpos == 0)
        {
            Debug_printf("Starting PNG line %d ... ",Ypos);
            c = 0;
            crc_value = rc_crc32(crc_value, c);
            adler_value = update_adler32(adler_value, c);
            fputc(c, _file);
            //printf("\nnew line %d ", c);

            img_pos++;
            blk_pos++;
        }

        // put byte from buffer
        c = buf[idx];
        crc_value = rc_crc32(crc_value, c);
        adler_value = update_adler32(adler_value, c);
        fputc(c, _file);
        //printf("%d ", c);

        Xpos++;
        idx++;
        img_pos++;
        blk_pos++;

        // check for end of's
        if (Xpos == width)
        {
            Debug_println("Finished PNG line.");
            Xpos = 0;
            Ypos++;
        }
        if (blk_pos == blkSize)
        {
            Debug_println("End of ZLIB block.");
            blk_pos = 0;
        }
    };

    if (img_pos == imgSize)
    {
        Debug_println("Writing ZLIB Adler checksum and PNG data CRC.");
        uint8_t data[] = {
            0, 0, 0, 0, // Adler32 Check value: 4 bytes
            0, 0, 0, 0  // CRC32: 4 bytes
        };
        uint32_to_array(adler_value, &data[0]);
        for (int i = 0; i < 4; i++)
            crc_value = rc_crc32(crc_value, data[i]); // crc_value = rc_crc32(crc_value, data[0],4); ?????????
        uint32_to_array(crc_value, &data[4]);
        fwrite(data, 1, 8, _file);
        png_end();
    }
}

void pngPrinter::png_end()
{
    Debug_println("Writing PNG footer.");
    unsigned char end[] = {
        0x00, 0x00, 0x00, 0x00,  // zero length
        'I', 'E', 'N', 'D',      // IEND
        0xAE, 0x42, 0x60, 0x82}; // crc32 - precompute because always the same
    fwrite(end, 1, 12, _file);
}

// TODO: Anything here?
void pngPrinter::pre_close_file()
{

}

void pngPrinter::post_new_file()
{
    // call PNG header routines
    png_signature();
    png_header();
    png_palette();
    // start IDAT chunk and now ready for data
    png_data();
}

bool pngPrinter::process_buffer(uint8_t n, uint8_t aux1, uint8_t aux2)
{
// copy buffer[] into linebuffer[]
    Debug_printf("%d bytes rx'd by PNG printer\r\n", n);
    uint16_t i = 0;
    while (i < n && img_pos < imgSize)
    {
        //Debug_println("processing buffer.");
        if (BOLflag)
        {
            Debug_println("Processing new line!");
            rep_code = buffer[i++];
            BOLflag = false;
        }
        else
        {
            line_buffer[line_index++] = buffer[i++];
        }
        if (line_index == 320)
        {
            while (rep_code-- > 0)
            {
                Debug_printf("Adding line %d\r\n", rep_code);
                png_add_data(&line_buffer[0], 320);
            }
            BOLflag = true;
            line_index = 0;
        }
    }
    return true;
}

