#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <ctype.h>

#include "fenders_bin.h"
#include "fenders_offsets.h"
#include "fendersdbl_bin.h"
#include "fendersdbl_offsets.h"

extern char *optarg;
extern int optind, opterr, optopt;

#ifndef VERSION
#define VERSION "???"
#endif

#define SELF "fenders"
#define BANNER SELF " v" VERSION " by B. Watson (WTFPL)\n"
#define DEFAULT_TITLE "atari arcade"
#define OPTIONS "hrsbcdiIt:v"

char *usage =
	BANNER
	"Install Fenders 3-sector loader in boot sectors of an ATR image\n"
	"Usage: " SELF " -[hrcsiv] [-t title] infile.atr [outfile.atr]\n"
	"  -h         Print this help message\n"
	"  -b         Create infile as a blank ATR, update it in-place.\n"
	"  -r         DON'T reboot (coldstart) the Atari if Reset is pressed\n"
	"  -c         Rotate colors during load\n"
	"  -d         Delete DOS.SYS and DUP.SYS\n"
	"  -s         Screen off after load\n"
	"  -i         In-place update (original renamed to end in ~)\n"
	"  -I         In-place update (NO backup)\n"
	"  -v         Set inverse video bit in title (blue/red text)\n"
	"  -t title   Set title (up to 20 chars, default: '" DEFAULT_TITLE "')\n";

typedef enum { SD, DD } density;

void set_title(char *title, density dens, int inverse) {
	int i;
	int offset;
	unsigned char *bin;
	int len = strlen(title);

	if(dens == SD) {
		offset = OFFSET_TITLE;
		bin = fenders_bin;
	} else {
		offset = OFFSET_TITLE_DD;
		bin = fendersdbl_bin;
	}

	/* zero out the title area first (zeroes are Atari spaces BTW) */
	memset(&bin[offset], 0, 20);

	if(len > 20) {
		len = 20;
		fprintf(stderr, SELF ": Truncating title to 20 characters\n");
	} else if(!len) {
		return;
	}

	/* convert ASCII to Atari screen codes (not the same as ATASCII)
		charset in GR.2 is space plus:
		!"#$%'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]^_

		Punctuation comes out orange (blue with -v)
		Digits and uppercase letters are orange/blue
		Lowercase letters are green/red
		` = green/red @
		{ = green/red [
		| = green/red ^
		} = green/red ]
		~ = green/red ^
	 */
	for(i=0; i<len; i++) {
		char c = title[i];

		if(c >= 32 && c <= 95)
			c -= 32;

		if(inverse)
			c |= 128;

		title[i] = c;
	}

	/* Center title in 20-byte title area */
	memcpy(&bin[offset + 10 - len / 2], title, len);

	/*
	fprintf(stderr, SELF ": Set title to \"");
	for(i=0; i<20; i++) {
		char c = bin[offset + i];
		if(!c) c = ' ';
		if(!isprint(c)) c = '.';
		putc(c, stderr);
	}

	putc('"', stderr);
	putc('\n', stderr);
	*/
}

/* for set_coldstart(), set_rot_color(), set_screen_off(),
	see fenders.dasm (and fendersdbl.dasm) to understand what's going on.
	Search for the OFFSET_* strings. */

void set_coldstart(density dens) {
	if(dens == SD) {
		fenders_bin[OFFSET_COLDST_1] = 1;    /* operand for LDY # (replaces 0) */
		fenders_bin[OFFSET_COLDST_2] = 0xea; /* NOP (replaces INY) */
	} else {
		fendersdbl_bin[OFFSET_COLDST_1_DD] = 1;
		fendersdbl_bin[OFFSET_COLDST_2_DD] = 0xea;
	}
}

void set_rot_color(density dens) {
	if(dens == SD)
		fenders_bin[OFFSET_ROTCOLOR] = 0x8d; /* STA abs (replace LDA abs) */
	else
		fendersdbl_bin[OFFSET_ROTCOLOR_DD] = 0x8d;
}

void set_screen_off(density dens) {
	if(dens == SD)
		fenders_bin[OFFSET_SCREENOFF] = 0x8d; /* STA abs (replace LDA abs) */
	else
		fendersdbl_bin[OFFSET_SCREENOFF_DD] = 0x8d;
}

/* TODO: error checking */
void call_axe(const char *args, const char *fname) {
	char cmd[8192];
	sprintf(cmd, "axe %s %s", args, fname);
	fprintf(stderr, SELF ": calling axe, cmd is: %s\n", cmd);
	system(cmd);
}

