#include <stdlib.h>
#include <stdio.h>
#include <png.h>

void save_png(const char *filename,
		unsigned char *pixels, int width, int height, int bpp,
		unsigned char *palette) {
	FILE *file;
	
	file = fopen(filename, "wb");
	if (!file) {
		return;
	}

	png_struct *png_write;
	png_info *png_info;

	png_write = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	png_info = png_create_info_struct(png_write);

	png_init_io(png_write, file);

	//png_set_compression_level(png_write, Z_BEST_COMPRESSION);
	png_set_IHDR(png_write, png_info, width, height,
	        bpp==4?4:8,
		bpp<=8?PNG_COLOR_TYPE_PALETTE:PNG_COLOR_TYPE_RGB_ALPHA,
		PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT,
		PNG_FILTER_TYPE_DEFAULT);
	
	if (bpp == 4) {
		png_color png_pal[16];
		for (int i = 0; i < 16; ++i) {
			png_pal[i].blue = *palette++;
			png_pal[i].green = *palette++;
			png_pal[i].red = *palette++;
			*palette++;
		}
		
		png_set_PLTE(png_write, png_info, png_pal, 16);
	} else if (bpp == 8) {
		png_color png_pal[256];
		for (int i = 0; i < 256; ++i) {
			png_pal[i].blue = *palette++;
			png_pal[i].green = *palette++;
			png_pal[i].red = *palette++;
			*palette++;
		}
		
		png_set_PLTE(png_write, png_info, png_pal, 256);
	} else {
		png_set_bgr(png_write);
	}

	png_byte *row_pointers[height];
	unsigned char *p;
	int orig_bpp = bpp;
	bpp /= 8;
	if (bpp == 0) {
	        bpp = 1;
	}
	
	p = pixels;
	for (int i = 0; i < height; ++i) {
		row_pointers[i] = (png_byte *)p;
		
		if (orig_bpp == 4) {
                        unsigned char *s = p;
                        unsigned char *b = p;
                        for (int j = 0; j < width; j += 2) {
                                unsigned char p1 = *s++;
                                unsigned char p2 = *s++;
                                
                                *b++ = (p2&0xf)|(p1<<4);
                        }
		}
		p += width*bpp;
	}
	
	png_write_info(png_write, png_info);
	png_write_image(png_write, row_pointers);
	png_write_end(png_write, NULL);
	png_destroy_write_struct(&png_write, &png_info);

	fclose(file);
}

struct bitstream_t {
        unsigned char *ptr;
        int bit;
        unsigned char *end;
};

int read_bit(bitstream_t *bs) {
        if (bs->ptr >= bs->end) {
                return 0;
        }
        
        int bit = *bs->ptr&(1<<(7-bs->bit))?1:0;
        
        bs->bit++;
        if (bs->bit == 8) {
                bs->ptr++;
                bs->bit = 0;
                
                if (bs->ptr == bs->end) {
                    printf("yar, something's wrong\n");
                }
        }
        
        return bit;
}

int read_bits(bitstream_t *bs, int count) {
        int output = 0;
        for (int i = 0; i < count; ++i) {
                output <<= 1;
                output |= read_bit(bs);
        }
        
        return output;
}

unsigned char colortable[16*16];
void init_colortable() {
        unsigned char *c = colortable;
        unsigned char v = 0;
        
        for (int i = 0; i < 16; ++i) {
                for (int j = 0; j < 16; ++j) {
                        *c++ = v & 0xf;
                        v--;
                }
                v++;
        }
}

int read_color_bits(bitstream_t *bs) {
        if (read_bit(bs)) {
                return read_bit(bs);
        }
        
        if (!read_bit(bs)) {
                return 2 + read_bit(bs);
        }
        
        if (!read_bit(bs)) {
                return 4 + read_bits(bs, 2);
        }
        
        return 8 + read_bits(bs, 3);
}

unsigned short read_one_color(bitstream_t *bs, unsigned short color) {
        color <<= 8;
        
        unsigned char new_color = read_color_bits(bs);
        // new value indexes the color table
        int coltable = ((color>>4)&0xf0) | (new_color&0xf);
        unsigned char real_color = colortable[coltable];
        
        color |= real_color;
        
        // if index > 0, shift indexed value to front of list
        if (new_color > 0) {
                unsigned char *src = &colortable[coltable-1];
                unsigned char *dest = &colortable[coltable];
        
                for (int i = 0; i < new_color; ++i) {
                        *dest-- = *src--;
                }
        
                *dest = real_color;
        }
        
        return color;
}

unsigned short read_colors(bitstream_t *bs, unsigned short color) {
        // read two colors, swap
        color = read_one_color(bs, color);
        color = read_one_color(bs, color);
        
        color = (color>>8) | (color<<8);
        
        return color;
}

