/*
   Copyright (C) Andrew Tridgell 1998-2003,
   Con Kolivas 2006-2009

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 2 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
*/
/* lrzip compression - main program */
#include "rzip.h"

struct rzip_control control;

static void usage(void)
{
	printf("lrzip version %d.%d%d\n", LRZIP_MAJOR_VERSION, LRZIP_MINOR_VERSION, LRZIP_MINOR_SUBVERSION);
	printf("Copyright (C) Con Kolivas 2006-2009\n\n");
	printf("Based on rzip ");
	printf("Copyright (C) Andrew Tridgell 1998-2003\n");
	printf("usage: lrzip [options] <file...>\n");
	printf(" Options:\n");
	printf("     -w size       compression window in hundreds of MB\n");
	printf("                   default chosen by heuristic dependant on ram and chosen compression\n");
	printf("     -d            decompress\n");
	printf("     -o filename   specify the output file name and/or path\n");
	printf("     -O directory  specify the output directory when -o is not used\n");
	printf("     -S suffix     specify compressed suffix (default '.lrz')\n");
	printf("     -f            force overwrite of any existing files\n");
	printf("     -D            delete existing files\n");
	printf("     -P            don't set permissions on output file - may leave it world-readable\n");
	printf("     -q            don't show compression progress\n");
	printf("     -L level      set lzma/bzip2/gzip compression level (1-9, default 7)\n");
	printf("     -n            no backend compression - prepare for other compressor\n");
	printf("     -l            lzo compression (ultra fast)\n");
	printf("     -b            bzip2 compression\n");
	printf("     -g            gzip compression using zlib\n");
	printf("     -z            zpaq compression (best, extreme compression, extremely slow)\n");
	printf("     -M            Maximum window and level - (all available ram and level 9)\n");
	printf("     -T value      Compression threshold with LZO test. (0 (nil) - 10 (high), default 1)\n");
	printf("     -N value      Set nice value to value (default 19)\n");
	printf("     -v[v]         Increase verbosity\n");
	printf("     -V            show version\n");
#if 0
	/* damn, this will be quite hard to do */
	printf("     -t          test compressed file integrity\n");
#endif
	printf("\nnote that lrzip cannot operate on stdin/stdout\n");
}

static void write_magic(int fd_in, int fd_out)
{
	struct stat st;
	char magic[24];
	int i;

	memset(magic, 0, sizeof(magic));
	strcpy(magic, "LRZI");
	magic[4] = LRZIP_MAJOR_VERSION;
	magic[5] = LRZIP_MINOR_VERSION;

	if (fstat(fd_in, &st) != 0)
		fatal("bad magic file descriptor!?\n");

	memcpy(&magic[6], &st.st_size, 8);

	/* save LZMA compression flags */
	if (LZMA_COMPRESS(control.flags)) {
		for (i = 0; i < 5; i++)
			magic[i+16] = (char)control.lzma_properties[i];
	}

	if (lseek(fd_out, 0, SEEK_SET) != 0)
		fatal("Failed to seek to BOF to write Magic Header\n");

	if (write(fd_out, magic, sizeof(magic)) != sizeof(magic))
		fatal("Failed to write magic header\n");
}

static void read_magic(int fd_in, i64 *expected_size)
{
	uint32_t v;
	char magic[24];
	int i;

	if (read(fd_in, magic, sizeof(magic)) != sizeof(magic))
		fatal("Failed to read magic header\n");

	*expected_size = 0;

	if (strncmp(magic, "LRZI", 4) != 0) {
		fatal("Not an lrzip file\n");
	}

	memcpy(&control.major_version, &magic[4], 1);
	memcpy(&control.minor_version, &magic[5], 1);

	/* Support the convoluted way we described size in versions < 0.40 */
	if (control.major_version == 0 && control.minor_version < 4) {
		memcpy(&v, &magic[6], 4);
		*expected_size = ntohl(v);
		memcpy(&v, &magic[10], 4);
		*expected_size |= ((i64)ntohl(v)) << 32;
	} else
		memcpy(expected_size, &magic[6], 8);

	/* restore LZMA compression flags only if stored */
	if ((int) magic[16]) {
		for (i = 0; i < 5; i++)
			control.lzma_properties[i]=magic[i+16];
	}
	if (control.flags & (FLAG_VERBOSITY | FLAG_VERBOSITY_MAX)) {
		printf("Detected lrzip version %d.%d file.\n", control.major_version, control.minor_version);
		fflush(stdout);
	}
}

