#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <png.h>
#include <stdbool.h>
#include <string.h>

typedef struct {
    int index;
    int distance;
} color_match_t;


#define ATARI_COLORS 256

const int palletePreferedLimitDefault = 8;

uint32_t atari_palette_pal[ATARI_COLORS] = 
{
	0x000000,0x111111,0x222222,0x333333,0x444444,0x555555,0x666666,0x777777,0x888888,0x999999,0xaaaaaa,0xbbbbbb,0xcccccc,0xdddddd,0xeeeeee,0xffffff,0x091900,0x192806,0x29370d,0x3a4714,0x4a561b,0x5a6522,0x6b7529,0x7b8430,0x8c9336,0x9ca33d,0xacb244,0xbdc14b,0xcdd152,0xdee059,0xeeef60,0xffff67,0x300000,0x3d1108,0x4b2211,0x593319,0x674422,0x75552a,0x826633,0x90773b,0x9e8844,0xac994c,0xbaaa55,0xc7bb5d,0xd5cc66,0xe3dd6e,0xf1ee77,0xffff80,0x4b0000,0x570f0c,0x631e18,0x6f2e24,0x7a3d30,0x874d3c,0x935c49,0x9f6b55,0xab7b61,0xb68a6d,0xc39a79,0xcfa986,0xdbb892,0xe6c89e,0xf3d7aa,0xffe7b7,0x550000,0x600e10,0x6b1c21,0x772a32,0x823843,0x8d4654,0x995465,0xa46276,0xaf7187,0xbb7f98,0xc68da9,0xd19bba,0xdda9cb,0xe8b7dc,0xf3c5ed,0xffd4fe,0x4c0047,0x570d53,0x631b5f,0x6f286b,0x7b3678,0x874384,0x935190,0x9f5e9c,0xab6ca9,0xb779b5,0xc387c1,0xcf94cd,0xdba2da,0xe7afe6,0xf3bdf2,0xffcbff,0x30007e,0x3b0b85,0x49198d,0x572796,0x65349f,0x7242a7,0x8050b0,0x8e5db8,0x9c6bc1,0xa979c9,0xb786d2,0xc594db,0xd3a2e3,0xe0afec,0xeebdf4,0xfccbfd,0x0a0097,0x1a0e9d,0x2a1da4,0x3b2cab,0x4b3ab2,0x5b49b9,0x6c58c0,0x7c67c7,0x8c75ce,0x9c84d5,0xad93dc,0xbda2e3,0xceb0ea,0xdebff1,0xeecef8,0xffddff,0x00008e,0x0c0d94,0x1b1e9c,0x2a2ea3,0x393eab,0x484eb2,0x575eba,0x666ec1,0x747ec9,0x838fd0,0x929fd8,0xa1afdf,0xb0bfe6,0xbfcfee,0xcedff5,0xddeffd,0x000e64,0x0c1e6e,0x192e78,0x263e83,0x324e8d,0x3f5e97,0x4c6ea2,0x587eac,0x658eb6,0x729ec1,0x7eaecb,0x8bbed5,0x98cee0,0xa4deea,0xb1eef4,0xbeffff,0x002422,0x09302e,0x153f3d,0x204d4c,0x2c5c5a,0x376a69,0x427978,0x4e8786,0x599695,0x65a4a4,0x70b3b2,0x7cc1c1,0x87d0d0,0x92dfde,0x9eeded,0xa9fcfc,0x003200,0x0b3f0e,0x164d1c,0x225b2b,0x2d6839,0x397648,0x448456,0x509164,0x5b9f73,0x67ad81,0x72ba90,0x7ec89e,0x89d6ac,0x95e3bb,0xa0f1c9,0xacffd8,0x003400,0x0c410a,0x194f14,0x265c1e,0x336a28,0x407732,0x4c853c,0x599246,0x66a050,0x73ad5a,0x80bb64,0x8cc86e,0x99d678,0xa6e382,0xb3f18c,0xc0ff97,0x002a00,0x0f3807,0x1e460e,0x2d5416,0x3c621d,0x4b7124,0x5a7f2c,0x698d33,0x799b3b,0x88a942,0x97b849,0xa6c651,0xb5d458,0xc4e260,0xd3f067,0xe3ff6f,0x0d1700,0x1d2606,0x2d350d,0x3d4514,0x4d541b,0x5d6422,0x6d7329,0x7d8330,0x8e9237,0x9ea23e,0xaeb145,0xbec14c,0xced053,0xdee05a,0xeeef61,0xffff68,0x330000,0x401008,0x4e2111,0x5b321a,0x694323,0x77542c,0x846535,0x92763e,0x9f8646,0xad974f,0xbba858,0xc8b961,0xd6ca6a,0xe3db73,0xf1ec7c,0xfffd85
};

