#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <zlib.h>

uint32_t endianswap32( uint32_t n )
{
	// if we're in a big endian system, we don't need this
	uint16_t testme = 0x1234;
	if ( *(uint8_t*)(&testme) == 0x12 ) return n;
	uint32_t on;
	for ( int i=0; i<4; i++ )
		*(((uint8_t*)(&on))+i) = *(((uint8_t*)(&n))+(3-i));
	return on;
}

uint16_t endianswap16( uint16_t n )
{
	// if we're in a big endian system, we don't need this
	uint16_t testme = 0x1234;
	if ( *(uint8_t*)(&testme) == 0x12 ) return n;
	uint16_t on;
	for ( int i=0; i<2; i++ )
		*(((uint8_t*)(&on))+i) = *(((uint8_t*)(&n))+(1-i));
	return on;
}

typedef struct
{
	uint8_t highbitcheck;
	char sig[3];
	char crlf[2];
	char dosstop;
	char lf;
} pnghead_t;

typedef struct
{
	uint32_t length;
	char type[4];
	uint8_t *data;
	uint32_t crc;
} pngchunk_t;

typedef struct
{
	uint32_t width, height;
	uint8_t depth, type, compression, filter, interlace;
} ihdr_t;

typedef struct
{
	uint32_t xppu, yppu;
	uint8_t unit;
} phys_t;

int verbose = 0;
int exitval = 0;
int bailout = 0;
FILE *f = 0;
int nchunk = 0;
pnghead_t hdr = {0,{0},{0},0,0};
pngchunk_t chk = {0,{0},0,0};	// make sure this is null at the very start
uint8_t coldepth = 0;
uint8_t coltype = 0;
uint32_t curcrc = 0;
int validcrc = 0;
uint8_t *tdat = 0;	// accumulated data for all IDAT chunks
uint32_t tdatsiz = 0;	// size of it

typedef struct
{
	char name[4];
	void(*func)(pngchunk_t *chk);
} pngfunc_t;

void dumpData( uint8_t *dat, uint32_t len )
{
	for ( uint32_t i=0; i<len; i++ )
	{
		printf(" %02x",dat[i]);
		if ( !((i+1)&15) || ((i+1) == len) ) putchar('\n');
	}
}

void dumpChunk( pngchunk_t *chk )
{
	printf(" !!parser for this chunk not implemented, hex dump follows\n");
	dumpData(chk->data,chk->length);
}

const char* getcoltype( uint8_t type )
{
	switch ( type )
	{
	case 0:
		return "Grayscale";
	case 2:
		return "RGB";
	case 3:
		return "Indexed";
	case 4:
		return "Grayscale Alpha";
	case 6:
		return "RGB Alpha";
	}
	return "Invalid";
}

const char* getcompression( uint8_t compression )
{
	if ( !compression )
		return "Deflate";
	return "Invalid";
}

const char* getfilter( uint8_t filter )
{
	if ( !filter )
		return "Adaptive";
	return "Invalid";
}

const char* getinterlace( uint8_t interlace )
{
	switch ( interlace )
	{
	case 0:
		return "None";
	case 1:
		return "Adam7";
	}
	return "Invalid";
}

void readIHDR( pngchunk_t *chk )
{
	if ( chk->length != 13 )
	{
		printf(" !!IHDR size is not 13, bailing out\n");
		bailout = 1;
		return;
	}
	ihdr_t* ihdr = (ihdr_t*)(chk->data);
	printf(" width: %u\n height: %u\n depth: %u\n type: %u (%s)\n"
		" compression: %u (%s)\n filter: %u (%s)\n interlace: %u (%s)\n",
		endianswap32(ihdr->width),endianswap32(ihdr->height),
		ihdr->depth,ihdr->type,getcoltype(ihdr->type),
		ihdr->compression,getcompression(ihdr->compression),
		ihdr->filter,getfilter(ihdr->filter),
		ihdr->interlace,getinterlace(ihdr->interlace));
	coldepth = ihdr->type;
	coltype = ihdr->type;
}

void readPLTE( pngchunk_t *chk )
{
	if ( chk->length%3 )
	{
		printf(" !!palette length not divisible by 3, may be corrupted, skipping\n");
		return;
	}
	unsigned ncol = chk->length/3;
	for ( unsigned i=0; i<ncol; i++ )
		printf(" %d: (%hhu,%hhu,%hhu)\n",i,chk->data[i*3],chk->data[i*3+1],chk->data[i*3+2]);
}

void readIDAT( pngchunk_t *chk )
{
	if ( tdat ) tdat = realloc(tdat,tdatsiz+chk->length);
	else tdat = malloc(chk->length);
	memcpy(tdat+tdatsiz,chk->data,chk->length);
	tdatsiz += chk->length;
}

