#include "ninox.h"

/*
 * Compress .fta archive files in the most efficient means possible
 *
 * Usage:
 *		./compress -o archive.ftz file1.fta file2.fta ...
 */

static int Compress(char *archive_name, char **input_list, int count);
static int DeCompress(char *archive_fname);

static int AddDirectory(char *dir, char *list[], int *count);
static FILE * StartOutput(char *name);
static int Error(int,char *,...);
static int _do_compress(struct Image *img, char *archive_fname, FILE *out);
static int _encode_frame(struct Image *img, char *buffer);
static int _encode_frame_8(struct Image *img, unsigned char *buffer);
static int _encode_frame_16(struct Image *img, unsigned short *buffer);
static int _write_output(struct Image *img, char *archive_fname, unsigned char *buffer, int nbytes, FILE *out);

static int _ftz_decode_frame_8(struct Image *img, unsigned char *buffer, int zbytes);
static int _ftz_decode_frame_16(struct Image *img, unsigned short *buffer, int zbytes);

static char *Output = NULL;
static FILE *Output_z = NULL;

#define ZRLE   1
#define Z12_16 2

#ifdef FTACOMPRESS_STANDALONE
void
Usage()
	{
	Print("Usage: compress -o archive.ftz file1.fta file2.fta ...\n");
	Print("Usage: compress -d archive.ftz  ...\n");
	Print("Usage: compress -z -o ... (additionally, use bzip2 to further compress the archive)\n");
	}

int
main(int argc, char *argv[])
	{
	char *archive_name = NULL;
	char *input_list[128],*tmp_archive_name;
	int i,count,decompress=0,do_bzip=0;

	count=0;
	for(i=1; i<argc; ++i) {
	   if (!strcmp(argv[i],"-z")) { do_bzip=1; }
	   else if (!strcmp(argv[i],"-o")) { ++i; archive_name = argv[i]; }
	   else if (!strcmp(argv[i],"-d")) { ++i; archive_name = argv[i]; decompress=1; }
	   else {
		if (decompress)
		   Error(1,"Cannot give source files when decompressing\n");

		if (isDirectory(argv[i]))
		     AddDirectory(argv[i],input_list,&count);
	   	else input_list[count++] = argv[i];
		}
	   }

	if (! archive_name || (decompress==0 && count==0)) {
	   Usage();
	   exit(1);
	   }

	if (decompress) 
	   DeCompress(archive_name);
	else {
	   Compress(archive_name,input_list,count);
	   if (do_bzip && isFile(archive_name)) {
		char cmd[1024];
		Print("bzipping %s...\n",archive_name);
		sprintf(cmd,"/usr/bin/bzip2 -f %s",archive_name);
		system(cmd);
		}
	   }

	exit(0);
	}
#endif

struct __d {
   char *name;
   int val;
   };

static int
dircomp(const void *a, const void *b)
	{
	struct __d *p1 = (struct __d *)a;
	struct __d *p2 = (struct __d *)b;
	return p1->val - p2->val;
	}

static int
AddDirectory(char *dir, char *list[], int *count)
	{
	DIR *D;
	struct dirent *e;
	struct __d dlist[32];
	int i,dcount=0,len=strlen(dir);

	D = opendir(dir);
	if (!D) Error(1,"Cannot open directory '%s'\n",dir);

	if (dir[len-1] == '/' || dir[len-1]=='\\')
	   dir[len-1]=0;

	while(e = readdir(D)) {
	   if (e->d_name[0] == '.') continue;
	   if (strlen(e->d_name) < 5) continue;

	   if (!strcmp(e->d_name + strlen(e->d_name)-4, ".fta")) {
		dlist[dcount].name = Malloc(strlen(dir) + strlen(e->d_name) + 2);
		sprintf(dlist[dcount].name,"%s/%s",dir,e->d_name);
		dcount++;
		}
	   }
	closedir(D);

	// Set the values for sorting
	for(i=0; i<dcount; ++i) {
	   char *ptr = dlist[i].name;
	   int val=0;
	   while(*ptr) {
	 	if (isdigit(*ptr)) val = (val*10) + (*ptr - '0');
		++ptr;
		}
	   dlist[i].val = val;
	   }

	// Sort the names
	qsort(dlist,dcount,sizeof(struct __d),dircomp);

	for(i=0; i<dcount; ++i) {
	   list[*count] = dlist[i].name;
	   ++(*count);
	   }

	return 1;
	}
	
