/*******************************************************************************

Svend Tofte, www.svendtofte.com
This code is in the public domain.

********************************************************************************

PNG format is as following:

[PNG header, 8 bytes]
[Chunk size N, 4 bytes]
[Chunk type in ASCII, 4 bytes]
[Chunk data, N bytes]
[Chunk CRC, 4 bytes]

Chunk size only describes the size of the data field. The integer is big-endian, 
4 bytes unsigned. The chunks repeat any number of times, in order to describe 
the whole file. The image itself is contained in one or more "IDAT" chunks.

Some of the chunks have a certain ordering:

IHDR  Shall be first
PLTE  Before first IDAT (optional)
cHRM  Before first IDAT (optional)
gAMA  Before first IDAT (optional)
sRGB  Before first IDAT (optional)
IDAT  Multiple IDATS may appear consecutively
IEND  Shall be last

More:
http://www.w3.org/TR/PNG/#5DataRep

*******************************************************************************/
#include <stdio.h>
#include <stdlib.h>

#define HEADER_SIZE 8
#define BUFFER_SIZE 1024
#define TMP_FILE_SUFFIX ".stripped.png"
#define CHUNK_PART_SIZE 4
#define CHUNK_HDR_GAMA 0x67414D41 // "gAMA" -> 0x67,0x41,0x4D,0x41
#define CHUNK_HDR_CHRM 0x6348524D // "cHRM" -> etc
#define CHUNK_HDR_SRGB 0x73524742
#define CHUNK_HDR_IEND 0x49454E44

// http://www.ibm.com/developerworks/aix/library/au-endianc/index.html
const int endian_test = 1;
#define is_bigendian() ( *((char*)&endian_test) == 0 )

unsigned char* std_header = "\x89PNG\xD\xA\x1A\xA"; // png file header
unsigned char chunk_size[CHUNK_PART_SIZE];
unsigned char chunk_type[CHUNK_PART_SIZE];
unsigned char chunk_data[BUFFER_SIZE];
unsigned char chunk_crc[CHUNK_PART_SIZE];
int last_chunk = 0;
FILE *infp, *outfp;

// forward decls
void write_chunk(char* s, int n, FILE* fp);
int read_chunk(char* s, int n, FILE* fp);
int skip_chunk(char* s);
int swap_endianness(int i);
void chunk_error();

void main(int argc, char* argv[]) {
    int i, j;
    unsigned int cdata_size;
    unsigned char c;
    unsigned char* copy;
    
    if (argc==1) {
        printf("No filename argument provided\n");
        exit(EXIT_FAILURE);
    }
    
    if (CHUNK_PART_SIZE != sizeof(i)) {
        printf("Need 4 byte integer size\n");
        exit(EXIT_FAILURE);
    }
    
    infp = fopen(argv[1],"rb");
    if (infp == NULL) {
        printf("Could not open %s\n", argv[1]);
        exit(EXIT_FAILURE);
    }
    
    for (i = 0; i < HEADER_SIZE; i++) {
        c = (unsigned char)fgetc(infp);
        if (c != std_header[i]) {
            printf("%s does not look like a PNG file\n", argv[1]);
            fclose(infp);
            exit(EXIT_FAILURE);
        }
    }
    
    copy = (unsigned char*)malloc(strlen(argv[1])+strlen(TMP_FILE_SUFFIX));
    *copy = '\0';
    
    strcat(copy, argv[1]);
    strcat(copy, TMP_FILE_SUFFIX);
    
    outfp = fopen(copy, "wb");
    
    // copy header to outfp    
    for (i = 0; i < HEADER_SIZE; fputc(std_header[i++], outfp));
    
    // start chunk processecing
    while (!last_chunk) {
        // size ////////////////////////////////////////////////////////////////
        i = read_chunk(chunk_size, CHUNK_PART_SIZE, infp);
        if (i != CHUNK_PART_SIZE) {
            chunk_error();
        }
        cdata_size = swap_endianness(*((int*)chunk_size));
        
        // type ////////////////////////////////////////////////////////////////
        i = read_chunk(chunk_type, CHUNK_PART_SIZE, infp);
        if (i != CHUNK_PART_SIZE) {
            chunk_error();
        }
        if (skip_chunk(chunk_type)) {
            // need to skip cdata_size+4 (for the CRC) to reach the next chunk.
            fseek(infp, cdata_size+4, SEEK_CUR);
            continue;
        }
        write_chunk(chunk_size, CHUNK_PART_SIZE, outfp);
        write_chunk(chunk_type, CHUNK_PART_SIZE, outfp);
        
        // data ////////////////////////////////////////////////////////////////
        
        while (cdata_size > 0) {
            j = min(BUFFER_SIZE, cdata_size);
            cdata_size -= j;
            i = read_chunk(chunk_data, j, infp);
            if (i != j) {
                chunk_error();
            }
            write_chunk(chunk_data, j, outfp);
        }
        
        // crc /////////////////////////////////////////////////////////////////
        i = read_chunk(chunk_crc, CHUNK_PART_SIZE, infp);
        if (i != CHUNK_PART_SIZE) {
            chunk_error();
        }
        write_chunk(chunk_crc, CHUNK_PART_SIZE, outfp);
    }
    
    // all done with in files
    fclose(infp);
    fclose(outfp);
    
    // delete input file, and rename output file to input file's name
    if (remove(argv[1])) { // non-zero on failure
        printf("Could not delete input file. Output written to %s\n", copy);
    } else if (rename(copy, argv[1])) { // non-zero on failure
        printf("Could not rename output file, Output written to %s\n", copy);
    }
    
    exit(EXIT_SUCCESS);
}

int read_chunk(char* s, int n, FILE* fp) {
    char c;
    int i = n;
    
    while (n > 0 && (c = fgetc(fp)) != EOF) {
        *s++=c;
        n--;
    }
    
    return i-n;
}

void chunk_error() {
    printf("Could not read expected data amount");
    fclose(infp);
    fclose(outfp);
    exit(EXIT_FAILURE);    
}

void write_chunk(char* s, int n, FILE* fp) {
    while (n-- > 0) {
        fputc(*s++, fp);
    }
}

int skip_chunk(char* s) {
    int i = swap_endianness(*(int*)s);

    switch (i) {
        case CHUNK_HDR_GAMA : 
            return 1;
        case CHUNK_HDR_CHRM : 
            return 1;
        case CHUNK_HDR_SRGB : 
            return 1;
        case CHUNK_HDR_IEND : 
            last_chunk = 1; /* fall through */
        default : 
            return 0;
    }
}

// assumes int is 4 bytes long! does not take sign into consideration.
// also, since we're reading from PNG (network order, big endian), we
// assume that the order will only need switching if platform is 
// little-endian such as windows. we return unmodified otherwise.
int swap_endianness(int i) {
    unsigned char c1, c2, c3, c4;

    if (is_bigendian()) {
        return i;
    } else {
        c1 = i & 255;
        c2 = (i >> 8) & 255;
        c3 = (i >> 16) & 255;
        c4 = (i >> 24) & 255;

        return ((int)c1 << 24) + ((int)c2 << 16) + ((int)c3 << 8) + c4;
    }
}