void readtRNS( pngchunk_t *chk )
{
	for ( unsigned i=0; i<chk->length; i++ )
		printf(" %d: %hhu\n",i,chk->data[i]);
}

void readgAMA( pngchunk_t *chk )
{
	if ( chk->length != 4 )
	{
		printf(" !!gAMA size is not 4, skipping\n");
		return;
	}
	uint32_t rg = endianswap32(*(uint32_t*)(chk->data));
	printf(" 1/%g\n",1./(rg/100000.));
}

void readcHRM( pngchunk_t *chk )
{
	if ( chk->length != 32 )
	{
		printf(" !!cHRM size is not 64, skipping\n");
		return;
	}
	uint32_t rwx = endianswap32(*(uint32_t*)(chk->data)),
		rwy = endianswap32(*(uint32_t*)(chk->data+4)),
		rrx = endianswap32(*(uint32_t*)(chk->data+8)),
		rry = endianswap32(*(uint32_t*)(chk->data+12)),
		rgx = endianswap32(*(uint32_t*)(chk->data+16)),
		rgy = endianswap32(*(uint32_t*)(chk->data+20)),
		rbx = endianswap32(*(uint32_t*)(chk->data+24)),
		rby = endianswap32(*(uint32_t*)(chk->data+28));
	printf(" White Point x: %g\n White Point y: %g\n"
		" Red x: %g\n Red y: %g\n"
		" Green x: %g\n Green y: %g\n"
		" Blue x: %g\n Blue y: %g\n",
		rwx/100000.,rwy/100000.,
		rrx/100000.,rry/100000.,
		rgx/100000.,rgy/100000.,
		rbx/100000.,rby/100000.);
}

void readsRGB( pngchunk_t *chk )
{
	if ( chk->length != 1 )
	{
		printf(" !!sRGB size is not 1, skipping\n");
		return;
	}
	switch ( chk->data[0] )
	{
	case 0:
		printf(" Perceptual\n");
		break;
	case 1:
		printf(" Relative colorimetric\n");
		break;
	case 2:
		printf(" Saturation\n");
		break;
	case 3:
		printf(" Absolute colorimetric\n");
		break;
	default:
		printf(" Unknown\n");
		break;
	}
}

void readiCCP( pngchunk_t *chk )
{
	// TODO process this
	dumpChunk(chk);
}

void readtEXt( pngchunk_t *chk )
{
	char *key = (char*)chk->data;
	char *val = key+strlen(key)+1;
	uint32_t vallen = chk->length-(val-key);
	printf(" '%s' :: '%.*s'\n",key,vallen,val);
}

void readzTXt( pngchunk_t *chk )
{
	// TODO decompress this stuff
	dumpChunk(chk);
}

void readiTXt( pngchunk_t *chk )
{
	char *key = (char*)chk->data;
	uint8_t cmpflag = chk->data[strlen(key)+1];
	if ( cmpflag )
	{
		printf(" !!compressed iTXt not supported, skipping");
		return;
	}
	char *lang = (char*)(chk->data+strlen(key)+3);
	char *kword = (char*)(chk->data+strlen(key)+3+strlen(lang)+1);
	char *text = (char*)(chk->data+strlen(key)+3+strlen(lang)+1+strlen(kword)+1);
	uint32_t textlen = chk->length-(text-key);
	printf(" '%s' :: ",key);
	if ( lang[0] ) printf("[%s] ",lang);
	if ( kword[0] ) printf("'%s' ",kword);
	if ( lang[0] || kword[0] ) printf(":: ");
	printf("'%.*s'\n",textlen,text);
}

void readbKGD( pngchunk_t *chk )
{
	switch ( coltype )
	{
	case 3:
		if ( chk->length != 1 )
		{
			printf(" !!bKGD size is not 1, skipping\n");
			return;
		}
		printf(" BG: %hhu\n",chk->data[0]);
		break;
	case 0:
	case 4:
		if ( chk->length != 2 )
		{
			printf(" !!bKGD size is not 2, skipping\n");
			return;
		}
		printf(" BG: %hu\n",endianswap16(*(uint16_t*)(chk->data)));
		break;
	case 2:
	case 6:
		if ( chk->length != 6 )
		{
			printf(" !!bKGD size is not 6, skipping\n");
			return;
		}
		printf(" BG: (%hu,%hu,%hu)\n",
			endianswap16(*(uint16_t*)(chk->data)),
			endianswap16(*(uint16_t*)(chk->data+2)),
			endianswap16(*(uint16_t*)(chk->data+4)));
		break;
	default:
		printf(" !!bKGD present for unhandled color type %d, skipping\n",coltype);
		break;
	}
}