static int
Compress(char *archive_name, char **input_list, int count)
	{
	char *tmp_archive_name;
	int i,err;
	extern int InputHistoStretch;
	FILE *out;

	//printf("Compressing %d files to '%s'\n",count,archive_name);

	tmp_archive_name = malloc(strlen(archive_name)+10);
	sprintf(tmp_archive_name,"%s.tmp",archive_name);

	out = StartOutput(tmp_archive_name);
	if (! out)
	   Error(1,"Error opening output '%s'\n",tmp_archive_name);

	InputHistoStretch = 0;

	for(i=0; i<count; ++i) {
	   int n=0;
	   FILE *in = OpenArchive(input_list[i]);
	   if (in == NULL) Error(1,"Error opening input '%s'\n",input_list[i]);

	   Print("\nCompressing %s -> %s\n",input_list[i],tmp_archive_name);
	   while(1) {
  		struct Image *img = ArchiveLoadNextImage(in,&err,ARCHIVE_LOAD);

           	if (err<0) {
                   Print("Fatal error reading archive '%s'\n",input_list[i]);
                   exit(1);
                   }
            	else if (err==0) {
                   break;
                   }
           	else if (err==1) {
                   if (!_do_compress(img,input_list[i],out))
			Error(1,"Error during compress, abort!\n");
                   }
           	else Error(1,"Unknown return value %d\n",err);

		if (n && (n%50)==0) { printf("."); fflush(stdout); }
		++n;
		}
           }

	fflush(out);
	fclose(out);
	Print("Finished, renaming %s -> %s\n",tmp_archive_name,archive_name);
	unlink(archive_name);
	rename(tmp_archive_name,archive_name);
	return(0);
	}

	
static int
Error(int errval, char *fmt, ...)
	{
        int n, size = 100;
        char *p;
        va_list ap;
        static int start_of_line = 1;
        FILE *fd = stdout;

	if (Output_z) fclose(Output_z);
	if (! errval) exit(errval);

	Print("Error %d occurred while writing to '%s'\n",errval,Output);
	unlink(Output);

        if ((p = (char *)malloc (size)) == NULL) {
            fprintf(fd,"Print: Out of memory!\n");
            fflush(fd);
            exit(1);
            }

        while (1) {
            /* Try to print in the allocated space. */
            va_start(ap, fmt);
            n = vsnprintf (p, size, fmt, ap);
            va_end(ap);

            /* If that worked, return the string. */
            if (n > -1 && n < size) {
                for(n=0; n<size && p[n]; ++n) {
                   if (start_of_line && CurrentFile) {fprintf(fd,"%s: ",CurrentFile); fflush(fd);}
                   if (putc(p[n],fd) == EOF) exit(1);
                   if (p[n] == '\n') start_of_line = 1;
                   else start_of_line = 0;
                   }

                free(p);
                fflush(fd);
                if (feof(fd)) exit(1);
                exit(errval);
                }

            /* Else try again with more space. */
            if (n > -1)    /* glibc 2.1 */
               size = n+1; /* precisely what is needed */
            else           /* glibc 2.0 */
               size *= 2;  /* twice the old size */
            if ((p = realloc (p, size)) == NULL) {
                fprintf(fd,"Print: Out of memory!\n"); fflush(fd);
                exit(1);
                }
             }

        // notreached
        fflush(fd);
        exit(errval);
        }

static FILE *
StartOutput(char *name)
	{
	Output_z = fopen(name,"wb");
	if (Output_z == NULL) return NULL;

	Output = strdup(name);
	return Output_z;
	}

static int
_do_compress(struct Image *img, char *archive_fname, FILE *out)
	{
	static char *buffer = NULL;
	static int bsize = 0;
	int width = img->width;
	int height = img->height;
	int bpp = img->depth/8;
	int nbytes = width * height * bpp;

	//Print("[%s]: ",img->dst_fname);

	if (bsize>0 && bsize != nbytes) {
	   Print("Buffer size change, not supported - abort!\n");
	   return 0;
	   }

	// Init the buffer
	if (buffer == NULL) buffer = Malloc(nbytes);

	nbytes = _encode_frame(img,buffer);

	if (! nbytes)
	   Error(1,"Encode %s/%s failed\n",archive_fname,img->dst_fname);

	if (! _write_output(img,archive_fname,buffer,nbytes,out))
	   Error(1,"Error writing output to %s (%s/%s)\n",Output,archive_fname,img->dst_fname);

	DestroyImage(img);
	return nbytes;
	}

