/*
 * stereojoin.cpp
 * 
 * PNG stereo image joiner/splitter. 
 * 
 * Copyright (c) 2003 by Wolfgang Wieser (wwieser at gmx doot de) 
 * 
 * This file may be distributed and/or modified under the terms of the 
 * GNU General Public License version 2 as published by the Free Software 
 * Foundation. 
 * 
 * This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 * 
 */

//------------------------------------------------------------------------------
// StereoJoin does the following: 
// 
// JOINING:
//  You have a series of frames l000.png,l001.png,...
//  and another series r000.png,r001.png,... 
//  StereoJoin combines these into a series of frames (f000.png,...) 
//  where the l*-frames are left of the r*-frames (and beween them a 
//  black gap of tunable width). 
// 
// SPLITTING:
//  Just the reverse of joining. You must specify the correct size 
//  of the black gap and get the left and right frames (equally-sized) 
//  out of the joined frames. 
// 
// Requires:
//   libpng must be installed (including header file <png.h>). 
// Compile stereojoin using: 
//   gcc -s -O2 -fno-rtti -fno-exceptions stereojoin.cc -o stereojoin -lpng
//   ln -s stereojoin stereosplit
// Install (the option "-d" keeps the symlink):
//   cp -d stereojoin stereosplit <where_you_want_them>
//
// KNOWN BUGS: 
//  - Probably does not work when fed with PNGs with depth!=8 and !=16. 
//------------------------------------------------------------------------------

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <ctype.h>
#include <errno.h>
#include <assert.h>
#include <png.h>

#include "stereojoin.h"

char *prg_name=NULL;

#define VERSION_STR "0.4"


static void PrintHelp()
{
	printf("Usage: %s [options]\n"
		"  -h --help  print this\n"
		"  --version  pint version info\n"
		"  -q         quiet operation\n"
		"  -J         join left + right -> joined (default when called stereojoin)\n"
		"  -S         split joined -> left + right (default wenn called stereosplit)\n"
		"  -d=NUM   number of black pixels between right & left (default: 8)\n"
		"  -s=NUM   start frame number (default: 0)\n"
		"  -n=NUM   number of frames (default: autodetect: until 3 frames fail)\n"
		"  -j=NUM   process every NUM frames (defalt: 1)\n"
		"  -z=NUM   set PNG compression level (0..9; default 6)\n"
		"  -l=PAT   left frame file name pattern (default: l%%07d.png)\n"
		"  -r=PAT   right frame file name pattern (default: r%%07d.png)\n"
		"  -o=PAT   joined frame file name pattern (default: f%%07d.png)\n"
		"PAT: path and file name with %%d/%%x/%%X placeholder\n"
		"     Only PNG format supported.\n" 
		"StereoJoin (c) 2003 by Wolgang Wieser\n",prg_name);
}

static void PrintVersion()
{
	printf("%s version %s\n",prg_name,VERSION_STR);
}


static void *CheckMalloc(void *ptr)
{
	if(!ptr)
	{  fprintf(stderr,"Allocation failure.\n");  abort();  }
	return(ptr);
}


// Return value: 0 -> okay; 1 -> illegal pattern
static int CheckFramePattern(const char *ptr)
{
	do {
		if(!ptr)  break;
		ptr=strchr(ptr,'%');
		if(!ptr)  break;
		++ptr;
		while(isdigit(*ptr))  ++ptr;
		if(*ptr!='x' && *ptr!='d' && *ptr!='X')
		{  ptr=NULL;  break;  }
		if(strchr(ptr,'%'))
		{  ptr=NULL;  break;  }
	}  while(0);
	return(ptr ? 0 : 1);
}



static int ReadInt(const char *str,const char *desc,int *dest,
	int nonneg)
{
	if(!*str)
	{  fprintf(stderr,"%s: value for %s omitted.\n",
		prg_name,desc);  return(1);  }
	char *end;
	int val=strtol(str,&end,0);
	if(*end || (val<0 && nonneg) || (val<=0 && nonneg==2))
	{  fprintf(stderr,"%s: illegal value for %s.\n",
		prg_name,desc);  return(1);  }
	*dest=val;
	return(0);
}

static int ReadPat(const char *str,const char *desc,char **dest)
{
	if(!*str)
	{  fprintf(stderr,"%s: value for %s omitted.\n",
		prg_name,desc);  return(1);  }
	
	// Check pattern: 
	if(CheckFramePattern(str))
	{  fprintf(stderr,"%s: illegal value for %s (\"%s\").\n",
		prg_name,desc,str);  return(1);  }
	
	if(*dest)  free(*dest);
	*dest=(char*)CheckMalloc(strdup(str));
	
	return(0);
}