void run_file(const char *filename) {
        FILE *file = fopen(filename, "rb");
        if (!file) {
                printf("can't open\n");
                return;
        }
        
        struct header_start_t {
                unsigned char id[8];
                unsigned char id_2[4];
                unsigned char extra_bytes[2];
                
                //unsigned char id_2[6];
                
                //unsigned char width_hi;
                //unsigned char width_lo;
                //unsigned char height_hi;
                //unsigned char height_lo;
        } header_start;
        
        struct header_t {
                unsigned char width_hi;
                unsigned char width_lo;
                unsigned char height_hi;
                unsigned char height_lo;
        } header;
        
        fread(&header_start, sizeof(header_start), 1, file);
        fseek(file, (header_start.extra_bytes[0]<<8)|header_start.extra_bytes[1], SEEK_CUR);
        fread(&header, sizeof(header), 1, file);
        
        // initialize palette
        unsigned char palette[256*4];
        unsigned char *p = palette;
        memset(palette, 0, sizeof(palette));
        
        for (int i = 0; i< 16; ++i) {
                fread(p, 3, 1, file);
                
                // swap r and b
                unsigned char a = p[0];
                p[0] = p[2];
                p[2] = a;
                
                p += 4;
        }
        
        unsigned int cur = ftell(file);
        fseek(file, 0, SEEK_END);
        unsigned int len = ftell(file) - cur;
        fseek(file, cur, SEEK_SET);
        
        // initialize bitstream
        unsigned char *input = new unsigned char[len];
        fread(input, len, 1, file);
        
        fclose(file);
        
        bitstream_t bs;
        bs.ptr = input;
        bs.bit = 0;
        bs.end = input + len;
        
        // initialize buffer
        unsigned int width = (header.width_hi << 8) | header.width_lo;
        unsigned int height = (header.height_hi << 8) | header.height_lo;
        
        // - two line buffer at the start filled with an initial color, for
        // referencing purposes
        unsigned char *buf = new unsigned char[(width*(height+2)) + 64];
        memset(buf, 0, width*(height+2));
        
        unsigned char *b = buf;
        unsigned char *b_end = buf + (width*(height+2));
        
        
        // ALGORITHM STARTS HERE
        
        // initialize starting colors for two line buffer
        init_colortable();
        
        unsigned short colors = 0;
        
        colors = read_colors(&bs, colors);
        
        // first color fills two lines
        for (unsigned int i = 0; i < width; ++i) {
                *(unsigned short *)b = colors;
                b += 2;
        }
        
        int last = 0;
        bool done = 0;
        
        while (!done) {
                int cursor = -width;
                
                // read place
                int place;
                if (read_bit(&bs)) {		// place 1x
                        if (read_bit(&bs)) {	// place 11
                                place = 0x110 | read_bit(&bs);
                        } else {
                                place = 0x010;
                        }
                } else if (read_bit(&bs)) {	// place 01
                        place = 0x001;
                } else {
                        place = 0x000;
                }
                
                switch(place) {
                case 0x000:
                        // cursor at -4, unless last two colors match,
                        // then cursor at -2
                        cursor = -4;
                        colors = *(unsigned short *)(b-2);
                        
                        if ((colors&0xff) == (colors>>8)) {
                                if (cursor != last) {
                                        cursor = -2;
                                }
                        }
                        break;
                case 0x001:
                        // one line up
                        break;
                case 0x010:
                        // two lines up
                        cursor += cursor;
                        break;
                case 0x110:
                        // one line up, right
                        cursor += 1;
                        break;
                case 0x111:
                        // one line up, left
                        cursor -= 1;
                        break;
                }
                
                // do whatever function we're supposed to do
                if (last != cursor) {
                        last = cursor;
                        if (!read_bit(&bs)) {
                                // copy one pixel pair from cursor
                                *(unsigned short *)b = *(unsigned short *)(b+cursor);
                                b += 2;
                                if (b >= b_end) {
                                        done = 1;
                                        break;
                                }
                        } else {
                                // read number of bits that we're going
                                // to be reading in
                                int pcount = 1;
                                int count = 1;
                                while (read_bit(&bs)) {
                                        count++;
                                }
                                
                                // read the count one bit at a time, note
                                // that the first 1 is implied
                                while (count-- > 0) {
                                        pcount <<= 1;
                                        pcount |= read_bit(&bs);
                                }
                                
                                // copy
                                while (pcount > 0) {
                                        *(unsigned short*)b = *(unsigned short *)(b+cursor);
                                        b+=2;
                                        
                                        if (b >= b_end) {
                                                done = 1;
                                                break;
                                        }
                                        
                                        --pcount;
                                }
                        }
                } else {
                        // 1 read color from b-1 into al
                        colors = b[-1];
                        do {
                                // 2 call read_color with color value
                                colors = read_colors(&bs, colors);
                                // 3 store color at b
                                *(unsigned short *)b = colors;
                                b+=2;
                                
                                // 4 al = ah
                                colors = (colors&0xff00)|(colors>>8);
                                
                                // if at block end, stop
                                if (b >= b_end) {
                                        done = 1;
                                        break;
                                }
                        
                                // get bit, if 1 then repeat last chunk starting at pos 2
                        } while (read_bit(&bs));
                        
                        // else reset cursor go to next loop
                        last = 0;
                }
        }
        
        // ALGORITHM ENDS HERE
        
        delete[] input;
        
        char out_filename[1024];
        sprintf(out_filename, "%s.png", filename);
        
        save_png(out_filename, buf+(width*2), width, height, 4, palette);
        
        delete[] buf;
        
        printf("OK %d x %d\n", width, height);
}

int main(int argc, char **argv) {
        int i;
        
        for (i = 1; i < argc; ++i) {
                printf("%s: ",argv[i]);
                run_file(argv[i]);
        }
        
        return 0;
}