static int
_encode_frame(struct Image *img, char *buffer)
	{
	switch(img->depth) {
	   case 8: 
		return _encode_frame_8(img,buffer);
		break;
	   case 16:
		return _encode_frame_16(img,(unsigned short *)buffer);
		break;
	   default:
		Print("_encode_frame: Depth %d not handled\n");
		return 0;
		break;
	   }

	Error(1,"Error, _encode_frame, notreached\n");
	}

static int
_encode_frame_8(struct Image *img, unsigned char *buffer)
	{
	Error(1,"_encode_frame_8: not implemented\n");
	}

static int
_estimate_background_16(struct Image *img)
	{
	unsigned short *ibuf = (unsigned short *)img->data;
	int width = img->width;
	int height = img->height;
	int npixels = width * height;
	double bgtotal;
	int bgcount;
	int tilesize = 5;  	// Use 5x5 tiles
	int zmin = 2;		// Must find at least zmin pixels with value 0
				// to count this tile towards the background
	int x1,y1,x,y,i;

	bgtotal=0; bgcount=0;
	
	for(y1=0; y1<=height-tilesize; y1+=tilesize)
	   for(x1=0; x1<=width-tilesize; x1+=tilesize) {
		int zcount=0;
		unsigned int total=0;

		for(y=y1,zcount=0; y<y1+tilesize; ++y)
		   for(x=x1,i=y*width+x; x<x1+tilesize; ++x,++i) {
			unsigned int n = ibuf[i];
			if (n==0) ++zcount;
			else total += n;
			}

		// Print("Tile (%d,%d) = %d zpixels, total=%d\n",x1,y1,zcount,total);
		if (zcount >= zmin) {
		   bgtotal += total;
		   bgcount += tilesize<<1;
		}
	   }

	return (int)(bgtotal/bgcount);
	}