JoinData::JoinData()
{
	mode=MNONE;
	f0=0;
	nframes=-1;
	fjump=1;
	black_delta=8;
	quiet=0;
	right_pat=NULL;
	left_pat=NULL;
	join_pat=NULL;
	compression_level=6;
}

JoinData::~JoinData()
{
	if(right_pat)  free(right_pat);
	if(left_pat)   free(left_pat);
	if(join_pat)   free(join_pat);
}


void png_error_func(png_structp,png_const_charp str)
{
	fprintf(stderr,"%s: PNG error: %s\n",prg_name,str);
}

void png_warning_func(png_structp,png_const_charp str)
{
	fprintf(stderr,"%s: PNG warning: %s\n",prg_name,str);
}


int PNGOutFile::WriteRow(char *row_buf)
{
	if(setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr,"(PNG error)\n");
		return(1);
	}
	
	png_write_row(png_ptr,(png_byte*)row_buf);
	
	return(0);
}


int PNGOutFile::Open(int compression_level)
{
	f=fopen(name,"w");
	if(!f)
	{  fprintf(stderr,"%s: while opening %s: %s\n",
		prg_name,name,strerror(errno));  return(1);  }
	
	png_ptr=png_create_write_struct(PNG_LIBPNG_VER_STRING,
		this,&png_error_func,&png_warning_func);
	CheckMalloc(png_ptr);
	
	info_ptr=png_create_info_struct(png_ptr);
	CheckMalloc(info_ptr);
	
	if(setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr,"(PNG error)\n");
		return(1);
	}
	
	png_init_io(png_ptr,f);
	
	png_set_compression_level(png_ptr,compression_level);
	
	return(0);
}


int PNGOutFile::WriteHeader(PNGInFile *left,int totwidth)
{
	if(setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr,"(PNG error)\n");
		return(1);
	}

	png_set_IHDR(png_ptr,info_ptr,totwidth,left->height,
		left->bit_depth,left->color_type,PNG_INTERLACE_NONE,
		PNG_COMPRESSION_TYPE_DEFAULT,PNG_FILTER_TYPE_DEFAULT);
	
	png_write_info(png_ptr,info_ptr);
	
	return(0);
}

int PNGOutFile::WriteTail()
{
	if(setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr,"(PNG error)\n");
		return(1);
	}
	
	png_write_end(png_ptr,info_ptr);
	
	return(0);
}


PNGOutFile::PNGOutFile(const char *pat,int frame)
{
	f=NULL;
	
	png_ptr=NULL;
	info_ptr=NULL;
	
	char tmp;
	int len=snprintf(&tmp,0,pat,frame);
	if(len<=0)  abort();
	name=(char*)CheckMalloc(malloc(len+1));
	snprintf(name,len+1,pat,frame);
}

PNGOutFile::~PNGOutFile()
{
	if(png_ptr)
	{  png_destroy_write_struct(&png_ptr,&info_ptr);  }
	if(f)
	{  fclose(f);  }
	if(name)
	{  free(name);  }
}


int PNGInFile::ReadRow(char *dest)
{
	if(setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr,"(PNG error)\n");
		return(1);
	}
	
	png_read_row(png_ptr,(png_byte*)dest,NULL);
	
	return(0);
}

int PNGInFile::Open()
{
	f=fopen(name,"r");
	if(!f)
	{  fprintf(stderr,"%s: while opening %s: %s\n",
		prg_name,name,strerror(errno));  return(1);  }
	
	int okay=0;
	do {
		if(fread(header,1,8,f)!=8)  break;
		if(png_sig_cmp((png_byte*)header,0,8))  break;
		okay=1;
	} while(0);
	if(!okay)
	{
		fprintf(stderr,"%s: %s is not a valid PNG file.\n",
			prg_name,name);
		return(1);
	}
	
	png_ptr=png_create_read_struct(
		PNG_LIBPNG_VER_STRING,this,
		&png_error_func,&png_warning_func);
	CheckMalloc(png_ptr);
	
	info_ptr=png_create_info_struct(png_ptr);
	if(!info_ptr)
	{  CheckMalloc(info_ptr);  }
	
	png_infop end_info=png_create_info_struct(png_ptr);
	if(!end_info)
	{  CheckMalloc(end_info);  }
	
	if(setjmp(png_jmpbuf(png_ptr)))
	{
		fprintf(stderr,"(PNG error)\n");
		return(1);
	}
	
	png_init_io(png_ptr,f);
	
	png_set_sig_bytes(png_ptr,8);
	
	png_read_info(png_ptr, info_ptr);
	
	int interlace_type;
	png_get_IHDR(png_ptr,info_ptr,&width, &height,
		&bit_depth,&color_type,&interlace_type,NULL,NULL);
	
	if(interlace_type!=PNG_INTERLACE_NONE)
	{  fprintf(stderr,"%s: %s: interlacing not supported\n",
		prg_name,name);  return(1);  }
	
	if(color_type!=PNG_COLOR_TYPE_GRAY && 
	   color_type!=PNG_COLOR_TYPE_GRAY_ALPHA && 
	   color_type!=PNG_COLOR_TYPE_RGB && 
	   color_type!=PNG_COLOR_TYPE_RGB_ALPHA)
	{  fprintf(stderr,"%s: %s unsupported color type %d\n",
		prg_name,name,color_type);  return(1);  }
	
	//if(color_type==PNG_COLOR_TYPE_PALETTE)
	//{  png_set_palette_to_rgb(png_ptr);  }
	
	channels=png_get_channels(png_ptr,info_ptr);
	rowbytes=png_get_rowbytes(png_ptr,info_ptr);
	
	return(0);
}


