/*
 * pnglite.c - pnglite library
 * For conditions of distribution and use, see copyright notice in pnglite.h
 */

#ifdef _STANDALONE
#include <sys/cdefs.h>
#include <stand.h>
#else
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#endif
#include <zlib.h>
#include "pnglite.h"

#ifndef abs
#define abs(x)  ((x) < 0? -(x):(x))
#endif

#define PNG_32b(b, s) ((uint32_t)(b) << (s))
#define PNG_U32(b1, b2, b3, b4) \
        (PNG_32b(b1, 24) | PNG_32b(b2, 16) | PNG_32b(b3, 8) | PNG_32b(b4, 0))

#define png_IDAT PNG_U32(73,  68,  65,  84)
#define png_IEND PNG_U32(73,  69,  78,  68)

static ssize_t
file_read(png_t *png, void *out, size_t size, size_t numel)
{
        ssize_t result;
        off_t offset = (off_t)(size * numel);

        if (offset < 0)
                return (PNG_FILE_ERROR);

        if (!out) {
                result = lseek(png->fd, offset, SEEK_CUR);
        } else {
                result = read(png->fd, out, size * numel);
        }

        return (result);
}

static int
file_read_ul(png_t *png, unsigned *out)
{
        uint8_t buf[4];

        if (file_read(png, buf, 1, 4) != 4)
                return (PNG_FILE_ERROR);

        *out = (buf[0]<<24) | (buf[1]<<16) | (buf[2]<<8) | buf[3];

        return (PNG_NO_ERROR);
}

static unsigned
get_ul(uint8_t *buf)
{
        unsigned result;
        uint8_t foo[4];

        memcpy(foo, buf, 4);

        result = (foo[0]<<24) | (foo[1]<<16) | (foo[2]<<8) | foo[3];

        return (result);
}

static int
png_get_bpp(png_t *png)
{
        int bpp;

        switch (png->color_type) {
        case PNG_GREYSCALE:
                bpp = 1; break;
        case PNG_TRUECOLOR:
                bpp = 3; break;
        case PNG_INDEXED:
                bpp = 1; break;
        case PNG_GREYSCALE_ALPHA:
                bpp = 2; break;
        case PNG_TRUECOLOR_ALPHA:
                bpp = 4; break;
        default:
                return (PNG_FILE_ERROR);
        }

        bpp *= png->depth / 8;

        return (bpp);
}

static int
png_read_ihdr(png_t *png)
{
        unsigned length = 0;
        unsigned orig_crc;
        unsigned calc_crc;
        uint8_t ihdr[13+4]; /* length should be 13, make room for type (IHDR) */

        if (file_read_ul(png, &length) != PNG_NO_ERROR)
                return (PNG_FILE_ERROR);

        if (length != 13)
                return (PNG_CRC_ERROR);

        if (file_read(png, ihdr, 1, 13+4) != 13+4)
                return (PNG_EOF_ERROR);

        if (file_read_ul(png, &orig_crc) != PNG_NO_ERROR)
                return (PNG_FILE_ERROR);

        calc_crc = crc32(0L, Z_NULL, 0);
        calc_crc = crc32(calc_crc, ihdr, 13+4);

        if (orig_crc != calc_crc) {
                return (PNG_CRC_ERROR);
        }

        png->width = get_ul(ihdr+4);
        png->height = get_ul(ihdr+8);
        png->depth = ihdr[12];
        png->color_type = ihdr[13];
        png->compression_method = ihdr[14];
        png->filter_method = ihdr[15];
        png->interlace_method = ihdr[16];

        if (png->color_type == PNG_INDEXED)
                return (PNG_NOT_SUPPORTED);

        if (png->depth != 8 && png->depth != 16)
                return (PNG_NOT_SUPPORTED);

        if (png->interlace_method)
                return (PNG_NOT_SUPPORTED);

        return (PNG_NO_ERROR);
}

