#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;
int list_details = 0;

/* CD image state */
MirageContext *mirage = NULL;
MirageDisc *disc = NULL;
MirageSession *session = NULL;
int output_track_number = 0;
int total_bytes = 0;
int total_sectors = 0;
int session_count = 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;
}

/* decide what to name a data track's output file. currently
	anything not recognized as Apple HFS/HFS+ is assumed to
	be an ISO. These 2 cases cover all the images I've ever
	run into in the wild. */
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;
}

void print_time(int sectors) {
	int mm, ss, fr;

	mm = sectors / (60 * 75);
	sectors %= 60 * 75;
	ss = sectors / 75;
	sectors %= 75;
	fr = (int)((double)sectors / 75.0L * 10.0L);

	printf("%d:%02d.%1d, ", mm, ss, fr);
}

/* list a track without reading the whole thing into memory */
void list_track(int t) {
	MirageTrack *track;
	MirageSector *sector;
	const guint8 *buf;
	int sector_type, sec = 0;
	int bytes;
	gint len;
	int tracklen;

	if(session_count == 1)
		printf("  %d: ", t + 1);
	else
		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:
			printf("data(%s): ", get_data_extension(track));
			break;
		case MIRAGE_SECTOR_AUDIO:
			printf("audio: ");
			break;
		case MIRAGE_SECTOR_MODE2_FORM2:
		case MIRAGE_SECTOR_MODE2_MIXED:
			printf("video(?): ");
			break;
		default:
			printf("<unknown %d>: ", sector_type);
			break;
	}

	sec = mirage_track_get_track_start(track);
	sector = mirage_track_get_sector(track, sec, 0, NULL);
	if(!(mirage_sector_get_data(sector, &buf, &len, &gerr)))
		die(NULL);
	tracklen = mirage_track_layout_get_length(track);

	bytes = len * tracklen;

	if(sector_type == MIRAGE_SECTOR_AUDIO)
		print_time(tracklen);

	if(list_details)
		printf("%sMB (%d = %d sectors of %d bytes)\n", human_mb(bytes), bytes, tracklen, len);
	else
		printf("%sMB\n", human_mb(bytes));

	total_bytes += bytes;
	total_sectors += tracklen;
}

/* helper for extract_track() */
SNDFILE *open_sndfile_output(const char *outfile) {
	SNDFILE *sfout = NULL;
	struct SF_INFO sfi;

	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))) {
		putchar('\n');
		die(sf_strerror(sfout));
	}
	sf_command(sfout, SFC_SET_VBR_ENCODING_QUALITY, &quality, sizeof(double));
	return sfout;
}

/* helper for extract_track() */
FILE *open_output(const char *outfile) {
	FILE *out;

	if(!(out = fopen(outfile, "wb"))) {
		putchar('\n');
		fprintf(stderr, "%s: %s: can't open output file\n", self, outfile);
		die(NULL);
	}

	return out;
}

/* extract a track. meat and potatoes here. */
void extract_track(int t) {
	MirageTrack *track;
	MirageSector *sector;
	int sector_type, sec = 0, bytes = 0, use_sndfile = 0;
	gint len;
	const guint8 *buf;
	const char *ext;
	char outfile[PATH_MAX + 1];
	FILE *out = NULL;
	SNDFILE *sfout = NULL;

	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);
			if(!want_data) return;
			break;
		case MIRAGE_SECTOR_AUDIO:
			if(!want_audio) return;
			if(outfmt[0] != 'c')
				use_sndfile = 1;
			ext = outfmt;
			break;
		case MIRAGE_SECTOR_MODE2_FORM2:
		case MIRAGE_SECTOR_MODE2_MIXED:
			if(!want_data) return;
			ext = "mpeg";
			break;
		default:
			if(!want_data) return;
			ext = "raw";
			break;
	}

	sprintf(outfile, "%s%02d.%s", outfilebase, output_track_number, ext);
	printf("     Extracting to %s...", outfile);
	fflush(stdout);

	if(use_sndfile)
		sfout = open_sndfile_output(outfile);
	else
		out = open_output(outfile);

	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(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("OK\n");

	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;
}

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

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] [-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);
}

/* 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] == '-') {
			/* avoid confusion, I hope */
			if(argv[0][1] && argv[0][2]) {
				fprintf(stderr, "%s: invalid option '%s', bundling not supported\n", self, argv[0]);
				exit(1);
			}
			char *nextarg = argv[1];
			switch(argv[0][1]) {
				case 'L':
					want_track = -1;
					list_details = 1;
					break;
				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;
						if(strlen(outfilebase) > PATH_MAX - 6)
							die("-b filename too long");
					} 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);

	session_count = mirage_disc_get_number_of_sessions(disc);
	if(session_count > 1)
		printf("This is a multi-session disc\n");

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

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

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

			list_track(t);

			if(want_track == 0 || want_track == output_track_number)
				extract_track(t);
		}
		if(!t) printf("  (no tracks in session)\n");
		g_object_unref(session);
	}

	if(list_details) {
		printf("Total size: %sMB (%d, %d sectors)\n", human_mb(total_bytes), total_bytes, total_sectors);
	} else {
		printf("Total size: %sMB\n", human_mb(total_bytes));
	}

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

	return 0;
}