/* preserve ownership and permissions where possible */
static void preserve_perms(int fd_in, int fd_out)
{
	struct stat st;

	if (fstat(fd_in, &st) != 0)
		fatal("Failed to fstat input file\n");
	if (fchmod(fd_out, (st.st_mode & 0777)) != 0)
		fatal("Failed to set permissions on %s\n", control.outfile);

	/* chown fail is not fatal */
	fchown(fd_out, st.st_uid, st.st_gid);
}

/*
  decompress one file from the command line
*/
static void decompress_file(void)
{
	int fd_in, fd_out = -1, fd_hist = -1;
	i64 expected_size;
	char *tmp, *tmpoutfile, *infilecopy;

	/* make sure infile has an extension. If not, add it
	  * because manipulations may be made to input filename, set local ptr
	*/
	if ((tmp = strrchr(control.infile,'.')) && strcmp(tmp,control.suffix)) {
		infilecopy = malloc(strlen(control.infile)+strlen(control.suffix)+1);
		if (infilecopy==NULL)
			fatal("failed to allocate memory for infile suffix\n");
		else {
			strcpy(infilecopy, control.infile);
			strcat(infilecopy, control.suffix);
		}
	} else
		infilecopy = strdup(control.infile);
	/* regardless, infilecopy has the input filename */

	/* if output name already set, use it */
	if (control.outname)
		control.outfile = strdup(control.outname);
	else {
		/* default output name from infilecopy
		 * test if outdir specified. If so, strip path from filename of
		 * infilecopy, then remove suffix.
		*/
		if (control.outdir && (tmp=strrchr(infilecopy,'/')))
			tmpoutfile=strdup(tmp+1);
		else
			tmpoutfile = strdup(infilecopy);

		/* remove suffix to make outfile name */
		if ((tmp = strrchr(tmpoutfile,'.')) && !strcmp(tmp,control.suffix))
			*tmp='\0';

		control.outfile = malloc((control.outdir==NULL?0:strlen(control.outdir)) +
					  strlen(tmpoutfile) + 1);
		if (!control.outfile)
			fatal("Failed to allocate outfile name\n");

		if (control.outdir) {	/* prepend control.outdir */
			strcpy(control.outfile, control.outdir);
			strcat(control.outfile, tmpoutfile);
		} else
			strcpy(control.outfile, tmpoutfile);
		free(tmpoutfile);
	}
	if (control.flags & FLAG_SHOW_PROGRESS) {
		printf("Output filename is: %s...Decompressing...\n", control.outfile);
		fflush(stdout);
	}

	fd_in = open(infilecopy,O_RDONLY);
	if (fd_in == -1) {
		fatal("Failed to open %s: %s\n",
		      infilecopy,
		      strerror(errno));
	}

	if ((control.flags & FLAG_TEST_ONLY) == 0) {
		if (control.flags & FLAG_FORCE_REPLACE)
			fd_out = open(control.outfile,O_WRONLY|O_CREAT|O_TRUNC,0666);
		else
			fd_out = open(control.outfile,O_WRONLY|O_CREAT|O_EXCL,0666);
		if (fd_out == -1)
			fatal("Failed to create %s: %s\n", control.outfile, strerror(errno));

		if (!(control.flags & FLAG_NO_SET_PERMS))
			preserve_perms(fd_in, fd_out);

		fd_hist = open(control.outfile,O_RDONLY);
		if (fd_hist == -1)
			fatal("Failed to open history file %s\n", control.outfile);
	}

	read_magic(fd_in, &expected_size);
	if (control.flags & FLAG_SHOW_PROGRESS) {
		printf("Decompressing...");
		fflush(stdout);
	}
	runzip_fd(fd_in, fd_out, fd_hist, expected_size);

	/* if we get here, no fatal errors during decompression */
	printf("\rOutput filename is: %s: [OK] - %lld bytes                                 \n", control.outfile, (long long)expected_size);

	if ((control.flags & FLAG_TEST_ONLY) == 0) {
		if (close(fd_hist) != 0 ||
		    close(fd_out) != 0)
			fatal("Failed to close files\n");
	}

	close(fd_in);

	if ((control.flags & (FLAG_KEEP_FILES | FLAG_TEST_ONLY)) == 0) {
		if (unlink(control.infile) != 0)
			fatal("Failed to unlink %s: %s\n", infilecopy, strerror(errno));
	}

	free(control.outfile);
	free(infilecopy);
}