void
png_print_info(png_t *png)
{
        printf("PNG INFO:\n");
        printf("\twidth:\t\t%d\n", png->width);
        printf("\theight:\t\t%d\n", png->height);
        printf("\tdepth:\t\t%d\n", png->depth);
        printf("\tcolor:\t\t");

        switch (png->color_type) {
        case PNG_GREYSCALE:
                printf("greyscale\n"); break;
        case PNG_TRUECOLOR:
                printf("truecolor\n"); break;
        case PNG_INDEXED:
                printf("palette\n"); break;
        case PNG_GREYSCALE_ALPHA:
                printf("greyscale with alpha\n"); break;
        case PNG_TRUECOLOR_ALPHA:
                printf("truecolor with alpha\n"); break;
        default:
                printf("unknown, this is not good\n"); break;
        }

        printf("\tcompression:\t%s\n",
            png->compression_method?
            "unknown, this is not good":"inflate/deflate");
        printf("\tfilter:\t\t%s\n",
            png->filter_method? "unknown, this is not good":"adaptive");
        printf("\tinterlace:\t%s\n",
            png->interlace_method? "interlace":"no interlace");
}

int
png_open(png_t *png, const char *filename)
{
        char header[8];
        int result;

        png->image = NULL;
        png->fd = open(filename, O_RDONLY);
        if (png->fd == -1)
                return (PNG_FILE_ERROR);

        if (file_read(png, header, 1, 8) != 8) {
                result = PNG_EOF_ERROR;
                goto done;
        }

        if (memcmp(header, "\x89\x50\x4E\x47\x0D\x0A\x1A\x0A", 8) != 0) {
                result = PNG_HEADER_ERROR;
                goto done;
        }

        result = png_read_ihdr(png);
        if (result == PNG_NO_ERROR) {
                result = png_get_bpp(png);
                if (result > 0) {
                        png->bpp = (uint8_t)result;
                        result = PNG_NO_ERROR;
                }
        }

done:
        if (result == PNG_NO_ERROR) {
                uint64_t size = png->width * png->height * png->bpp;

                if (size < UINT_MAX)
                        png->image = malloc(size);
                if (png->image == NULL)
                        result = PNG_MEMORY_ERROR;
        }

        if (result == PNG_NO_ERROR)
                result = png_get_data(png, png->image);

        if (result != PNG_NO_ERROR) {
                free(png->image);
                (void) close(png->fd);
                png->fd = -1;
                return (result);
        }

        return (result);
}

int
png_close(png_t *png)
{
        (void) close(png->fd);
        png->fd = -1;
        free(png->image);
        png->image = NULL;

        return (PNG_NO_ERROR);
}

static int
png_init_inflate(png_t *png)
{
        z_stream *stream;
        png->zs = calloc(1, sizeof (z_stream));

        stream = png->zs;

        if (!stream)
                return (PNG_MEMORY_ERROR);

        if (inflateInit(stream) != Z_OK) {
                free(png->zs);
                png->zs = NULL;
                return (PNG_ZLIB_ERROR);
        }

        stream->next_out = png->png_data;
        stream->avail_out = png->png_datalen;

        return (PNG_NO_ERROR);
}

static int
png_end_inflate(png_t *png)
{
        z_stream *stream = png->zs;
        int rc = PNG_NO_ERROR;

        if (!stream)
                return (PNG_MEMORY_ERROR);

        if (inflateEnd(stream) != Z_OK) {
                printf("ZLIB says: %s\n", stream->msg);
                rc = PNG_ZLIB_ERROR;
        }

        free(png->zs);
        png->zs = NULL;

        return (rc);
}

static int
png_inflate(png_t *png, uint8_t *data, int len)
{
        int result;
        z_stream *stream = png->zs;

        if (!stream)
                return (PNG_MEMORY_ERROR);

        stream->next_in = data;
        stream->avail_in = len;

        result = inflate(stream, Z_SYNC_FLUSH);

        if (result != Z_STREAM_END && result != Z_OK) {
                printf("%s\n", stream->msg);
                return (PNG_ZLIB_ERROR);
        }

        if (stream->avail_in != 0)
                return (PNG_ZLIB_ERROR);

        return (PNG_NO_ERROR);
}