/**
 * @brief Calculates the squared Euclidean distance between two RGB colors.
 *
 * This function is used to measure the "closeness" of two colors in RGB space
 * without computing the actual Euclidean distance (avoids sqrt for performance).
 *
 * @param r1 Red component of the first color.
 * @param g1 Green component of the first color.
 * @param b1 Blue component of the first color.
 * @param r2 Red component of the second color.
 * @param g2 Green component of the second color.
 * @param b2 Blue component of the second color.
 * @return Squared distance between the two colors.
 */

static inline int color_distance_sq(uint8_t r1, uint8_t g1, uint8_t b1,
                                    uint8_t r2, uint8_t g2, uint8_t b2) 
{
    int dr = (int)r1 - r2;
    int dg = (int)g1 - g2;
    int db = (int)b1 - b2;
    return dr * dr + dg * dg + db * db;
}

/**
 * @brief Converts an 8-bit value to a binary string representation.
 *
 * The result is a static string of 8 characters, representing the binary form of the input byte.
 *
 * @param in The input 8-bit value.
 * @return Pointer to a static 9-character null-terminated string.
 */

static inline char *toBinString(uint8_t in)
{
	static char out[9] = "01234567";
	for(int i=0;i<8;i++)
	{
		out[i] = ((in>>7)==0)?'0':'1';
		in=in<<1;
	}
	return out;
}

/**
 * @brief Compare color distances
 *
 * @param a
 * @param b
*/

static int compare_color_match(const void *a, const void *b) 
{
    color_match_t *ca = (color_match_t *)a;
    color_match_t *cb = (color_match_t *)b;
    return ca->distance - cb->distance;
}

/**
 * @brief Sorts all Atari palette indices based on distance to the input RGB color.
 *
 * Uses Euclidean color distance (squared) to rank the palette entries by closeness.
 * The result is stored in a provided array, sorted from best to worst match.
 *
 * @param r Red component of the input color.
 * @param g Green component of the input color.
 * @param b Blue component of the input color.
 * @param palette Pointer to a 256-entry palette of 32-bit Atari color values (0xRRGGBB).
 * @param sorted_indices Output array of size 256; will contain palette indices sorted by distance.
 */
 
void rgb_to_atari_index_sorted(uint8_t r, uint8_t g, uint8_t b, uint32_t *palette, int *sorted_indices)
{
	color_match_t matches[ATARI_COLORS];

	for (int i = 0; i < ATARI_COLORS; i++) 
	{
	uint32_t color = palette[i];
	uint8_t pr = (color >> 16) & 0xff;
	uint8_t pg = (color >> 8) & 0xff;
	uint8_t pb = color & 0xff;

	matches[i].index = i;
	matches[i].distance = color_distance_sq(r, g, b, pr, pg, pb);
	}

	qsort(matches, ATARI_COLORS, sizeof(color_match_t), compare_color_match);

	for (int i = 0; i < ATARI_COLORS; i++)
	{
		sorted_indices[i] = matches[i].index;
	}
}

typedef enum
{
	GF_LINEAR,
	GF_8x8,
	GF_4x8
	
} gfx_mode_t;

typedef enum
{
	CM_1BIT, 
	CM_2BIT,
	CM_4BIT
	
} color_mode_t;

typedef enum
{
	OUTFMT_ASM, 
	OUTFMT_RAW
	
} out_fmt_t;

/**
 * @brief Writes a single byte to a file in either raw or ASM format.
 *
 * In ASM mode, the byte is printed as a binary string with the `.byte` directive.
 * In RAW mode, it is written directly as binary data.
 *
 * @param byte The byte to write.
 * @param fp_outfile File pointer to the output file.
 * @param out_fmt Output format (raw or ASM).
 * @return true if write was successful, false otherwise.
 */

static bool write_byte(uint8_t byte,FILE *fp_outfile,out_fmt_t out_fmt)
{
	if(!fp_outfile)
	{
		return false;
	}
	
	if(out_fmt == OUTFMT_RAW)
	{
		fwrite(&byte, 1, 1, fp_outfile);
	}
	else if(out_fmt == OUTFMT_ASM)
	{
		fprintf(fp_outfile,"\t.byte %%%s\n",toBinString(byte)); // maybe with rows later
	}
	return true;
}