PNGInFile::PNGInFile(const char *pat,int frame)
{
	f=NULL;
	width=0;
	height=0;
	bit_depth=-1;
	color_type=0;
	channels=0;
	rowbytes=0;
	
	png_ptr=NULL;
	info_ptr=NULL;
	end_info=NULL;
	
	char tmp;
	int len=snprintf(&tmp,0,pat,frame);
	if(len<=0)  abort();
	name=(char*)CheckMalloc(malloc(len+1));
	snprintf(name,len+1,pat,frame);
}

PNGInFile::~PNGInFile()
{
	if(png_ptr)
	{  png_destroy_read_struct(&png_ptr,&info_ptr,&end_info);  }
	if(f)
	{  fclose(f);  }
	if(name)
	{  free(name);  }
}


// Returns 1 if the passed mem chunk is filled with zeros. 
static inline bool mem_test_zero(const char *_d,size_t len)
{
	unsigned char *d=(unsigned char*)_d;
	const unsigned char *e=d+len;
	for(;d<e;d++)
	{  if(*d)  return(0);  }
	return(1);
}


int DoJoinFrame(JoinData *d,int frame)
{
	PNGInFile left_file(d->left_pat,frame);
	PNGInFile right_file(d->right_pat,frame);
	PNGOutFile out_file(d->join_pat,frame);
	
	// Open input files: 
	if(left_file.Open() || right_file.Open())
	{  return(1);  }
	
	if(	left_file.height!=right_file.height || 
		left_file.bit_depth!=right_file.bit_depth || 
		left_file.color_type!=right_file.color_type ||
		left_file.channels!=right_file.channels)
	{  fprintf(stderr,"%s: Input files %s and %s incompatible.\n",
		prg_name,left_file.name,right_file.name);  return(1);  }
	
	assert((left_file.rowbytes*8)%left_file.width==0);
	size_t bits_per_pixel=(left_file.rowbytes*8)/left_file.width;
	size_t middlelen=d->black_delta*bits_per_pixel/8;
	if((d->black_delta*bits_per_pixel) % 8)
	{  fprintf(stderr,"%s: illegal combination: "
		"black delta=%d, depth=%d\n",
		prg_name,d->black_delta,left_file.bit_depth);  return(1);  }
	
	if(out_file.Open(d->compression_level))
	{  return(1);  }
	
	// Allocate mem for a image row of the dest image. 
	size_t lbufsize=left_file.rowbytes+right_file.rowbytes;
	//fprintf(stderr,"lwidth=%d, rwidth=%d, bit_depth=%d, lbufsize=%u\n",
	//	left_file.width,right_file.width,left_file.bit_depth,lbufsize);
	assert((left_file.width+right_file.width)*bits_per_pixel==8*lbufsize);
	lbufsize+=middlelen;
	char *row_buf=(char*)CheckMalloc(malloc(lbufsize));
	
	int rv=1;
	
	char *left_buf=row_buf;
	char *middle_buf=left_buf+left_file.rowbytes;
	char *right_buf=middle_buf+middlelen;
	if(d->black_delta)
	{  memset(middle_buf,0,middlelen);  }
	
	if(out_file.WriteHeader(&left_file,
		left_file.width+d->black_delta+right_file.width))
	{  goto doreturn;  }
	
	// Process each line: 
	for(unsigned int yy=0; yy<left_file.height; yy++)
	{
		if(left_file.ReadRow(left_buf))  goto doreturn;
		if(right_file.ReadRow(right_buf))  goto doreturn;
		
		if(out_file.WriteRow(row_buf))  goto doreturn;
	}
	
	if(out_file.WriteTail())
	{  goto doreturn;  }
	
	rv=0;
doreturn:;
	if(row_buf)
	{  free(row_buf);  }
	return(rv);
}