static int
png_read_idat(png_t *png, unsigned length)
{
        unsigned orig_crc;
        unsigned calc_crc;
        ssize_t len = length;

        if (!png->readbuf || png->readbuflen < length) {
                png->readbuf = realloc(png->readbuf, length);
                png->readbuflen = length;
        }

        if (!png->readbuf)
                return (PNG_MEMORY_ERROR);

        if (file_read(png, png->readbuf, 1, length) != len)
                return (PNG_FILE_ERROR);

        calc_crc = crc32(0L, Z_NULL, 0);
        calc_crc = crc32(calc_crc, (uint8_t *)"IDAT", 4);
        calc_crc = crc32(calc_crc, (uint8_t *)png->readbuf, length);

        if (file_read_ul(png, &orig_crc) != PNG_NO_ERROR)
                return (PNG_FILE_ERROR);

        if (orig_crc != calc_crc)
                return (PNG_CRC_ERROR);

        return (png_inflate(png, png->readbuf, length));
}

static int
png_process_chunk(png_t *png)
{
        int result = PNG_NO_ERROR;
        unsigned type;
        unsigned length;

        if (file_read_ul(png, &length) != PNG_NO_ERROR)
                return (PNG_FILE_ERROR);

        if (file_read_ul(png, &type) != PNG_NO_ERROR)
                return (PNG_FILE_ERROR);

        /*
         * if we found an idat, all other idats should be followed with no
         * other chunks in between
         */
        if (type == png_IDAT) {
                if (!png->png_data) {   /* first IDAT */
                        png->png_datalen = png->width * png->height *
                            png->bpp + png->height;
                        png->png_data = malloc(png->png_datalen);
                }

                if (!png->png_data)
                        return (PNG_MEMORY_ERROR);

                if (!png->zs) {
                        result = png_init_inflate(png);
                        if (result != PNG_NO_ERROR)
                                return (result);
                }

                return (png_read_idat(png, length));
        } else if (type == png_IEND)
                return (PNG_DONE);
        else
                (void) file_read(png, 0, 1, length + 4); /* unknown chunk */

        return (result);
}

static void
png_filter_sub(unsigned stride, uint8_t *in, uint8_t *out, unsigned len)
{
        unsigned i;
        uint8_t a = 0;

        for (i = 0; i < len; i++) {
                if (i >= stride)
                        a = out[i - stride];

                out[i] = in[i] + a;
        }
}

static void
png_filter_up(unsigned stride __unused, uint8_t *in, uint8_t *out,
    uint8_t *prev_line, unsigned len)
{
        unsigned i;

        if (prev_line) {
                for (i = 0; i < len; i++)
                        out[i] = in[i] + prev_line[i];
        } else
                memcpy(out, in, len);
}

static void
png_filter_average(unsigned stride, uint8_t *in, uint8_t *out,
    uint8_t *prev_line, unsigned len)
{
        unsigned int i;
        uint8_t a = 0;
        uint8_t b = 0;
        unsigned int sum = 0;

        for (i = 0; i < len; i++) {
                if (prev_line)
                        b = prev_line[i];

                if (i >= stride)
                        a = out[i - stride];

                sum = a;
                sum += b;

                out[i] = in[i] + sum/2;
        }
}

static uint8_t
png_paeth(uint8_t a, uint8_t b, uint8_t c)
{
        int p = (int)a + b - c;
        int pa = abs(p - a);
        int pb = abs(p - b);
        int pc = abs(p - c);

        int pr;

        if (pa <= pb && pa <= pc)
                pr = a;
        else if (pb <= pc)
                pr = b;
        else
                pr = c;

        return (pr);
}

static void
png_filter_paeth(unsigned stride, uint8_t *in, uint8_t *out, uint8_t *prev_line,
    unsigned len)
{
        unsigned i;
        uint8_t a;
        uint8_t b;
        uint8_t c;

        for (i = 0; i < len; i++) {
                if (prev_line && i >= stride) {
                        a = out[i - stride];
                        b = prev_line[i];
                        c = prev_line[i - stride];
                } else {
                        if (prev_line)
                                b = prev_line[i];
                        else
                                b = 0;

                        if (i >= stride)
                                a = out[i - stride];
                        else
                                a = 0;

                        c = 0;
                }

                out[i] = in[i] + png_paeth(a, b, c);
        }
}