static int
_encode_frame_16(struct Image *img, unsigned short *buffer)
	{
	unsigned short *ibuf = (unsigned short *)img->data;
	struct fits_info *fi = (struct fits_info *)img->info;
	int npixels = img->width * img->height;
	int x,y,n=0,total,ltotal;
	int width = img->width;
	int height = img->height;
	int maxval=0,minval,minval1;
	int i,zcount,do_shift,tmpval;
	int state,mode,bgval;

	// Estimate the background value and store it in the header
	minval = _estimate_background_16(img);
	minval1 = minval * 2;

	//Print("Background level: %d (%d/256)\n",minval,minval/256);
	for(i=0; i<2880; i += 80) if (! strncmp(fi->header+i,"END ",4)) {
	   char strtmp[32];
	   sprintf(strtmp,"(bg=%d)",minval);
	   strncpy(fi->header+i+4,strtmp,strlen(strtmp));
	   }

	// Detect false zeroes and map to zero.
	// Use buffer[] as the target

	for(y=0; y<height; ++y) for(x=0,i=y*width+x; x<width; ++x,++i) {
	   n = ibuf[i];

	   if (y<1 || x<1 || x>=width-1 || y>=height-1) n=0;
	   else if (n<=minval) n=0;
	   else if (n<=minval1) {
		int zcount=0;
		if (ibuf[i-width-1]<=minval) ++zcount;
		if (ibuf[i-width]<=minval) ++zcount;
		if (ibuf[i-width+1]<=minval) ++zcount;

 		if (ibuf[i-1]<=minval) ++zcount;
 		if (ibuf[i+1]<=minval) ++zcount;

		if (ibuf[i+width-1]<=minval) ++zcount;
		if (ibuf[i+width]<=minval) ++zcount;
		if (ibuf[i+width+1]<=minval) ++zcount;
		if (zcount>=4) n=0;
		}

	   if (n>maxval) maxval=n;
	   buffer[i] = n;
	   }

	for(i=0; i<npixels; ++i) ibuf[i] = buffer[i];

	// Step 3, calculate the data shift required so that maxval
	// fits in 12 bits, leaving the top 4 bits free for use by the 
	// compressor

	do_shift=0;
	tmpval = maxval;
	while(tmpval>=4096) { tmpval>>=1; ++do_shift; }
	//printf("%d bit data shift required\n",do_shift);

	// Step 4 - scan the input buffer and encode pixels using
	// either ZRLE or Z12_16 schemes. The magic 16 bit value 0 switches
	// between the schemes.

	total=ltotal=zcount=state=0;
	mode = ZRLE;

	buffer[total++] = do_shift;

	for(i=0; i<npixels-5; ++i) {
	   if (mode == ZRLE) {
	      if (ibuf[i] || zcount>=2047) {
		if (zcount<15) buffer[total++] = (zcount<<12) | (ibuf[i]>>do_shift);
		else {
		   buffer[total++] = (15<<12) | zcount;
		   buffer[total++] = ibuf[i];
		   }
		zcount=0;
		if (ibuf[i+1] && ibuf[i+2] && ibuf[i+3] && ibuf[i+4]) {
		   buffer[total++] = 0; mode=Z12_16;
		   }
		}
	      else ++zcount;
	      }
	   else { // mode is Z12_16
		if (state==0) {
		   tmpval=ibuf[i]>>do_shift;
		   if (tmpval&0xf000) Error(1,"tmpval %d too large (%d >> %d)\n",tmpval,ibuf[i],do_shift);
		   state++;
		   }
		else if (state==1) {
		   unsigned short v = ibuf[i]; v>>=do_shift;
		   buffer[total++] = v | ((tmpval&0xf)<<12); state++; 
		   }
		else if (state==2) {
		   unsigned short v = ibuf[i]; v>>=do_shift;
		   buffer[total++] = v | ((tmpval&0xf0)<<8); state++; 
		   }
		else if (state==3) {
		   unsigned short v = ibuf[i]; v>>=do_shift;
		   buffer[total++] = v | ((tmpval&0xf00)<<4); 
		   state=0; 
		   if (i<npixels-2 && ((ibuf[i+1]&0xf)==0 && (ibuf[i+2]>>do_shift)==0)) {mode=ZRLE; buffer[total++]=0; }
		   }
		}
	   }

	//printf("%d/%d data bytes = %lf\n",total,npixels,(double)total/npixels);
	//if (mode==ZRLE) Print("%d trailing black pixels skipped\n",zcount);

	return total * 2;  // return number of bytes
	}

static int
_write_output(struct Image *img, char *archive_fname,unsigned char *buffer, int nbytes, FILE *out)
	{
	struct fits_info *fi = img->info;

	if (! fwrite((void *)archive_fname,strlen(archive_fname)+1,1,out))
	   Error(1,"Write failed on %s\n",Output);

	if (! fwrite((void *)img->src_fname,strlen(img->src_fname)+1,1,out))
	   Error(1,"Write failed on %s\n",Output);

	if (! fwrite((void *)fi->header,2880,1,out))
		Error(1,"Write failed on %s\n",Output);

	if (! fwrite((void *)&nbytes,sizeof(nbytes),1,out))
		Error(1,"Write failed on %s\n",Output);

	if (! fwrite((void *)buffer,nbytes,1,out))
		Error(1,"Write failed on %s\n",Output);

	fflush(out);
	return nbytes;
	}

//------------------------------------------------------------------------------------------

