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

#include <sndfile.h>
#include <mirage/mirage.h>

#ifndef PATH_MAX
#  define PATH_MAX 1024
#endif

/* This should be big enough to hold any CD sector */
#define BUF_SIZE 4096

/* main() sets this to our executable name */
const char *self;

/* command-line args/opts set these */
char *imagefile = NULL;
char *outfilebase = "track";
char *outfmt = "wav";
int want_track = 0; /* -1 = no tracks, 0 = all, 1-99 = just that one */
double quality = 0.7l;
int swap_bytes = 0;
int autoname = 0;
int want_data = 1;
int want_audio = 1;
char *password = NULL;

/* CD image state */
MirageContext *mirage = NULL;
MirageDisc *disc = NULL;
MirageSession *session = NULL;
int output_track_number = 0;
int total_bytes = 0;

/* libmirage errors set this */
GError *gerr = NULL;

/* print executable name: error message, then exit with failure status. */
void die(const char *msg) {
	if(msg) {
		/* do nothing, use msg as-is */
	} else if(gerr && gerr->message) {
		/* msg was NULL, if there's a libmirage error message, use it */
		msg = gerr->message;
	} else {
		/* msg was NULL and there's nothing from libmirage, fallback: */
		msg = strerror(errno);
	}
	fprintf(stderr, "%s: %s\n", self, msg);
	exit(1);
}

/* either swap bytes into a static buffer and return a pointer
	to it, or don't swap bytes and return the original buffer.
	Paranoia, because mirage_sector_get_data() returns a const
	buffer (probably safe to cast away the const, but...) */
const guint8 *apply_s_opt(const guint8 *buf, gint len) {
	int i;
	static guint8 newbuf[BUF_SIZE];

	if(!swap_bytes) return buf;

	for(i = 0; i < len; i += 2) {
		newbuf[i] = buf[i + 1];
		newbuf[i + 1] = buf[i];
	}

	return (const guint8 *)newbuf;
}

/* each call to spinner() prints the next character and a backspace */
void spinner(void) {
	static char spinchars[] = "-/|\\";
	static int count = 0;

	putchar(spinchars[count++]);
	putchar('\b');
	fflush(stdout);

	if(count == 4) count = 0;
}

/* format human-readable quantities (567.5MB is a lot easier to read
	than 595024576). these are 2^10 byte megabytes, as the gods intended. */
char *human_mb(int bytes) {
	static char buf[10];
	sprintf(buf, "%2.1f", (float)bytes / 1048576.0);
	return buf;
}

const char *get_data_extension(MirageTrack *track) {
	MirageSector *sector;
	const char *ext = "iso";
	int sec;
	gint len;
	const guint8 *buf;

	/* examine first sector, see if it's got Apple magic */
	sec = mirage_track_get_track_start(track);

	if(!(sector = mirage_track_get_sector(track, sec, 0, &gerr)))
		die(NULL);

	if(!(mirage_sector_get_data(sector, &buf, &len, &gerr)))
		die(NULL);

	if(buf[0] == 'E' && buf[1] == 'R' && buf[3] == 0 && (buf[2] == 0x02 || buf[2] == 0x08))
	{
		ext = "hfs";
		/* if it also has ISO9660 magic, it's a hybrid */
		if((sector = mirage_track_get_sector(track, sec + 16, 0, NULL)) &&
				mirage_sector_get_data(sector, &buf, &len, NULL) &&
				(memcmp("\x01" "CD001", buf, 6) == 0))
			ext = "hfs.iso";
	}

	return ext;
}


/* list and (optionally) extract a track. this should be broken up into
	smaller functions for clarity. */