/**
 * @brief Handles buffered output of bytes with optional console printing.
 *
 * Increments an internal flush counter. When the counter reaches a threshold,
 * the current byte is written to the output file and optionally printed to the console.
 *
 * @param fullPrint If true, print binary string to stdout.
 * @param flushCountMax Threshold for flush counter.
 * @param flushCount Pointer to the flush counter (must be initialized).
 * @param fp_outfile Output file pointer.
 * @param out_fmt Output format (RAW or ASM).
 * @param byte The byte to potentially write.
 * @return true if the operation completed successfully, false otherwise.
 */

static bool bufferedWrite(bool fullPrint,int flushCountMax,int *flushCount,FILE *fp_outfile,out_fmt_t out_fmt,uint8_t byte)
{
	if(!fp_outfile || !flushCount)
	{
		return false;
	}
	
	(*flushCount)++;
	if(*flushCount==flushCountMax)
	{
		if(fullPrint)
		{
			printf("%s",toBinString(byte));
		}
		*flushCount = 0;
		
		write_byte(byte,fp_outfile,out_fmt);
	}
	
	return true;
}

/**
 * @brief Prints usage instructions for the converter tool.
 *
 * Displays command-line options and usage examples to the console.
 *
 * @param progname The name of the program (typically argv[0]).
 */

static void print_usage(const char *progname) 
{
	printf(
		"Usage: %s Options -i input.png -o output.asm\n"
		"\n"
		"Options:\n"
		"  -gm LINEAR|8x8|4x8  Convert mode Linear or 8x8/4x8 tiles (default LINEAR)\n"
		"  -cm 1BIT/2BIT/4BIT  Color mode 1,2 or 4 Bit (default 2 Bit)\n"
		"  -outfmt <asm|raw>   Output format (default: asm)\n"
		"  -full               Optional print all informations (default is off)\n"
		"  -palPref cnt        Optional count of prefered colors to display (default is %d)\n"
		"\n"
		"Examples:\n"
		"  %s -gm LINEAR -cm 1BIT -outfmt raw gfx.s gfx.raw\n"
		"  %s -gm 4x8 -cm 2BIT -outfmt raw gfx.s gfx.raw\n"
		"\n", progname, palletePreferedLimitDefault, progname,progname);
}


/**
 * @brief Entry point for the Atari graphics converter tool.
 *
 * This function parses command-line arguments to determine input/output files,
 * conversion modes (graphics layout, color depth, output format), and optionally prints
 * palette usage statistics. It loads a PNG image (palette-based), maps the colors to a
 * predefined Atari palette, and converts the image data into either linear or tile-based
 * formats with optional output as raw binary or assembler-compatible text.
 *
 * Supported options:
 * - `-i <input.png>`: Specifies the input PNG file (palette-based only).
 * - `-o <output>`: Specifies the output file (.asm or .raw).
 * - `-gm <LINEAR|8x8|4x8>`: Selects the graphics mode (default: LINEAR).
 * - `-cm <1BIT|2BIT>`: Selects the color mode depth (default: 2BIT).
 * - `-outfmt <asm|raw>`: Specifies the output format (default: asm).
 * - `-full`: Enables full binary printout to stdout for debugging/visual inspection.
 *
 * @param argc Argument count.
 * @param argv Argument vector.
 * @return Exit code: 0 on success, 1 on failure or invalid input.
 */

