#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>

/* FIXME: This code is terrible.
	I should have written an ATR library for atrsize, atr2xfd, and xfd2atr
	to use, instead of the copy/paste mess I've created. */

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

#define SELF "atrsize"
#define USAGE \
	SELF " v" VERSION " by B. Watson (WTFPL)\n" \
	"Usage: " SELF " -[bB] image.atr [sectors]\n" \
	"  -b  Create image.atr as a new, blank, single-density image\n" \
	"  -B  Create image.atr as a new, blank, double-density image\n" \
	"Original image.atr is backed up as image.atr~\n" \
	"Without [sectors], size is rounded up to next 'standard' size\n"

char *usage = USAGE;

int main(int argc, char **argv) {
	struct stat st;
	int filesize = 0;
	int headersize = 0;
	int blank = 0;
	int len, sec;
	unsigned int secsize, seccount, paras, newseccount = 0, last_data_sec = 0;
	char *newsize, infile[4096], *outfile;
	unsigned char buf[384];
	FILE *in, *out;

	if(argc < 2 || argc > 4) {
		fprintf(stderr, usage);
		exit(1);
	}

	if(strcmp(argv[1], "-b") == 0) {
		blank = 128;
		argv++;
		argc--;
	} else if(strcmp(argv[1], "-B") == 0) {
		blank = 256;
		argv++;
		argc--;
	}

	newsize = argv[2];

	outfile = argv[1];
	strcpy(infile, outfile);
	if( !(in = fopen(infile, "rb")) ) {
		if(!blank) {
			fprintf(stderr,
					SELF ": %s: %s\n", infile, strerror(errno));
			exit(1);
		}

		fprintf(stderr,
				SELF ": %s does not exist, creating blank image\n", infile);

		secsize = blank;
		seccount = 720;

		memset(buf, 0, 16);

		buf[0] = 0x96;
		buf[1] = 0x02;
		buf[4] = secsize & 0xff;
		buf[5] = (secsize >> 8) & 0xff;

		filesize = headersize = 128 * 720 + (secsize - 128) * 717;
		newseccount = seccount = 720;

		paras = headersize / 16;
		buf[2] = paras & 0xff;
		buf[3] = (paras >> 8) & 0xff;
		buf[6] = (paras >> 16) & 0xff;
	}

	if(in) {
		if(blank) {
			fprintf(stderr,
					SELF ": won't create blank image %s: file exists\n", infile);
			exit(1);
		}

		if(fstat(fileno(in), &st)) {
			fprintf(stderr, "Can't determine size of %s: %s\n",
					infile, strerror(errno));
			exit(1);
		} else {
			filesize = st.st_size - 16;
			if(filesize < 384) {
				fprintf(stderr,
						SELF ": %s is too small to be a valid ATR image\n",
						infile);
				exit(1);
			}
		}

		if(fread(buf, 1, 16, in) != 16) {
			fprintf(stderr, SELF ": %s: %s\n", infile, strerror(errno));
			exit(1);
		}
	}

	if( !(buf[0] == 0x96 && buf[1] == 0x02) ) {
		fprintf(stderr,
				SELF ": (fatal) %s not an ATR file (no NICKATARI signature)!\n",
				infile);
		exit(2);
	}

	secsize = buf[4] + (buf[5] << 8);

	if( !(secsize == 128 || secsize == 256) ) {
		fprintf(stderr,
				SELF ": (fatal) %s has invalid sector size %d\n",
				infile, secsize);
		exit(2);
	}

	paras = buf[2] + (buf[3] << 8) + (buf[6] << 16);
	headersize = paras * 16;

	if(filesize && (filesize != headersize)) {
		fprintf(stderr,
				SELF ": warning: %s file size (%d bytes) doesn't agree with "
				"ATR header (%d bytes). File may be truncated or corrupt.\n",
				infile,
				filesize,
				headersize);
	}

	if(secsize == 128) {
		fprintf(stderr, SELF ": single density image (128 bytes/sector)\n");
		seccount = filesize / secsize;
	} else {
		fprintf(stderr, SELF ": double density image (256 bytes/sector)\n");
		seccount = (filesize - 128 * 3) / secsize + 3;
	}

	newseccount = seccount;
	if(newsize == NULL) {
		/* figure out appropriate new size for this image */
		if(secsize == 128) {
			if(seccount < 720)
				newseccount = 720;
			else if(seccount > 720 && seccount < 1040)
				newseccount = 1040;
			else if(seccount >= 1040) {
				fprintf(stderr, SELF ": must specify a sector count "
						"for SD images >= 1040 sectors.\n");
				exit(1);
			}
		} else {
			if(seccount < 720)
				newseccount = 720;
			else if(seccount > 720 && seccount < 1440)
				newseccount = 1440;
			else if(seccount >= 1440) {
				fprintf(stderr, SELF ": must specify a sector count "
						"for DD images >= 1440 sectors.\n");
				exit(1);
			}
		}

		/*
		if(in) {
			if(newseccount == seccount) {
				fprintf(stderr,
						SELF ": image is already %d sectors, nothing to do\n", seccount);
				exit(0);
			} else if(!newseccount) {
				fprintf(stderr,
						SELF ": image is already standard size, nothing to do\n");
				exit(0);
			}
		}
		*/
	} else {
		newseccount = atoi(newsize);
		if(newseccount < 3 || newseccount > 65535) {
			fprintf(stderr,
					SELF ": invalid sector count (must be 3 - 65535)\n");
			exit(1);
		}
	}

	if(newseccount < 368) {
		fprintf(stderr, SELF ": warning: "
				"Output image will not have Atari/MyDOS directory sectors.\n");
	}

	/* fix up ATR header */
	if(secsize == 128)
		headersize = secsize * newseccount;
	else
		headersize = (128 * 3) + (newseccount - 3) * secsize;

	paras = headersize / 16;
	buf[2] = paras & 0xff;
	buf[3] = (paras >> 8) & 0xff;
	buf[6] = (paras >> 16) & 0xff;

	if(in) {
		/* make backup file */
		len = strlen(infile);
		infile[len] = '~';
		infile[len + 1] = '\0';

		unlink(infile);
		if(link(outfile, infile)) {
			fprintf(stderr, SELF ": can't create %s: %s\n",
					infile, strerror(errno));
			exit(1);
		}

		if(unlink(outfile)) {
			fprintf(stderr, SELF ": can't delete %s: %s\n",
					outfile, strerror(errno));
			exit(1);
		}
	}

	if( !(out = fopen(outfile, "wb")) ) {
		fprintf(stderr, SELF ": %s: %s\n", outfile, strerror(errno));
		exit(1);
	}

	/* write ATR header */
	if(fwrite(buf, 1, 16, out) != 16) {
		fprintf(stderr, SELF ": %s: %s\n", outfile, strerror(errno));
		exit(1);
	}

	/* read first 3 sectors (always present, always SD) */
	if(in) {
		if(fread(buf, 1, 384, in) != 384) {
			fprintf(stderr, SELF ": %s: %s\n", infile, strerror(errno));
			exit(1);
		}
	} else {
		memset(buf, 0, 384);
	}

	/* write first 3 sectors */
	if(fwrite(buf, 1, 384, out) != 384) {
		fprintf(stderr, SELF ": %s: %s\n", outfile, strerror(errno));
		exit(1);
	}

	for(sec = 4; (sec <= seccount) && (sec <= newseccount); sec++) {
		int i, has_data = 0;

		if(in) {
			if(fread(buf, 1, secsize, in) != secsize) {
				fprintf(stderr, SELF ": %s: %s\n", infile, strerror(errno));
				exit(1);
			}
		}

		for(i=0; i<secsize; i++)
			has_data |= buf[i];

		if(has_data)
			last_data_sec = sec;

		if(fwrite(buf, 1, secsize, out) != secsize) {
			fprintf(stderr, SELF ": %s: %s\n", outfile, strerror(errno));
			exit(1);
		}
	}

	if(last_data_sec)
		fprintf(stderr, SELF ": last non-empty sector was %d\n", last_data_sec);
	else
		fprintf(stderr, SELF ": image is blank (no data)\n");

	if(newseccount < seccount) {
		fprintf(stderr, SELF ": %s truncated to %d sectors, OK\n",
				outfile, newseccount);
	} else if(newseccount == seccount) {
		fprintf(stderr, SELF ": %s rewritten at %d sectors, OK\n",
				outfile, newseccount);
	} else {
		memset(buf, 0, secsize);
		for( ; sec <= newseccount; sec++) {
			if(fwrite(buf, 1, secsize, out) != secsize) {
				fprintf(stderr, SELF ": %s: %s\n", outfile, strerror(errno));
				exit(1);
			}
		}
		fprintf(stderr, SELF ": %s extended to %d sectors, OK\n",
				outfile, newseccount);
	}

	if(in) fclose(in);
	fclose(out);
	exit(0);
}