void process_track(int t, int extract) {
	MirageTrack *track;
	MirageSector *sector;
	int sector_type, sec = 0;
	gint len;
	const guint8 *buf;
	const char *ext;
	char outfile[PATH_MAX + 1];
	FILE *out = NULL;
	SNDFILE *sfout = NULL;
	struct SF_INFO sfi;
	int bytes = 0;

	printf("  %d (-t %d): ", t + 1, output_track_number);

	if(!(track = mirage_session_get_track_by_index(session, t, &gerr)))
		die(NULL);

	sector_type = mirage_track_get_sector_type(track);

	switch(sector_type) {
		case MIRAGE_SECTOR_MODE1:
		case MIRAGE_SECTOR_MODE2_FORM1:
			ext = get_data_extension(track);
			printf("data(%s): ", ext);
			if(!want_data) extract = 0;
			break;
		case MIRAGE_SECTOR_AUDIO:
			printf("audio: ");
			if(!want_audio) extract = 0;
			ext = outfmt;
			break;
		default:
			printf("<unknown>: ");
			ext = "raw";
			if(!want_data) extract = 0;
			break;
	}

	sprintf(outfile, "%s%02d.%s", outfilebase, output_track_number, ext);

	if(extract) {
		if((sector_type == MIRAGE_SECTOR_AUDIO) && (outfmt[0] != 'c')) {
			sfi.samplerate = 44100;
			sfi.channels = 2;

			switch(outfmt[0]) {
				case 'o':
					sfi.format = SF_FORMAT_OGG | SF_FORMAT_VORBIS;
					break;
				case 'f':
					sfi.format = SF_FORMAT_FLAC | SF_FORMAT_PCM_16;
					break;
				case 'w':
					sfi.format = SF_FORMAT_WAV | SF_FORMAT_PCM_16;
					break;
				default:
					/* this should never happen */
					break;
			}

			if(!(sfout = sf_open(outfile, SFM_WRITE, &sfi)))
				die(sf_strerror(sfout));
			sf_command(sfout, SFC_SET_VBR_ENCODING_QUALITY, &quality, sizeof(double));

		} else {
			if(!(out = fopen(outfile, "wb"))) {
				die("can't open output file");
			}
		}
	}

	sec = mirage_track_get_track_start(track);
	while((sector = mirage_track_get_sector(track, sec++, 0, NULL))) {
		if(!(mirage_sector_get_data(sector, &buf, &len, &gerr)))
			die(NULL);
		bytes += len;
		if(extract) {
			if(sfout) {
				int items = len / 2;
				if(sf_write_short(sfout, (const short *)apply_s_opt(buf, len), items) != items)
					die(sf_strerror(sfout));
			} else {
				if(fwrite(buf, len, 1, out) != 1)
					die(strerror(errno));
			}
			if(isatty(1) && !(sec % 50)) spinner();
		}
		g_object_unref(sector);
	}

	if(sfout) sf_close(sfout);
	if(out) fclose(out);

	// printf("%d bytes (%sMB, %d sectors)\n", bytes, human_mb(bytes), sec);
	printf("%d bytes (%sMB)\n", bytes, human_mb(bytes));
	if(extract) printf("    Extracted to %s\n", outfile);
	total_bytes += bytes;

	return;
}

/* assign base filename (without extension), based on image filename */
void set_auto_name(void) {
	static char autoname[PATH_MAX + 1];
	char *p;

	strcpy(autoname, imagefile);
	if((p = strrchr(autoname, '.'))) {
		*p = '\0';
	}

	outfilebase = autoname;
}

/* return true if outfmt is a supported output format */
int outfmt_ok(char *outfmt) {
	static const char *supported[] = { "cdda", "wav", "flac", "ogg", NULL };
	const char **p;

	for(p = supported; *p; p++)
		if(strcmp(*p, outfmt) == 0) return 1;

	return 0;
}

void usage(void) {
	char spaces[42];
	int i, count;

	count = strlen(self) + 8;
	if(count >= 40) count = 40;

	for(i = 0; i < count; i++)
		spaces[i] = ' ';
	spaces[i] = '\0';

	printf(
			"miragextract v" VERSION " by B. Watson, WTFPL\n"
			"Usage: %s [-l] [-s] [-n] [-a] [-d] [-t track] [-b base]\n"
			"%s[-f fmt] [-q quality] [-p password] image-file\n"
			"See man page for details\n", self, spaces);
	exit(0);
}