int main(int argc, char **argv) {
	int coldstart = 1, rot_color = 0, screen_off = 0;
	int c, res, size, inverse = 0;
	char title[21];
	int in_place = 0, rm_backup = 0, rm_dos_dup = 0, create_blank = 0;
	char rename_to[4096];
	unsigned char buf[384], *bin;
	char *infile = "-", *outfile = "-";
	FILE *in = stdin, *out = stdout;
	density dens;

	/* initialize title (may be changed by -t option) */
	strcpy(title, DEFAULT_TITLE);

	/* parse options */
	while( (c = getopt(argc, argv, OPTIONS)) != -1) {
		switch(c) {
			case 'h':
				printf(usage);
				exit(0);
				break;

			case 'r':
				coldstart = 0;
				break;

			case 'b':
				create_blank = 1;
				in_place = 1;
				rm_backup = 1;

			case 'c':
				rot_color = 1;
				break;

			case 'd':
				rm_dos_dup = 1;

			case 's':
				screen_off = 1;
				break;

			case 't':
				strcpy(title, optarg);
				break;

			case 'i':
				in_place = 1;
				break;

			case 'I':
				in_place = 1;
				rm_backup = 1;
				break;

			case 'v':
				inverse = 1;
				break;

			default:
				fprintf(stderr, usage);
				exit(1);
				break;
		}
	}

	/* get input filename if present */
	if(optind < argc)
		infile = argv[optind++];

	/* get output filename if present */
	if(!in_place && optind < argc)
		outfile = argv[optind++];

	if(optind < argc) {
		fprintf(stderr, SELF
				": Ignoring extra junk '%s ...' on command line.\n",
				argv[optind]);
	}

	if(in_place) {
		/* rename infile to infile~, set outfile to old infile */
		int len = strlen(infile);

		if(strcmp(infile, "-") == 0) {
			fprintf(stderr, SELF
					": Can't use in-place mode with standard input.\n");
			exit(1);
		}

		strcpy(rename_to, infile);
		rename_to[len] = '~';
		rename_to[len + 1] = '\0';

		if(!rm_backup)
			fprintf(stderr, SELF ": Backing up %s to %s\n", infile, rename_to);

		if(create_blank) {
			call_axe("-b", infile);
		}

		if(link(infile, rename_to)) {
			perror("link()");
			exit(1);
		}

		if(unlink(infile)) {
			perror("unlink()");
			exit(1);
		}

		outfile = infile;
		infile = rename_to;
	}

	/* open input and output files, if not stdin/stdout */
	if(strcmp(infile, "-") != 0) {
		in = fopen(infile, "rb");
		if(!in) {
			perror(infile);
			exit(1);
		}
		if(rm_backup) unlink(infile);
	}

	/* read ATR header */
	res = fread(buf, 1, 16, in);
	if(res < 16) {
		perror(infile);
		exit(1);
	}

	/* make sure it's an ATR image */
	if( !(buf[0] == 0x96 && buf[1] == 0x02) ) {
		fprintf(stderr, SELF
				": %s not an ATR file (no NICKATARI signature)!\n"
				"If this is an XFD file, try xfd2atr\n",
				infile);
		exit(2);
	}

	/* get sector size. The single- and double-density versions of
		the loader are totally different, so pick the one we need. */
	if( (buf[4] == 0x80 && buf[5] == 0x00) ) {
		dens = SD;
		bin = fenders_bin;
	} else if( (buf[4] == 0x00 && buf[5] == 0x01) ) {
		dens = DD;
		bin = fendersdbl_bin;
	} else {
		fprintf(stderr, SELF ": ATR image must have 128 or 256 byte sectors\n");
		exit(2);
	}

	/* modify the loader according to the user's options */
	set_title(title, dens, inverse);
	if(coldstart) set_coldstart(dens);
	if(rot_color) set_rot_color(dens);
	if(screen_off) set_screen_off(dens);

	/* size of ATR image in bytes (minus the header). We don't support
		DD images less than 720 sectors. */
	size = (buf[2] + (buf[3] << 8) + (buf[6] << 16)) * 16;
	if(dens == SD && size < (128 * 369)) {
		fprintf(stderr, SELF
				": ATR is single density < 369 sectors, not supported\n"
				"Use atrsize to grow the image.\n");
		exit(2);
	} else if(dens == SD && size > (128 * 720)) {
		fprintf(stderr, SELF
				": ATR is single density > 720 sectors; "
				"some files may not appear in menu.\n");
	} else if(dens == DD && size < (128 * 3 + 256 * 717)) {
		fprintf(stderr, SELF
				": ATR is double density < 720 sectors, not supported\n"
				"Use atrsize to grow the image.\n");
		exit(2);
	} else if(dens == DD && size > (128 * 3 + 256 * 717)) {
		/* 20071005 bkw: whoops, we were truncating large images to 180K.
			The bootloader doesn't work with MyDOS >180K formats anyway, so
			don't try.
		fprintf(stderr, SELF
				": ATR file is double density > 720 sectors; some files may not"
				"appear in the menu or load correctly\n");
				*/

		/* Abort instead */
		fprintf(stderr, SELF
				": ATR is double density > 720 sectors, not supported "
				"by bootloader (try MyPicoDOS)\n");
		exit(2);
	}

	/* Input looks OK, open the output... */
	if(strcmp(outfile, "-") != 0) {
		out = fopen(outfile, "wb");
		if(!out) {
			fclose(in);
			perror(outfile);
			exit(1);
		}
	} else if(isatty(fileno(stdout))) {
		/* don't scare the n00bs! */
		fprintf(stderr,
				SELF ": Standard output is a terminal, not writing binary data.\n"
				"Either redirect to a file or set the output filename.\n");
		exit(1);
	}

	/* write ATR header */
	res = fwrite(buf, 1, 16, out);
	if(res < 16) {
		perror(outfile);
		exit(1);
	}

	/* read (and ignore) first 3 sectors. A DD image still uses SD
		sectors for the first 3 (boot) sectors on the disk. */
	res = fread(buf, 1, 384, in);
	if(res < 384) {
		perror(infile);
		exit(1);
	}

	/* Write the loader. For SD disks, this is the whole thing.
		For DD disks, this is the first 384 bytes (3 boot sectors),
		and we'll have to write the rest of the object code to sector 720 */
	res = fwrite(bin, 1, 384, out);
	if(res < 384) {
		perror(outfile);
		exit(1);
	}

	if(dens == SD) {
		/* single density can just use a simple copy loop */
		while( (c = getc(in)) != EOF )
			putc(c, out);
	} else {
		/* double density: copy sectors 4-719 as-is... */
		for(c=4; c<720; c++) {
			if(fread(buf, 1, 256, in) < 256) {
				if(feof(in)) {
					fprintf(stderr, SELF
							": got premature EOF (bad/truncated ATR image).\n");
				} else {
					perror(infile);
				}
				exit(1);
			}

			if(fwrite(buf, 1, 256, out) < 256) {
				perror(outfile);
				exit(1);
			}
		}

		/* TODO: check the VTOC and/or look for non-zero data in sector
			720, warn the user if the sector was in use. */

		/* TODO: fix bootloader to work with MyDOS-style sector link bytes.
			Also, store last part of bootloader somewhere not used by MyDOS,
			maybe sector 369 (last directory sector, unused on disks with
			less than 56 files... and 56 files is way too many to fit on
			screen in GR.1). This should happen on both single and double
			density images. */

		/* TODO: examine directory sectors, look for:
			- DOS 2.5 extended files. Either warn about them, or clear the
			extended flag (which causes the bootloader to load them just fine).
			- MyDOS subdirectories. It's probably best to abort in that case.
			- Non-DOS-compatible disk formats (boot disks, SpartaDOS, Atari
			DOS 3 or 4).
		*/

		/* write the rest of the loader code code to sector 720.
		 	The code in the boot sectors will load sector 720. */
		if(fwrite(bin+384, 1, 256, out) < 256) {
			perror(outfile);
			exit(1);
		}
	}

	/* set return value: 0 for success, 1 for failure */
	c = 0;

	if(ferror(in)) {
		perror(infile);
		c = 1;
	}

	if(ferror(out)) {
		perror(outfile);
		c = 1;
	}

	fclose(in);
	fclose(out);

	if(rm_dos_dup) {
		/* TODO: check for enhanced density! */
		if(dens != SD) {
			fprintf(stderr, SELF
					": -d option only works on single-density.\n");
			exit(2);
		}
		if(strcmp(outfile, "-") == 0) {
			fprintf(stderr, SELF
					": Can't use -d option with standard output.\n");
			exit(2);
		}
		call_axe("-D DOS.SYS", outfile);
		call_axe("-D DUP.SYS", outfile);
		call_axe("-D AUTORUN.SYS", outfile);
	}

	/* ...and I'm spent! */
	return c;
}