/*
  compress one file from the command line
*/
static void compress_file(void)
{
	int fd_in, fd_out;
	const char *tmp, *tmpinfile; 	/* we're just using this as a proxy for control.infile.
				 	 * Spares a compiler warning
			       		*/
	char header[24];

	memset(header,0,sizeof(header));

	/* is extension at end of infile? */
	if ((tmp = strrchr(control.infile,'.')) && !strcmp(tmp,control.suffix)) {
		printf("%s: already has %s suffix. Skipping...\n", control.infile, control.suffix);
		return;
	}

	if (control.outname) {
		/* check if outname has control.suffix */
		if (*(control.suffix) == '\0') /* suffix is empty string */
			control.outfile = strdup(control.outname);
		else if ((tmp=strrchr(control.outname,'.')) && strcmp(tmp, control.suffix)) {
			control.outfile = malloc(strlen(control.outname) +
						  strlen(control.suffix) + 1);
			if (!control.outfile)
				fatal("Failed to allocate outfile name\n");
			strcpy(control.outfile, control.outname);
			strcat(control.outfile, control.suffix);
			printf("Suffix added to %s.\nFull pathname is: %s\n", control.outname, control.outfile);
		} else	/* no, already has suffix */
			control.outfile = strdup(control.outname);
	} else {
		/* default output name from control.infile
		 * test if outdir specified. If so, strip path from filename of
		 * control.infile
		*/
		if (control.outdir && (tmp=strrchr(control.infile,'/')))
			tmpinfile = tmp + 1;
		else
			tmpinfile = control.infile;

		control.outfile = malloc((control.outdir==NULL?0:strlen(control.outdir)) +
					  strlen(tmpinfile) + strlen(control.suffix) + 1);
		if (!control.outfile)
			fatal("Failed to allocate outfile name\n");

		if (control.outdir) {	/* prepend control.outdir */
			strcpy(control.outfile, control.outdir);
			strcat(control.outfile, tmpinfile);
		} else
			strcpy(control.outfile, tmpinfile);
		strcat(control.outfile, control.suffix);
		if ( control.flags & FLAG_SHOW_PROGRESS )
			printf("Output filename is: %s\n", control.outfile);
	}

	fd_in = open(control.infile,O_RDONLY);
	if (fd_in == -1)
		fatal("Failed to open %s: %s\n", control.infile, strerror(errno));

	if (control.flags & FLAG_FORCE_REPLACE)
		fd_out = open(control.outfile,O_WRONLY|O_CREAT|O_TRUNC,0666);
	else
		fd_out = open(control.outfile,O_WRONLY|O_CREAT|O_EXCL,0666);
	if (fd_out == -1)
		fatal("Failed to create %s: %s\n", control.outfile, strerror(errno));

        if (!(control.flags & FLAG_NO_SET_PERMS))
		preserve_perms(fd_in, fd_out);

	/* write zeroes to 24 bytes at beginning of file */
	if (write(fd_out, header, sizeof(header)) != sizeof(header))
		fatal("Cannot write file header\n");
	rzip_fd(fd_in, fd_out);

	/* write magic at end b/c lzma does not tell us properties until it is done */
	write_magic(fd_in, fd_out);

	if (close(fd_in) != 0 || close(fd_out) != 0)
		fatal("Failed to close files\n");

	if ((control.flags & FLAG_KEEP_FILES) == 0) {
		if (unlink(control.infile) != 0)
			fatal("Failed to unlink %s: %s\n", control.infile, strerror(errno));
	}

	free(control.outfile);
}