gchar *password_callback(gpointer user_data) {
	if(!password) return NULL;
	return g_strdup(password);
}

/* this should probably use something like getopts(), but I'm trying
	to avoid GNU extensions. */
void parse_args(int argc, char **argv) {
	if(argc == 1) usage();
	if(argc > 1 && strcmp(argv[1], "--help") == 0) usage();

	while(++argv, --argc) {
		if(argv[0][0] == '-') {
			char *nextarg = argv[1];
			switch(argv[0][1]) {
				case 'l':
					want_track = -1;
					break;
				case 'a':
					want_audio = 1;
					want_data = 0;
					break;
				case 'd':
					want_audio = 0;
					want_data = 1;
					break;
				case 's':
					swap_bytes = 1;
					break;
				case 't':
					if(nextarg) {
						argv++, argc--;
						want_track = atoi(nextarg);
						if(want_track < 1 || want_track > 99)
							die("invalid track number for -t (must be 1-99)");
					} else
						die("-t option requires track argument");
					break;
				case 'b':
					if(nextarg) {
						argv++, argc--;
						outfilebase = nextarg;
					} else
						die("-b option requires base filename argument");
					break;
				case 'f':
					if(nextarg) {
						argv++, argc--;
						outfmt = nextarg;
						if(!outfmt_ok(outfmt))
							die("unsupported output format for -f");
					} else
						die("-f option requires format argument");
					break;
				case 'p':
					if(nextarg) {
						argv++, argc--;
						password = nextarg;
					} else
						die("-p option requires password");
					break;
				case 'q':
					if(nextarg) {
						argv++, argc--;
						quality = (double)atoi(nextarg) / 10.0l;
						if(quality < 0.0l || quality > 1.0l)
							die("invalid quality for -q (must be 1-10)");
					} else
						die("-t option requires track argument");
					break;
				case 'n':
					autoname = 1;
					break;
				default:
					fprintf(stderr, "unrecognized option '%s', try --help\n", *argv);
					exit(1);
					break;
			}
		} else if(imagefile) {
			die("too many arguments on command line");
		} else {
			imagefile = *argv;
		}
	}

	if(!imagefile) die("missing image file argument");

	if(autoname) set_auto_name();
}

/* open the disc image, iterate over the tracks, calling
	process_track() on each one. */
int main(int argc, char **argv) {
	const char *p;
	gchar *fn[2];
	int s, t;

	/* get rid of any path elements in our executable name */
	self = argv[0];
	for(p = self; *p; p++)
		if(p[0] == '/' && p[1]) self = p + 1;

	parse_args(argc, argv);

	if(!((mirage = g_object_new(MIRAGE_TYPE_CONTEXT, NULL))))
		die("can't get mirage context");

	if(!mirage_initialize(&gerr))
		die(NULL);

	mirage_context_set_password_function(mirage, password_callback, NULL, NULL);

	/* TODO: find out what image formats require multiple filenames, if any */
	fn[0] = (gchar *)imagefile;
	fn[1] = NULL;

	if(!(disc = mirage_context_load_image(mirage, fn, &gerr)))
		die(NULL);

	for(s = 0; s < mirage_disc_get_number_of_sessions(disc); s++) {
		if(!(session = mirage_disc_get_session_by_index(disc, s, &gerr)))
			die(NULL);

		printf("Session %d tracks:\n", s + 1);

		for(t = 0; t < mirage_session_get_number_of_tracks(session); t++) {
			int extract = 0;
			output_track_number++;

			if(want_track == 0 || want_track == output_track_number)
				extract = 1;

			process_track(t, extract);
		}
		if(!t) printf("  (no tracks in session)\n");
		g_object_unref(session);
	}

	g_object_unref(disc);
	g_object_unref(mirage);
	mirage_shutdown(NULL);

	printf("Total size: %d bytes (%sMB)\n", total_bytes, human_mb(total_bytes));

	return 0;
}