static int
png_unfilter(png_t *png, uint8_t *data)
{
        unsigned i;
        unsigned pos = 0;
        unsigned outpos = 0;
        uint8_t *filtered = png->png_data;
        unsigned stride = png->bpp;

        while (pos < png->png_datalen) {
                uint8_t filter = filtered[pos];

                pos++;

                if (png->depth == 16) {
                        for (i = 0; i < png->width * stride; i += 2) {
                                *(short *)(filtered+pos+i) =
                                    (filtered[pos+i] << 8) | filtered[pos+i+1];
                        }
                }

                switch (filter) {
                case 0: /* none */
                        memcpy(data+outpos, filtered+pos, png->width * stride);
                        break;
                case 1: /* sub */
                        png_filter_sub(stride, filtered+pos, data+outpos,
                            png->width * stride);
                        break;
                case 2: /* up */
                        if (outpos) {
                                png_filter_up(stride, filtered+pos, data+outpos,
                                    data + outpos - (png->width*stride),
                                    png->width*stride);
                        } else {
                                png_filter_up(stride, filtered+pos, data+outpos,
                                    0, png->width*stride);
                        }
                        break;
                case 3: /* average */
                        if (outpos) {
                                png_filter_average(stride, filtered+pos,
                                    data+outpos,
                                    data + outpos - (png->width*stride),
                                    png->width*stride);
                        } else {
                                png_filter_average(stride, filtered+pos,
                                    data+outpos, 0, png->width*stride);
                        }
                        break;
                case 4: /* paeth */
                        if (outpos) {
                                png_filter_paeth(stride, filtered+pos,
                                    data+outpos,
                                    data + outpos - (png->width*stride),
                                    png->width*stride);
                        } else {
                                png_filter_paeth(stride, filtered+pos,
                                    data+outpos, 0, png->width*stride);
                        }
                        break;
                default:
                        return (PNG_UNKNOWN_FILTER);
                }

                outpos += png->width * stride;
                pos += png->width * stride;
        }

        return (PNG_NO_ERROR);
}

int
png_get_data(png_t *png, uint8_t *data)
{
        int result = PNG_NO_ERROR;

        png->zs = NULL;
        png->png_datalen = 0;
        png->png_data = NULL;
        png->readbuf = NULL;
        png->readbuflen = 0;

        while (result == PNG_NO_ERROR)
                result = png_process_chunk(png);

        if (png->readbuf) {
                free(png->readbuf);
                png->readbuflen = 0;
        }
        if (png->zs)
                (void) png_end_inflate(png);

        if (result != PNG_DONE) {
                free(png->png_data);
                return (result);
        }

        result = png_unfilter(png, data);

        free(png->png_data);

        return (result);
}

char *
png_error_string(int error)
{
        switch (error) {
        case PNG_NO_ERROR:
                return ("No error");
        case PNG_FILE_ERROR:
                return ("Unknown file error.");
        case PNG_HEADER_ERROR:
                return ("No PNG header found. Are you sure this is a PNG?");
        case PNG_IO_ERROR:
                return ("Failure while reading file.");
        case PNG_EOF_ERROR:
                return ("Reached end of file.");
        case PNG_CRC_ERROR:
                return ("CRC or chunk length error.");
        case PNG_MEMORY_ERROR:
                return ("Could not allocate memory.");
        case PNG_ZLIB_ERROR:
                return ("zlib reported an error.");
        case PNG_UNKNOWN_FILTER:
                return ("Unknown filter method used in scanline.");
        case PNG_DONE:
                return ("PNG done");
        case PNG_NOT_SUPPORTED:
                return ("The PNG is unsupported by pnglite, too bad for you!");
        case PNG_WRONG_ARGUMENTS:
                return ("Wrong combination of arguments passed to png_open. "
                    "You must use either a read_function or supply a file "
                    "pointer to use.");
        default:
                return ("Unknown error.");
        };
}