/*
 * Returns ram size on linux/darwin.
 */
#ifdef __APPLE__
static i64 get_ram(void)
{
	int mib[2];
	i64 len;
	int *p, ramsize;

	mib[0] = CTL_HW;
	mib[1] = HW_PHYSMEM;
	sysctl(mib, 2, NULL, &len, NULL, 0);
	p = malloc(len);
	sysctl(mib, 2, p, &len, NULL, 0);
	ramsize = *p / 1024; // bytes -> KB

	/* Darwin can't overcommit as much as linux so we return half the ram
	   size to fudge it to use smaller windows */
	ramsize /= 2;
	return ramsize;
}
#else
static i64 get_ram(void)
{
	return (i64)sysconf(_SC_PHYS_PAGES) * (i64)sysconf(_SC_PAGE_SIZE) / 1024;
}
#endif

int main(int argc, char *argv[])
{
	extern int optind;
	int c, i;
	struct timeval start_time, end_time;
	int hours,minutes;
	double seconds,total_time; // for timers
	char *eptr; /* for environment */

	memset(&control, 0, sizeof(control));

	control.flags = FLAG_SHOW_PROGRESS | FLAG_KEEP_FILES;
	control.suffix = ".lrz";
	control.outdir = NULL;

	if (strstr(argv[0], "lrunzip")) {
		control.flags |= FLAG_DECOMPRESS;
	}

	control.compression_level = 7;
	control.ramsize = get_ram() / 104858; /* hundreds of megabytes */
	control.window = 0;
	control.threshold = 0.95;	/* default lzo test compression threshold (level 1) with LZMA compression */
	/* for testing single CPU */
#ifndef NOTHREAD
	control.threads = sysconf(_SC_NPROCESSORS_ONLN);	/* get CPUs for LZMA */
#else
	control.threads = 1;
#endif

	control.nice_val=19;

	/* generate crc table */
	CrcGenerateTable();

	/* Get Preloaded Defaults from lrzip.conf
	 * Look in ., $HOME/.lrzip/, /etc/lrzip.
	 * If LRZIP=NOCONFIG is set, then ignore config
	 */

	eptr = getenv("LRZIP");
	if (eptr == NULL)
		read_config(&control);
	else if (!strstr(eptr,"NOCONFIG"))
		read_config(&control);

	while ((c = getopt(argc, argv, "L:hdS:tVvDfqo:w:nlbMO:T:N:gPz")) != -1) {
		switch (c) {
		case 'L':
			control.compression_level = atoi(optarg);
			if (control.compression_level < 1 || control.compression_level > 9)
				fatal("Invalid compression level (must be 1-9)\n");
			break;
		case 'w':
			control.window = atol(optarg);
			break;
		case 'd':
			control.flags |= FLAG_DECOMPRESS;
			break;
		case 'S':
			control.suffix = optarg;
			break;
		case 'o':
			if (control.outdir)
				fatal("Cannot have -o and -O together\n");
			control.outname = optarg;
			break;
		case 't':
			fatal("integrity checking currently not implemented\n");
			control.flags |= FLAG_TEST_ONLY;
			break;
		case 'f':
			control.flags |= FLAG_FORCE_REPLACE;
			break;
		case 'D':
			control.flags &= ~FLAG_KEEP_FILES;
			break;
		case 'v':
			/* set verbosity flag */
			if ( !(control.flags & FLAG_VERBOSITY) && !(control.flags & FLAG_VERBOSITY_MAX) )
				control.flags |= FLAG_VERBOSITY;
			else if ( (control.flags & FLAG_VERBOSITY) ) {
				control.flags &= ~FLAG_VERBOSITY;
				control.flags |= FLAG_VERBOSITY_MAX;
			}
			break;
		case 'q':
			control.flags &= ~FLAG_SHOW_PROGRESS;
			break;
		case 'V':
			printf("lrzip version %d.%d%d\n",
				LRZIP_MAJOR_VERSION, LRZIP_MINOR_VERSION, LRZIP_MINOR_SUBVERSION);
			exit(0);
			break;
		case 'l':
			if (control.flags & FLAG_NOT_LZMA)
				fatal("Can only use one of -l, -b, -g, -z or -n\n");
			control.flags |= FLAG_LZO_COMPRESS;
			break;
		case 'b':
			if (control.flags & FLAG_NOT_LZMA)
				fatal("Can only use one of -l, -b, -g, -z or -n\n");
			control.flags |= FLAG_BZIP2_COMPRESS;
			break;
		case 'n':
			if (control.flags & FLAG_NOT_LZMA)
				fatal("Can only use one of -l, -b, -g, -z or -n\n");
			control.flags |= FLAG_NO_COMPRESS;
			break;
		case 'M':
			control.compression_level = 9;
			control.window = control.ramsize / 10 * 9 ? : 1;
			break;
		case 'O':
			if (control.outname)	/* can't mix -o and -O */
				fatal("Cannot have options -o and -O together\n");
			control.outdir=malloc(strlen(optarg)+2);
			if (control.outdir == NULL)
				fatal("Failed to allocate for outdir\n");
			strcpy(control.outdir,optarg);
			if (strcmp(optarg+strlen(optarg)-1,"/")) 	/* need a trailing slash */
				strcat(control.outdir,"/");
			break;
		case 'T':
			/* invert argument, a threshold of 1 means that the compressed result can be
			 * 90%-100% of the sample size
			*/
			control.threshold = atoi(optarg);
			if (control.threshold < 0 || control.threshold > 10)
				fatal("Threshold value must be between 0 and 10\n");
			control.threshold = 1.05-control.threshold / 20;
			break;
		case 'N':
			control.nice_val = atoi(optarg);
			if (control.nice_val < -20 || control.nice_val > 19)
				fatal("Invalid nice value (must be -20..19)\n");
			break;
		case 'g':
			if (control.flags & FLAG_NOT_LZMA)
				fatal("Can only use one of -l, -b, -g, -z or -n\n");
			control.flags |= FLAG_ZLIB_COMPRESS;
			break;
		case 'P':
			control.flags |= FLAG_NO_SET_PERMS;
			break;
		case 'z':
			if (control.flags & FLAG_NOT_LZMA)
				fatal("Can only use one of -l, -b, -g, -z or -n\n");
			control.flags |= FLAG_ZPAQ_COMPRESS;
			break;
		default:
		case 'h':
			usage();
			return -1;
		}
	}

	argc -= optind;
	argv += optind;

	if (control.outname && argc > 1)
		fatal("Cannot specify output filename with more than 1 file\n");

	if (( (control.flags & FLAG_VERBOSITY) || (control.flags & FLAG_VERBOSITY_MAX) )
			&& !(control.flags & FLAG_SHOW_PROGRESS)) {
		printf("Cannot have -v and -q options. -v wins.\n");
		control.flags |= FLAG_SHOW_PROGRESS;
	}

	if (argc < 1) {
		usage();
		exit(1);
	}

	if (control.window > control.ramsize)
		printf("Compression window has been set to larger than ramsize, proceeding at your request. If you did not mean this, abort now.\n");

	/* The control window chosen is the largest that will not cause
	   massive swapping on the machine (60% of ram). Most of the pages
	   will be shared by lzma even though it uses just as much ram itself
	   */
	if (!control.window)
		control.window = control.ramsize / 3 * 2 ? : 1;

	/* malloc limited to 2GB on 32bit */
	if (sizeof(long) == 4 && control.window > 20) {
		control.window = 20;
		if (control.flags & (FLAG_VERBOSITY | FLAG_VERBOSITY_MAX))
			printf("Limiting control window to 2GB due to 32bit limitations.\n");
	}

	/* OK, if verbosity set, print summary of options selected */
	if ((control.flags & FLAG_VERBOSITY) || (control.flags & FLAG_VERBOSITY_MAX)) {
		printf("The following options are in effect for this %s.\n",
			control.flags & FLAG_DECOMPRESS ? "DECOMPRESSION" : "COMPRESSION");
		if (LZMA_COMPRESS(control.flags))
			printf("Threading is %s. Number of CPUs detected: %lu\n", control.threads > 1 ?"ENABLED":"DISABLED",
				control.threads);
		printf("Nice Value: %d\n", control.nice_val);
		if (control.flags & FLAG_SHOW_PROGRESS)
			printf("Show Progress\n");
		if (control.flags & FLAG_VERBOSITY)
			printf("Verbose\n");
		else if (control.flags & FLAG_VERBOSITY_MAX)
			printf("Max Verbosity\n");
		if (control.flags & FLAG_FORCE_REPLACE)
			printf("Overwrite Files\n");
		if (!(control.flags & FLAG_KEEP_FILES))
			printf("Remove input files on completion\n");
		if (control.outdir)
			printf("Output Directory Specified: %s\n", control.outdir);
		else if (control.outname)
			printf("Output Filename Specified: %s\n", control.outname);

		/* show compression options */
		if (!(control.flags & FLAG_DECOMPRESS)) {
			printf("Compression mode is: ");
			if (LZMA_COMPRESS(control.flags))
				printf("LZMA. LZO Test Compression Threshold: %.f\n",
				       (control.threshold < 1 ? 20 - control.threshold * 20 : 1));
			else if (control.flags & FLAG_LZO_COMPRESS)
				printf("LZO\n");
			else if (control.flags & FLAG_BZIP2_COMPRESS)
				printf("BZIP2. LZO Test Compression Threshold: %.f\n",
				       (control.threshold < 1 ? 20 - control.threshold * 20 : 1));
			else if (control.flags & FLAG_ZLIB_COMPRESS)
				printf("GZIP\n");
			else if (control.flags & FLAG_ZPAQ_COMPRESS)
				printf("ZPAQ. LZO Test Compression Threshold: %.f\n",
				       (control.threshold < 1 ? 20 - control.threshold * 20 : 1));
			else if (control.flags & FLAG_NO_COMPRESS)
				printf("RZIP\n");
			printf("Compression Window: %lld = %lldMB\n", control.window, control.window * 104858);
			printf("Compression Level: %d\n", control.compression_level);
		}
		printf("\n");
	}

	if (setpriority(PRIO_PROCESS, 0, control.nice_val) == -1)
		fatal("Unable to set nice value\n");

	for (i = 0; i < argc; i++) {
		control.infile = argv[i];
		gettimeofday(&start_time, NULL);

		if (control.flags & (FLAG_DECOMPRESS | FLAG_TEST_ONLY))
			decompress_file();
		else
			compress_file();

		/* compute total time */
		gettimeofday(&end_time, NULL);
		total_time = (end_time.tv_sec + (double)end_time.tv_usec / 1000000) -
			      ( start_time.tv_sec + (double)start_time.tv_usec / 1000000);
		hours = (int)total_time / 3600;
		minutes = (int)(total_time - hours * 3600) / 60;
		seconds = total_time - hours * 60 - minutes * 60;
		if (control.flags & FLAG_SHOW_PROGRESS)
			printf("Total time: %02d:%02d:%06.3f\n", hours, minutes, seconds);
	}

	return 0;
}