void readpHYs( pngchunk_t *chk )
{
	if ( chk->length != 9 )
	{
		printf(" !!pHYs size is not 9, skipping\n");
		return;
	}
	phys_t* phys = (phys_t*)(chk->data);
	printf(" xppu: %u\n yppu: %u\n unit: %u\n",endianswap32(phys->xppu),endianswap32(phys->yppu),phys->unit);
}

void readsBIT( pngchunk_t *chk )
{
	switch ( coltype )
	{
	case 0:
		if ( chk->length != 1 )
		{
			printf(" !!sBIT size is not 1, skipping\n");
			return;
		}
		printf(" %hhu\n",chk->data[0]);
		break;
	case 2:
	case 3:
		if ( chk->length != 3 )
		{
			printf(" !!sBIT size is not 3, skipping\n");
			return;
		}
		printf(" (%hhu,%hhu,%hhu)\n",chk->data[0],chk->data[1],chk->data[2]);
		break;
	case 4:
		if ( chk->length != 2 )
		{
			printf(" !!sBIT size is not 2, skipping\n");
			return;
		}
		printf(" (%hhu,%hhu)\n",chk->data[0],chk->data[1]);
		break;
	case 6:
		if ( chk->length != 4 )
		{
			printf(" !!sBIT size is not 4, skipping\n");
			return;
		}
		printf(" (%hhu,%hhu,%hhu,%hhu)\n",chk->data[0],chk->data[1],chk->data[2],chk->data[3]);
		break;
	default:
		printf(" !!sBIT present for unhandled color type %d, skipping\n",coltype);
		break;
	}
}

void readsPLT( pngchunk_t *chk )
{
	char *name = (char*)chk->data;
	uint8_t depth = chk->data[strlen(name)+1];
	printf(" suggested palette: %s\n sample depth: %hhu\n",name,depth);
	int div;
	if ( depth == 16 )
		div = 10;
	else if ( depth == 8 )
		div = 6;
	else
	{
		printf(" !!invalid depth of %u, skipping\n",depth);
		return;
	}
	uint32_t psiz = (chk->length-(strlen(name)+2));
	if ( psiz%div )
	{
		printf(" !!palette length not divisible by %d, skipping\n",div);
		return;
	}
	unsigned ncol = psiz/div;
	uint8_t *pdat = chk->data+(strlen(name)+2);
	for ( unsigned i=0; i<ncol; i++ )
	{
		if ( div == 10 )
			printf(" %u: (%hu,%hu,%hu) [%hu]\n",i,
				endianswap16(*(uint16_t*)(pdat+i*10)),
				endianswap16(*(uint16_t*)(pdat+i*10+2)),
				endianswap16(*(uint16_t*)(pdat+i*10+4)),
				endianswap16(*(uint16_t*)(pdat+i*10+6)));
		else if ( div == 6 )
			printf(" %u: (%hhu,%hhu,%hhu) [%hu]\n",i,
				*(uint8_t*)(pdat+i*6),
				*(uint8_t*)(pdat+i*6+1),
				*(uint8_t*)(pdat+i*6+2),
				endianswap16(*(uint16_t*)(pdat+i*6+3)));
	}
}

void readhIST( pngchunk_t *chk )
{
	if ( chk->length%2 )
	{
		printf(" !!hIST size not divisible by two, skipping\n");
		return;
	}
	unsigned ncol = chk->length/2;
	for ( unsigned i=0; i<ncol; i++ )
		printf(" %u: %hu\n",i,endianswap16(*(uint16_t*)(chk->data+i*2)));
}

void readtIME( pngchunk_t *chk )
{
	if ( chk->length != 7 )
	{
		printf(" !!tIME size is not 7, skipping\n");
		return;
	}
	printf(" %hu-%hhu-%hhu %hhu:%hhu:%hhu\n",
		endianswap16(*(uint16_t*)(chk->data)),chk->data[2],chk->data[3],
		chk->data[4],chk->data[5],chk->data[6]);
}

void readgrAb( pngchunk_t *chk )
{
	if ( chk->length != 8 )
	{
		printf(" !!grAb size is not 8, skipping\n");
		return;
	}
	uint32_t rx = endianswap32(*(uint32_t*)(chk->data)), ry = endianswap32(*(uint32_t*)(chk->data+4));
	int32_t x = *(int32_t*)(&rx), y = *(int32_t*)(&ry);
	printf(" %+d,%+d\n",x,y);
}

#define NUMREADERS 20