int DoSplitFrame(JoinData *d,int frame)
{
	PNGInFile in_file(d->join_pat,frame);
	PNGOutFile left_file(d->left_pat,frame);
	PNGOutFile right_file(d->right_pat,frame);
	
	// Open input file: 
	if(in_file.Open())
	{  return(1);  }
	
	int split_width=-1;
	do {
		if(in_file.width<(unsigned int)d->black_delta)  break;
		int tmpw=in_file.width-d->black_delta;
		if(tmpw%2)  break;
		// This will split into equally-sized files (because of check above). 
		split_width=tmpw/2;
	} while(0);
	if(split_width<1)
	{  fprintf(stderr,"%s: illegal combination: "
		"black delta=%d, width=%d\n",
		prg_name,d->black_delta,int(in_file.width));  return(1);  }
	
	// See how many bytes of black delta have to be cut out: 
	assert((in_file.rowbytes*8)%in_file.width==0);
	size_t bits_per_pixel=(in_file.rowbytes*8)/in_file.width;
	size_t middlelen=d->black_delta*bits_per_pixel/8;
	if((d->black_delta*bits_per_pixel) % 8)
	{  fprintf(stderr,"%s: illegal combination: "
		"black delta=%d, depth=%d\n",
		prg_name,d->black_delta,in_file.bit_depth);  return(1);  }
	
	size_t lr_row_bytes=split_width*bits_per_pixel/8;
	if((split_width*bits_per_pixel) % 8)
	{  fprintf(stderr,"%s: illegal combination: "
		"split width=%d, depth=%d\n",
		prg_name,split_width,in_file.bit_depth);  return(1);  }
	
	// Open output files: 
	if(left_file.Open(d->compression_level) || 
	   right_file.Open(d->compression_level) )
	{  return(1);  }
	
	// Allocate mem for a image row of the source image. 
	char *row_buf=(char*)CheckMalloc(malloc(in_file.rowbytes));
	
	int rv=1;
	int warned_zero=0;
	
	char *left_buf=row_buf;
	char *middle_buf=left_buf+lr_row_bytes;
	char *right_buf=middle_buf+middlelen;
	
	if(left_file.WriteHeader(&in_file,split_width) || 
	   right_file.WriteHeader(&in_file,split_width) )
	{  goto doreturn;  }
	
	// Process each line: 
	for(unsigned int yy=0; yy<in_file.height; yy++)
	{
		if(in_file.ReadRow(row_buf))  goto doreturn;
		
		if(left_file.WriteRow(left_buf))  goto doreturn;
		if(right_file.WriteRow(right_buf))  goto doreturn;
		
		if(!warned_zero && !mem_test_zero(middle_buf,middlelen))
		{  fprintf(stderr,"%s: Frame %d: warning: middle part not entirely "
			"black.\n",prg_name,frame);  warned_zero=1;  }
	}
	
	if(left_file.WriteTail())  goto doreturn;
	if(right_file.WriteTail())  goto doreturn;
	
	rv=0;
doreturn:;
	if(row_buf)
	{  free(row_buf);  }
	return(rv);
}


int DoProcess(JoinData *d)
{
	if(!d->quiet)
	{
		char tmp[24];
		if(d->nframes<0)
		{  strcpy(tmp,"*");  }
		else
		{  snprintf(tmp,24,"%d",d->f0+d->nframes);  }
		if(d->mode==MJOIN)
		{  printf("%s: Joining [ %s | %s ] --> %s  (%d..%s, j=%d)\n",
			prg_name,d->left_pat,d->right_pat,d->join_pat,
			d->f0,tmp,d->fjump);  }
		else if(d->mode==MSPLIT)
		{  printf("%s: Splitting %s -> [ %s | %s ] (%d..%s, j=%d)\n",
			prg_name,d->join_pat,d->left_pat,d->right_pat,
			d->f0,tmp,d->fjump);  }
		else
		{  assert(0);  abort();  }
	}
	
	int errors=0;
	int failed_in_seq=0;
	int need_nl=0;
	for(int frame=d->f0; (frame<d->nframes+d->f0 || d->nframes<0); 
		frame+=d->fjump)
	{
		if(!d->quiet)
		{
			printf("\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\bFrame %7d...",frame);
			fflush(stdout);
			need_nl=1;
		}
		
		int rv;
		if(d->mode==MJOIN)
		{  rv=DoJoinFrame(d,frame);  }
		else if(d->mode==MSPLIT)
		{  rv=DoSplitFrame(d,frame);  }
		else
		{  assert(0);  abort();  }
		if(rv)
		{
			need_nl=0;
			++errors;
			++failed_in_seq;
			if(failed_in_seq>=3)
			{
				fprintf(stderr,"%s: %d frames failed in sequence. "
					"Giving up.\n",prg_name,failed_in_seq);
				break;
			}
		}
		else
		{  failed_in_seq=0;  }
	}
	if(need_nl && !d->quiet)
	{  printf(" OK\n");  fflush(stdout);  need_nl=0;  }
	
	return(errors ? 1 : 0);
}