int main(int argc, char *argv[]) 
{
	/* defaults */
	
	gfx_mode_t 	gfx_mode = GF_LINEAR;
	color_mode_t	color_mode = CM_2BIT;
	out_fmt_t	out_fmt = OUTFMT_ASM;
	bool fullPrint = false;
	bool printUsage = false;
	int palletePreferedLimit = palletePreferedLimitDefault;
	
	char *input_file = NULL;
	char *output_file = NULL;
	
	/* parse arguments */
	
	for (int i = 1; i < argc; ++i) 
	{
		if (!strcmp(argv[i], "-i") && i + 1 < argc) 
		{
		    input_file = argv[++i];
		} 
		else if (!strcmp(argv[i], "-o") && i + 1 < argc) 
		{
		    output_file = argv[++i];
		}
		else if (!strcmp(argv[i], "-palPref") && i + 1 < argc) 
		{
		    palletePreferedLimit = atoi(argv[++i]);
		    palletePreferedLimit = palletePreferedLimit<1?1:palletePreferedLimit;
		    palletePreferedLimit = palletePreferedLimit>=ATARI_COLORS?(ATARI_COLORS-1):palletePreferedLimit;
		}		
		else if (!strcmp(argv[i], "-full")) 
		{
		    fullPrint = true;
		}
		else if (!strcmp(argv[i], "-outfmt") && i + 1 < argc) 
		{
		    const char *val = argv[++i];
		    if (!strcmp(val, "asm")) out_fmt = OUTFMT_ASM;
		    else if (!strcmp(val, "raw")) out_fmt = OUTFMT_RAW;
		    else 
		    {
			fprintf(stderr, "Unknown output format: %s\n", val);
			return 1;
		    }
		}
		else if (!strcmp(argv[i], "-cm") && i + 1 < argc) 
		{
		    const char *val = argv[++i];
		    if (!strcasecmp(val, "1BIT")) color_mode = CM_1BIT;
		    else if (!strcasecmp(val, "2BIT")) color_mode = CM_2BIT;
		    else if (!strcasecmp(val, "4BIT")) color_mode = CM_4BIT;
		    else 
		    {
			fprintf(stderr, "Unknown color format: %s\n", val);
			return 1;
		    }
		}
		else if (!strcmp(argv[i], "-gm") && i + 1 < argc) 
		{
		    const char *val = argv[++i];
		    if (!strcasecmp(val, "LINEAR")) gfx_mode = GF_LINEAR;
		    else if (!strcasecmp(val, "8x8")) gfx_mode = GF_8x8;
		    else if (!strcasecmp(val, "4x8")) gfx_mode = GF_4x8;
		    else 
		    {
			fprintf(stderr, "Unknown gfx format: %s\n", val);
			return 1;
		    }
		}
	}
	
	/* check */
	
	if(argc != 1)
	{
		if(!input_file)
		{
			fprintf(stderr, "Error: No input file\n");
			printUsage = true;
		}
		
		if(!output_file)
		{
			fprintf(stderr, "Error: No output file\n");
			printUsage = true;
		}
	}
	else
	{
		printUsage = true;
	}
	
	if(printUsage)
	{
		print_usage(argv[0]);
		return 0;
	}
	
	/* load */

	FILE *fp_infile = fopen(input_file, "rb");
	if (!fp_infile) 
	{
		perror("fopen PNG");
		return 1;
	}

	png_structp png = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	png_infop info = png_create_info_struct(png);
	if (setjmp(png_jmpbuf(png))) 
	{
		printf("Error loading PNG\n");
		fclose(fp_infile);
		return 1;
	}

	png_init_io(png, fp_infile);
	png_read_info(png, info);
	int width = png_get_image_width(png, info);
	int height = png_get_image_height(png, info);
	png_byte color_type = png_get_color_type(png, info);
	png_byte bit_depth = png_get_bit_depth(png, info);
	

	printf("; PNG info: \n; Width=%d Height=%d ColorType=%d BitDepth=%d\n",width,height,(int)color_type,(int)bit_depth);

	if(color_type != PNG_COLOR_TYPE_PALETTE && color_mode != CM_4BIT)
	{
		printf("Error only palette types supported\n");
		fclose(fp_infile);
		return 1;	
	}
	
	png_set_packing(png); // 1 pixel is 1 byte
	
	png_read_update_info(png, info);

	png_colorp palette = NULL; 
	int num_palette = 0;
	
	if(color_mode == CM_4BIT)
	{
		// no palette
	}
	else
	{
		if (png_get_PLTE(png, info, &palette, &num_palette)) 
		{
			int colorsPrefered[ATARI_COLORS];
			printf("; PNG Palette (%d colors):\n", num_palette);
			for (int i = 0; i <  num_palette; i++) 
			{
				int r = palette[i].red;
				int g = palette[i].green;
				int b = palette[i].blue;
				
				// calculate distances best to worst
				
				rgb_to_atari_index_sorted(r, g, b,atari_palette_pal,colorsPrefered);
				
				printf("; Color %3d: RGB $%2.2x%2.2x%2.2x ATPalette (best is first): ",i,r,g,b);
				for(int j=0;j<palletePreferedLimit;j++)
				{
					printf("$%2.2x%s",colorsPrefered[j],(j!=(palletePreferedLimit-1))?",":"");
				}
				printf("\n");
			}
		} 
		else 
		{
			printf("Error no palette found.\n");
			fclose(fp_infile);
			return 1;
		}
	}
	
	FILE *fp_outfile = fopen(output_file, out_fmt == OUTFMT_RAW ? "wb" : "w");
	if (!fp_outfile) 
	{
		perror("Cannot open output file");
		fclose(fp_infile);
		return 1;
	}
	
	// read into memory
	
	png_bytep *rows = malloc(sizeof(png_bytep) * height);
	for (int y = 0; y < height; y++)
	{
		//printf("%d\n",(int)png_get_rowbytes(png, info));
		rows[y] = malloc(png_get_rowbytes(png, info));
	}
	png_read_image(png, rows);
	fclose(fp_infile);
	
	// check indexing
	
	int colorHistory[ATARI_COLORS];
	for (int i = 0; i < ATARI_COLORS; i++) 
	{
		colorHistory[i] = 0;
	}
	for (int ty = 0; ty < height; ty++) 
	{
		for (int tx = 0; tx < width; tx++) 
		{		   
			png_bytep pixel = &(rows[ty][tx]);
			colorHistory[*pixel]++;
		}  
	}
	
	int usedColors = 0;
	printf("; Check color usage\n");
	for (int i = 0; i < ATARI_COLORS; i++) 
	{
		if(colorHistory[i] != 0)
		{
			usedColors++;
			printf("; Color %d = %d used\n",i,colorHistory[i]);
		}
	}
	printf("; Used colors %d\n",usedColors);
	if( (usedColors>2 && color_mode==CM_1BIT) || (usedColors>4 && color_mode==CM_2BIT) )
	{
		printf("; WARNING more colors used as output format allows, clip higher colors\n");
	}

	/* convert */
	
	if(gfx_mode == GF_LINEAR)
	{
		printf("; generate linear\n");
		
		int flushCountMax = color_mode==CM_1BIT?8:
			            color_mode==CM_2BIT?4:2;
				    
		int flushCount = 0;
		uint8_t byte = 0;
		
		for (int ty = 0; ty < height; ty++) 
		{
			if(fullPrint)
			{
				printf(";");
			}
			
			
			for (int tx = 0; tx < width; tx++) 
			{		   
				png_bytep pixel = &(rows[ty][tx]);
				uint8_t cindex = (uint8_t) *pixel;
				
				if(color_mode==CM_1BIT)
				{
					byte=(byte<<1) | (cindex&1);				
				}
				else if(color_mode==CM_2BIT)
				{
					byte=(byte<<2) | (cindex&3);				
				}
				else if(color_mode==CM_4BIT)
				{
					byte=(byte<<4) | (cindex&0xf);		
				}
				
				bufferedWrite(fullPrint,flushCountMax,&flushCount,fp_outfile,out_fmt,byte);
			}
					
			if(fullPrint)
			{
				printf("\n");
			}
		}
	}
	else if(gfx_mode == GF_8x8 || gfx_mode == GF_4x8)
	{
		int twidth = gfx_mode == GF_8x8?8:4;
		int theight = 8;
		
		printf("; tile size (%d x %d)\n",twidth,theight);
		
		if(width & twidth)
		{
			printf("; Width is not dividible by %d without rest (clip)\n",twidth);
		}
		
		if(height & theight)
		{
			printf("; Height is not dividible by %d without rest (clip)\n",theight);
		}
		
		int tcountx = width / twidth;
		int tcounty = height / theight;
		
		if(tcountx == 0 || tcounty == 0)
		{
			perror("Error Input picture is smaller then tile size\n");
		}
		else
		{
			int flushCountMax = color_mode==CM_1BIT?8:
					    color_mode==CM_2BIT?4:2;
			int flushCount = 0;
			uint8_t byte = 0;
			
			printf("; generate %d tiles (%dx%d)\n",tcountx*tcounty,tcountx,tcounty);
			
			for (int ty = 0; ty < height; ty += theight) 
			{
				for (int tx = 0; tx < width; tx += twidth) 
				{	
					if(fullPrint)
					{
						printf(";");
					}
			
					for(int iy=0;iy<theight;iy++)
					{						
						for(int ix=0;ix<twidth;ix++)
						{
							png_bytep pixel = &(rows[ty+iy][tx+ix]);
							uint8_t cindex = (uint8_t) *pixel;
				
							if(color_mode==CM_1BIT)
							{
								byte=(byte<<1) | (cindex&1);				
							}
							else if(color_mode==CM_2BIT)
							{
								byte=(byte<<2) | (cindex&3);				
							}
							else if(color_mode==CM_4BIT)
							{
								byte=(byte<<4) | (cindex&0xf);		
							}
							
							bufferedWrite(fullPrint,flushCountMax,&flushCount,fp_outfile,out_fmt,byte);
						}
					}
					
					if(fullPrint)
					{
						printf("\n");
					}
				}
			}
		}
		
	}
        
	// Cleanup
	
	fclose(fp_outfile);
	
	for (int y = 0; y < height; y++)
	{
		free(rows[y]);
	}
	free(rows);
	

	printf("Done. Output written to %s\n", output_file);
	return 0;
}