static int
DeCompress(char *fname)
	{
	struct Image *img = CreateImage();
	struct fits_info *fi;
	unsigned char *buffer = NULL;
	unsigned char *obuf = NULL;
	char archive_fname[1024],tmpstr[128],*ptr;
	char current_archive_fname[1024];
	int ntmp,nbytes,zbytes;
	FILE *in,*out=NULL;
	extern int HistoStretch;

	HistoStretch = 0;

	Print("Decompressing from '%s'\n",fname);

	if ((in = fopen(fname,"rb")) == NULL)
	   Error(1,"Error opening file '%s'\n",fname);

	current_archive_fname[0]=0;

	while(1) {

            // read null-terminated real filename
            ptr = archive_fname; *ptr=0;
            while(1) {
              int ch = ReadByte(in);
              if (ch==EOF) { 
		   Print("Finished reading input\n");
	  	   return 1;
		   }
              if (ch<0) { return 0;}

	      if (strlen(archive_fname) >= 1023) 
              	   Error(1,"Archive entry '%s/%s' error: filename too long\n",archive_fname);

              *(ptr++) = ch; *ptr=0; 
              if (ch==0) break;
              }


            // read null-terminated filename inside
	    // the archive
            ptr = tmpstr; *ptr=0;
            while(strlen(tmpstr)<32) {
              int ch = ReadByte(in);
              if (ch==EOF) { return 0;}
              if (ch<0) { return 0;}
              *(ptr++) = ch; *ptr=0; 
              if (ch==0) break;
              }

	   img->type = IMG_FIT;
	   img->src_fname = strdup(tmpstr);
	   img->dst_fname = strdup(tmpstr);

	   Print("Decoding Image: [%s/%s] ",archive_fname,img->src_fname);

	   // Read the 2880 (36 x 80) header records
	   if (! load_image_header(in,img)) Error(1,"End of input data\n");

	   if (! obuf) { obuf = Malloc(img->width*img->height*2); }
	   if (! buffer) { buffer = Malloc(img->width*img->height*2); }

	   if (img->dst_width==0)  img->dst_width = img->width;
           if (img->dst_height==0) img->dst_height = img->height;

	   Print("%dx%d (%d bpp)\n",img->width,img->height,img->depth);

	   if (Fread((unsigned char *)&zbytes,sizeof(zbytes),1,in) != 1)
	   	Error(1,"Error: Short read on zbytes\n");

	   // Print("%d bytes of compressed data\n",zbytes);
	   // Print("File Offset %ld\n",lseek(fileno(in),0,SEEK_CUR));
	   if (Fread(buffer, zbytes, 1, in) != 1)
	   	Error(1,"Short read on data\n");

	   img->data = obuf;
	   nbytes = _ftz_decode_frame(img,buffer,zbytes);

	   if (! nbytes) Error(1,"Error decoding data\n");

	   if (strcmp(archive_fname,current_archive_fname)) {
		char *ptr = archive_fname;

	   	Print("Opening %s for Writing\n",archive_fname);
		strcpy(current_archive_fname,archive_fname);

		if (out) fclose(out);

		// Check and make sure parent directories exist
		while(*ptr) {
		   if (*(ptr+1) == '/' || *(ptr+1)=='\\') {
			*(ptr+1)=0;
			if (! isDirectory(archive_fname)) {
			   Print("Mkdir(%s)\n",archive_fname);
			   Mkdir(archive_fname);
			   }
			*(ptr+1) = '/';
			}
		   ++ptr;
		   }
		
		out = fopen(archive_fname,"wb");
		if (! out) Error(1,"Error: cannot open '%s' for writing\n",archive_fname);
		}

	   fwrite(img->src_fname,strlen(img->src_fname)+1,1,out);
	   fi = (struct fits_info *)img->info;
	   fwrite(fi->header,2880,1,out);
	   ntmp = write_image_data(out,img);

	   //printf("wrote %d/%d bytes for %s\n",ntmp,nbytes,img->dst_fname);
	   }
	
	// notreached
	return 0;
	}

int
_ftz_decode_frame(struct Image *img, char *buffer,int zbytes)
	{
	switch(img->depth) {
	   case 8: 
		return _ftz_decode_frame_8(img,buffer,zbytes);
		break;
	   case 16:
		return _ftz_decode_frame_16(img,(unsigned short *)buffer,zbytes);
		break;
	   default:
		Print("_ftz_decode_frame: Depth %d not handled\n");
		return 0;
		break;
	   }

	Error(1,"Error, _ftz_decode_frame, notreached\n");
	return 0;
	}

static int
_ftz_decode_frame_8(struct Image *img, unsigned char *buffer, int zbytes)
	{
	Error(1,"_decode_frame_8: not implemented\n");
	return 0;
	}