pngfunc_t readers[NUMREADERS] =
{
	// critical
	{"IHDR",readIHDR},
	{"PLTE",readPLTE},
	{"IDAT",readIDAT},
	{"IEND",0},
	// transparency
	{"tRNS",readtRNS},
	// colorspace
	{"gAMA",readgAMA},
	{"cHRM",readcHRM},
	{"sRGB",readsRGB},
	{"iCCP",readiCCP},
	// text
	{"tEXt",readtEXt},
	{"zTXt",readzTXt},
	{"iTXt",readiTXt},
	// misc
	{"bKGD",readbKGD},
	{"pHYs",readpHYs},
	{"sBIT",readsBIT},
	{"sPLT",readsPLT},
	{"hIST",readhIST},
	{"tIME",readtIME},
	// zdoom
	{"grAb",readgrAb},
	{"alPh",0}
};

void readchunkdata( pngchunk_t* chk )
{
	for ( int i=0; i<NUMREADERS; i++ )
	{
		if ( strncmp(chk->type,readers[i].name,4) ) continue;
		if ( readers[i].func ) readers[i].func(chk);
		return;
	}
	printf(" !!no reader for this chunk type\n");
}

// TODO process the whole thing (might not bother actually)
void processIDAT( void )
{
	if ( !verbose ) return;
	printf("full datastream dump follows:\n");
	dumpData(tdat,tdatsiz);
}

int main( int argc, char **argv )
{
	if ( (argc > 1) && (!strncmp(argv[1],"-v",2) || !strncmp(argv[1],"--verbose",9)) )
		verbose = 1;
	if ( argc < (2+verbose) )
	{
		printf("No file supplied.\n");
		exitval = 0;
		goto endmii;
	}
	f = fopen(argv[1+verbose],"rb");
	if ( !f )
	{
		printf("Failed to open file.\n");
		exitval = 1;
		goto endmii;
	}
	fread(&hdr,1,sizeof(hdr),f);
	if ( hdr.highbitcheck != 0x89 )
	{
		printf("High bit check failed.\n");
		exitval = 2;
		goto endmii;
	}
	if ( strncmp(hdr.sig,"PNG",3) )
	{
		printf("PNG signature check failed.\n");
		exitval = 4;
		goto endmii;
	}
	if ( strncmp(hdr.crlf,"\r\n",2) )
	{
		printf("CRLF check failed.\n");
		exitval = 8;
		goto endmii;
	}
	if ( hdr.dosstop != '\032' )
	{
		printf("EOF check failed.\n");
		exitval = 16;
		goto endmii;
	}
	if ( hdr.lf != '\n' )
	{
		printf("LF check failed.\n");
		exitval = 32;
		goto endmii;
	}
readchunk:
	fread(&chk,1,8,f);	// read size and type first
	chk.length = endianswap32(chk.length);	// swap from BE
	if ( (nchunk == 0) && strncmp(chk.type,"IHDR",4) )
	{
		printf("First chunk is not IHDR.\n");
		exitval = 64;
		goto endmii;
	}
	// (re)allocate data
	if ( chk.data ) chk.data = realloc(chk.data,chk.length);
	else chk.data = malloc(chk.length);
	fread(chk.data,1,chk.length,f);	// read data
	fread(&(chk.crc),1,4,f);	// read CRC
	chk.crc = endianswap32(chk.crc);
	curcrc = crc32(0,(uint8_t*)chk.type,4);
	if ( chk.length ) curcrc = crc32(curcrc,chk.data,chk.length);
	validcrc = (curcrc == chk.crc);
	printf("%.4s chunk of length %u (CRC32: %08X, %s).\n",chk.type,chk.length,chk.crc,validcrc?"OK":"FAILED");
	if ( !validcrc ) printf("  CRC32 ACTUAL: %08X\n",curcrc);
	else readchunkdata(&chk);
	nchunk++;
	if ( bailout ) goto endmii;
	if ( strncmp(chk.type,"IEND",4) ) goto readchunk;
	if ( tdat ) processIDAT();
	long cpos = ftell(f);
	fseek(f,0,SEEK_END);
	long epos = ftell(f);
	if ( epos > cpos )
	{
		printf(" !!there are %lu bytes of extra data after IEND\n",epos-cpos);
		tdatsiz = epos-cpos;
		if ( tdat ) tdat = realloc(tdat,tdatsiz);
		else tdat = malloc(tdatsiz);
		fseek(f,cpos,SEEK_SET);
		fread(tdat,1,tdatsiz,f);
		if ( verbose ) dumpData(tdat,tdatsiz);
	}
endmii:
	if ( f ) fclose(f);
	if ( chk.data ) free(chk.data);
	if ( tdat ) free(tdat);
	return exitval;
}