int main(int argc,char **arg)
{
	prg_name=strrchr(arg[0],'/');
	if(prg_name)  ++prg_name;
	else  prg_name=arg[0];
	
	JoinData jd;
	
	int errors=0;
	for(int i=1; i<argc; i++)
	{
		if(*arg[i]!='-')
		{
			fprintf(stderr,"%s: illegal arg \"%s\".\n",prg_name,arg[i]);
			++errors;
			continue;
		}
		if(!strcmp(arg[i],"--help") || !strcmp(arg[i],"-help"))
		{  PrintHelp();  return(0);  }
		if(!strcmp(arg[i],"--version") || !strcmp(arg[i],"-version"))
		{  PrintVersion();  return(0);  }
		if(!strncmp(arg[i],"-d=",3))
		{  errors+=ReadInt(arg[i]+3,"-d (number of pixels)",
			&jd.black_delta,1);  }
		else if(!strncmp(arg[i],"-s=",3))
		{  errors+=ReadInt(arg[i]+3,"-s (startframe)",&jd.f0,1);  }
		else if(!strncmp(arg[i],"-n=",3))
		{  errors+=ReadInt(arg[i]+3,"-n (number of frames)",&jd.nframes,1);  }
		else if(!strncmp(arg[i],"-j=",3))
		{  errors+=ReadInt(arg[i]+3,"-j (frame jump)",&jd.fjump,2);  }
		else if(!strncmp(arg[i],"-l=",3))
		{  errors+=ReadPat(arg[i]+3,"-l (left frames)",&jd.left_pat);  }
		else if(!strncmp(arg[i],"-r=",3))
		{  errors+=ReadPat(arg[i]+3,"-r (right frames)",&jd.right_pat);  }
		else if(!strncmp(arg[i],"-j=",3))
		{  errors+=ReadPat(arg[i]+3,"-j (joined frames)",&jd.join_pat);  }
		else if(!strncmp(arg[i],"-z=",3))
		{  errors+=ReadInt(arg[i]+3,"-z (compression level)",
			&jd.compression_level,1);  }
		else if(strchr(arg[i],'='))
		{
			fprintf(stderr,"%s: illegal parameter \"%s\".\n",
				prg_name,arg[i]);
			++errors;
			continue;
		}
		else for(char *c=arg[i]+1; *c; c++)
		{
			switch(*c)
			{
				case 'h':  PrintHelp();  return(0);  break;
				case 'q':  jd.quiet=1;      break;
				case 'J':  jd.mode=MJOIN;   break;
				case 'S':  jd.mode=MSPLIT;  break;
				default:
					fprintf(stderr,"%s: illegal option %c in "
						"arg \"%s\".\n",prg_name,*c,arg[i]);
					return(1);
			}
		}
	}
	
	if(errors)
	{  return(1);  }
	
	if(jd.mode==MNONE)
	{
		if(!strcmp(prg_name,"stereojoin"))
		{  jd.mode=MJOIN;  }
		else if(!strcmp(prg_name,"stereosplit"))
		{  jd.mode=MSPLIT;  }
		else
		{  fprintf(stderr,"%s: Please use either -J or -S or call me "
			"either \"stereojoin\" or \"stereosplit\".\n",prg_name);
			++errors;  }
	}
	
	if(errors)
	{  return(1);  }
	
	if(!jd.right_pat)
	{  jd.right_pat=(char*)CheckMalloc(strdup("r%07d.png"));  }
	if(!jd.left_pat)
	{  jd.left_pat=(char*)CheckMalloc(strdup("l%07d.png"));  }
	if(!jd.join_pat)
	{  jd.join_pat=(char*)CheckMalloc(strdup("f%07d.png"));  }
	
	return(DoProcess(&jd));
}