static int
_ftz_decode_frame_16(struct Image *img, unsigned short *ibuf, int zwords)
	{
	struct fits_info *fi = (struct fits_info *)img->info;
	unsigned short *buffer = (unsigned short *)img->data;
	unsigned short *tmpbuf;
	int width = img->width;
	int npixels = width * img->height;
	int x,y,n,total;
	int i,do_shift,tmpval;
	int mode,black=ZBlack<<8,black2;
	unsigned int zcount;

	// See if the default bg is encoded into the header
	if (! black) for(i=0; i<80*36; i+=80) if (!strncmp(fi->header+i,"END (",5)) {
	   sscanf(fi->header+i+4,"(bg=%d)",&black);
	   }

	// Step 1 - recover the shift and nbytes data from the start
	// of the stream

	i=0;
	zwords >>= 1;
	do_shift=ibuf[i++];

	//printf("%d zwords, %d bit data shift required\n",zwords,do_shift);

	// Step 2 - scan the input buffer and decode pixels using
	// either ZRLE or Z12_16 schemes. The magic 16 bit value 0 switches
	// between the schemes.

	total=zcount=0;
	mode = ZRLE;;

	while(i<zwords && total<npixels-5) {
	   if (mode == ZRLE) {
	      if (! ibuf[i]) { mode = Z12_16; ++i; continue; }

	      zcount = ibuf[i]>>12;
	      if (zcount==15) {zcount = ibuf[i++]&0xfff; n = ibuf[i]; }
	      else n = (ibuf[i]&0xfff) << do_shift;;

	      while(zcount--) buffer[total++]=black;
	      buffer[total++] = n;
	      }
	   else { // mode is Z12_16
		int v;
		tmpval=ibuf[i]>>12;
		v = (ibuf[i++]&0xfff) << do_shift; if (0 && v<black) v=black; buffer[++total] = v;
		tmpval |= (ibuf[i]&0xf000)>>8;
		v = (ibuf[i++]&0xfff) << do_shift; if (0 && v<black) v=black; buffer[++total] = v;
		tmpval |= (ibuf[i]&0xf000)>>4;
		v = (ibuf[i++]&0xfff) << do_shift; if (0 && v<black) v=black; buffer[++total] = v;
		tmpval <<= do_shift;
		if (0 && tmpval<black) tmpval=black;
		buffer[total-3] = tmpval;
		total++;
		if (! ibuf[i]) { mode = ZRLE; i++; }
		continue;
		}

	   ++i;

	   if (total > npixels) Error(1,"oops, too many decode pixels\n");
	   }

	//Print("reconstituted %d/%d pixels (%lf)\n",total,npixels,(double)total/npixels);

	if (total < npixels) {
	   //Print("Padding with %d black pixels\n",npixels-total);
	   while(total<npixels) buffer[total++]=0;
	   }

	if (total != npixels)
	   Error(1,"Decode error, pixel count mismatch! wanted %d, found %d\n",npixels,total);

	// Soften the black transitions
	tmpbuf = ZeroMalloc(npixels * 2);
	black2 = black * 2;
	for(y=1; y<img->height-1; ++y) for(x=1,i=y*width+x; x<width-1; ++x,++i) {
	   if (buffer[i]) {
		int v = buffer[i];
		int j = i-width-1;
		if (buffer[j] <= black) tmpbuf[j]=(tmpbuf[j]+v)/2; ++j;
		if (buffer[j] <= black) tmpbuf[j]=(tmpbuf[j]+v)/2; ++j;
		if (buffer[j] <= black) tmpbuf[j]=(tmpbuf[j]+v)/2;

		j = i-1;
		if (buffer[j] <= black) tmpbuf[j]=(tmpbuf[j]+v)/2; j+=2;
		if (buffer[j] <= black) tmpbuf[j]=(tmpbuf[j]+v)/2;

		j = i+width-1;
		if (buffer[j] <= black) tmpbuf[j]=(tmpbuf[j]+v)/2; ++j;
		if (buffer[j] <= black) tmpbuf[j]=(tmpbuf[j]+v)/2; ++j;
		if (buffer[j] <= black) tmpbuf[j]=(tmpbuf[j]+v)/2;
		}
	   }

	for(y=1; y<img->height-1; ++y) for(x=1,i=y*width+x; x<width-1; ++x,++i) {
	   if (tmpbuf[i]) buffer[i]=tmpbuf[i];
	   }

	free(tmpbuf);

	return total * 2;  // return number of bytes
	}