/*
 * luttool.c
 *
 * Utility to apply LUTs to images
 *
 * Author: Dario Rodriguez antartica@whereismybit.com
 * This program is licensed under the terms of the MIT license,
 * with an exception: the license text may be substituted for a
 * link to it. MIT license text: https://spdx.org/licenses/MIT.html
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <png.h>

#define DEFAULT_IDENTITYSIZE 12 /* Same as RawTherapee's HaldCLUTs, 1728x1728 AKA (12*12)*12x12*(12*12), equivalent to a cube with 144px wide sides*/
#define DEFAULT_QUALITY 85
/* #define JPEG_SUPPORTED */

int luttool_applylut(char *lutfilenamelist,char *infilename, char *outfilename,int quality,int identitysize);

int util_pnginfo(char *filename, int *width, int *height);
int util_pngload(char *filename, char **bufimg, int *sizebufimg, int *width, int *height);
int util_pngsave(char *filename, char *bufimg, int width, int height);
int util_identitygen(char **bufimg, int *sizebufimg, int identitysize, int *width, int *height);
int luttool_rgba2lut4096(char *rgba, int width, int height, char **lut4096, int *sizelut4096);
int luttool_rgba2lut4096_interpolated(char *rgba, int width, int height, char **lut4096, int *sizelut4096);
int luttool_applylut4096(char *lut4096,char *imgbuf,int imgwidth,int imgheight);

int
main(int argc, char *argv[])
{
        int firstarg,outfilenamearg;
        int identitysize;
        int quality;
        char *outfilename;
        char *lutlist;
        int i;
        int flagseveralinfile;
        char bufoutfilename[1024];
        char *currentoutfilename;
        if(argc<2 || strcmp(argv[argc-1],"--help")==0) {
#ifdef JPEG_SUPPORTED
                printf("Syntax: %s [-q jpegquality] [-n identitylevel] {identity | lutfile[,lutfile[,...]]} { identity | infile1 } [infile2 [...]] -o outfile\n",argv[0]);
                printf("Example: %s haldclut.png sprite.png -o result.png\n",argv[0]);
                printf("Example: %s -n 16 identity identity -o identity.png\n",argv[0]);
                printf("Example: %s -q 85 haldclut.png infile.png -o outfile.jpg\n",argv[0]);
#else
                printf("Syntax: %s [-n identitylevel] {identity | lutfile[,lutfile[,...]]} { identity | infile1 } [infile2 [...]] -o outfile\n",argv[0]);
                printf("Example: %s haldclut.png sprite.png -o result.png\n",argv[0]);
                printf("Example: %s -n 16 identity identity -o identity.png\n",argv[0]);
#endif
                return(1);
        }
        /* parse arguments */
        identitysize=DEFAULT_IDENTITYSIZE;
        quality=DEFAULT_QUALITY;
        for(firstarg=1;firstarg<(argc-1) && (strcmp(argv[firstarg],"-n")==0 || (strcmp(argv[firstarg],"-q")==0));) {
                if(strcmp(argv[firstarg],"-n")==0) {
                        identitysize=atoi(argv[firstarg+1]);
                        if(identitysize<2 || identitysize>16) {
                                fprintf(stderr,"%s: ERROR: the parameter to -n must be between 2 and 16 (default: %i)\n",argv[0],DEFAULT_IDENTITYSIZE);
                                return(1);
                        }
                } else if(strcmp(argv[firstarg],"-q")==0) {
                        quality=atoi(argv[firstarg+1]);
                        if(quality<1 || quality>100) {
                                fprintf(stderr,"%s: ERROR: the parameter to -q must be between 1 and 100 (default: %i)\n",argv[0],DEFAULT_QUALITY);
                                return(1);
                        }
                }
                firstarg+=2;
        }
        for(outfilenamearg=-1,outfilename=NULL,i=1;i<(argc-1);i++) {
                if(strcmp(argv[i],"-o")==0) {
                        outfilename=argv[i+1];
                        outfilenamearg=i;
                        break;
                }
        }
        if(outfilename==NULL) {
                fprintf(stderr,"%s: ERROR: no output file specified\n",argv[0]);
                return(1);
        }
        firstarg+=((firstarg==outfilenamearg)?2:0);
        lutlist=argv[firstarg];
        firstarg++;
        flagseveralinfile=((argc-firstarg-((outfilenamearg>=firstarg)?2:0))>1)?1:0;
        /* process files sequentially, skipping the "-o outputfilename" */
        for(i=firstarg;i<argc;i++) {
                if(i==outfilenamearg) {
                        i++;
                        continue;
                }
                currentoutfilename=outfilename;
                if(flagseveralinfile) {
                        char *ext;
                        int l;
                        strncpy(bufoutfilename,argv[i],sizeof(bufoutfilename));
                        bufoutfilename[sizeof(bufoutfilename)-1]='\0';
                        if((ext=strrchr(bufoutfilename,'.'))!=NULL)
                                *ext='\0';
                        l=strlen(bufoutfilename);
                        snprintf(bufoutfilename+l,sizeof(bufoutfilename)-l,"_%s",outfilename);
                        bufoutfilename[sizeof(bufoutfilename)-1]='\0';
                        currentoutfilename=bufoutfilename;
                }
                if(luttool_applylut(lutlist,argv[i],currentoutfilename,quality,identitysize)!=0) {
                        fprintf(stderr,"%s: ERROR: couldn't apply LUTs \"%s\" to \"%s\" or couldn't write \"%s\"cified\n",argv[0],lutlist,argv[i],currentoutfilename);
                        return(1);
                }
        }
        return(0);
}

int
luttool_applylut(char *lutlist,char *infilename, char *outfilename,int quality,int identitysize)
{
        char *curlut,*next;
        char filename[1024];
        int l;
        int in_error;
        char *lutbuf,*imgbuf,*rgblutbuf;
        int sizelutbuf,lutwidth,lutheight;
        int sizeimgbuf,imgwidth,imgheight;
        int sizergblutbuf;
        lutbuf=imgbuf=rgblutbuf=NULL;
        sizelutbuf=sizeimgbuf=sizergblutbuf=0;
        lutwidth=lutheight=imgwidth=imgheight=0;
        if(lutlist==NULL || infilename==NULL || outfilename==NULL)
                return(-1); /* sanity check error */
        in_error=0;
        /* load the image to process */
        if(strcmp(infilename,"identity")==0) {
                 if(util_identitygen(&imgbuf,&sizeimgbuf,identitysize,&imgwidth,&imgheight)==-1)
                         in_error=1;
        } else {
                if(util_pngload(infilename,&imgbuf,&sizeimgbuf,&imgwidth,&imgheight)!=0)
                        in_error=1;
        }
        /* apply all luts */
        for(curlut=lutlist,next=strchr(curlut,','),next=((next==NULL)?(curlut+strlen(curlut)):next)
          ;in_error==0 && *curlut!='\0'
          ;curlut=next+((*next==',')?1:0),next=strchr(curlut,','),next=((next==NULL)?(curlut+strlen(curlut)):next)) {
                l=next-curlut;
                l=(l>(sizeof(filename)-1))?(sizeof(filename)-1):l;
                memcpy(filename,curlut,l);
                filename[l]='\0';
                if(strcmp(filename,"identity")==0) {
                        if(util_identitygen(&lutbuf,&sizelutbuf,identitysize,&lutwidth,&lutheight)==-1) {
                                in_error=1;
                                break; /* couldn't generate identity lut */
                        }
                } else {
                        if(util_pngload(filename,&lutbuf,&sizelutbuf,&lutwidth,&lutheight)==-1) {
                                in_error=1;
                                break; /* couldn't load lut */
                        }
                }
                if(luttool_rgba2lut4096(lutbuf, lutwidth, lutheight, &rgblutbuf, &sizergblutbuf)!=0) {
                        in_error=1;
                        break; /* couldn't process lut */
                }
                if(luttool_applylut4096(rgblutbuf,imgbuf,imgwidth,imgheight)!=0) {
                        in_error=1;
                        break; /* couldn't apply lut */
                }
        }
        /* save the result */
        if(in_error==0) {
                if(util_pngsave(outfilename,imgbuf,imgwidth,imgheight)!=0)
                        in_error=1; /* couldn't write */
        }
        /* cleanup and exit */
        if(lutbuf!=NULL)
                free(lutbuf),lutbuf=NULL;
        if(imgbuf!=NULL)
                free(imgbuf),imgbuf=NULL;
        if(rgblutbuf!=NULL)
                free(rgblutbuf),rgblutbuf=NULL;
        return((in_error!=0)?-1:0);
}



int
util_pnginfo(char *filename, int *width, int *height)
{
        return(util_pngload(filename,NULL,NULL,width,height));
}

int
util_pngload(char *filename, char **bufimg, int *sizebufimg, int *width, int *height)
{
        FILE *f;
        png_structp png;
        png_infop info;
        png_byte color_type,bit_depth;
        png_bytep *rowpointers;
        int reqsize;
        int i,pitch;
        if(filename==NULL || (bufimg==NULL && sizebufimg!=NULL) || (bufimg!=NULL && sizebufimg==NULL) || width==NULL || height==NULL)
                return(-1); /* sanity check failed */
        if((f=fopen(filename,"rb"))==NULL)
                return(-1); /* couldn't open file */
        if((png=png_create_read_struct(PNG_LIBPNG_VER_STRING,NULL,NULL,NULL))==NULL
          || (info=png_create_info_struct(png))==NULL
          || setjmp(png_jmpbuf(png))) {
                if(png!=NULL)
                        png_destroy_read_struct(&png,&info,NULL),png=NULL,info=NULL;
                fclose(f),f=NULL;
                return(-1); /* couldn't init png library */
        }
        png_init_io(png,f);
        png_read_info(png,info);
        *width=png_get_image_width(png,info);
        *height=png_get_image_height(png,info);
        /* special case: if bufimg==NULL, we only want the width/height */
        if(bufimg==NULL) {
                png_destroy_read_struct(&png,&info,NULL),png=NULL,info=NULL;
                fclose(f),f=NULL;
                return(0); /* all done */
        }
        /* make sure we have enough space for the new image in 8-bit RGBA */
        reqsize=(*width)*(*height)*4;
        if(*sizebufimg<reqsize) {
                char *newbufimg;
                if((newbufimg=realloc(*bufimg,reqsize))==NULL) {
                      png_destroy_read_struct(&png,&info,NULL),png=NULL,info=NULL;
                      fclose(f),f=NULL;
                      return(-1); /* insuf. mem. for resulting image */
                }
                *bufimg=newbufimg;
                *sizebufimg=reqsize;
        }
        /* make sure we get 8-bit RGBA data, regardless of the file's bit_depth and color_type */
        if((bit_depth=png_get_bit_depth(png,info))==16)
                png_set_strip_16(png);
        color_type=png_get_color_type(png,info);
        if(color_type==PNG_COLOR_TYPE_PALETTE)
                png_set_palette_to_rgb(png);
        if(color_type==PNG_COLOR_TYPE_GRAY && bit_depth<8)
                png_set_expand_gray_1_2_4_to_8(png);
        if(png_get_valid(png,info,PNG_INFO_tRNS))
                png_set_tRNS_to_alpha(png);
        if(color_type==PNG_COLOR_TYPE_RGB
          || color_type==PNG_COLOR_TYPE_GRAY
          || color_type==PNG_COLOR_TYPE_PALETTE) {
                png_set_filler(png,0xff,PNG_FILLER_AFTER);
        }
        if(color_type==PNG_COLOR_TYPE_GRAY
          || color_type==PNG_COLOR_TYPE_GRAY_ALPHA) {
                png_set_gray_to_rgb(png);
        }
        /* read the image */
        png_read_update_info(png,info);         
        if((rowpointers=(png_bytep *)malloc(sizeof(png_bytep)*(*height)))==NULL) {
                png_destroy_read_struct(&png,&info,NULL),png=NULL,info=NULL;
                fclose(f),f=NULL;
                return(-1); /* insuf. mem. for temporary data */
        }
        for(i=0,pitch=(*width)*4;i<(*height);i++)
                rowpointers[i]=(png_byte *)((*bufimg)+i*pitch);
        png_read_image(png,rowpointers);
        /* cleanup */
        free(rowpointers),rowpointers=NULL;
        fclose(f),f=NULL;
        png_destroy_read_struct(&png,&info,NULL),png=NULL,info=NULL;
        return(0);
}

int
util_pngsave(char *filename, char *bufimg, int width, int height)
{
        FILE *f;
        png_structp png;
        png_infop info;
        int i,pitch;
        char *ptr;
        if((f=fopen(filename,"wb"))==NULL)
                return(-1); /* couldn't open file for writing */
        if((png=png_create_write_struct(PNG_LIBPNG_VER_STRING,NULL,NULL,NULL))==NULL
          || (info=png_create_info_struct(png))==NULL
          || setjmp(png_jmpbuf(png))) {
                if(png!=NULL)
                        png_destroy_write_struct(&png,&info),png=NULL,info=NULL;
                fclose(f),f=NULL;
                return(-1); /* couldn't init png library */
        }
        png_init_io(png,f);
        png_set_IHDR(png,info,width,height,8,PNG_COLOR_TYPE_RGBA,
          PNG_INTERLACE_NONE,PNG_COMPRESSION_TYPE_BASE,PNG_FILTER_TYPE_BASE);
        png_write_info(png,info);
        for(i=0,ptr=bufimg,pitch=width*4;i<height;i++,ptr+=pitch)
                png_write_row(png,(png_const_bytep)ptr);
        png_write_end(png,NULL);
        fclose(f),f=NULL;
        png_destroy_write_struct(&png,&info),png=NULL,info=NULL;
        return(0);
}

int
util_identitygen(char **bufimg, int *sizebufimg, int identitysize, int *width, int *height)
{
        int x,y,r,g,b;
        int isizesq,isizesqminusone;
        int osizesqminusone;
        int reqsize;
        int l;
        unsigned char *ptr;
        if(bufimg==NULL || sizebufimg==NULL || identitysize<=1 || identitysize>16 || width==NULL || height==NULL)
                return(-1);
        l=*width=*height=identitysize*identitysize*identitysize;
        reqsize=(*width)*(*height)*4;
        if(*sizebufimg<reqsize) {
                char *newbufimg;
                if((newbufimg=realloc(*bufimg,reqsize))==NULL)
                      return(-1); /* insuf. mem. for resulting image */
                *bufimg=newbufimg;
                *sizebufimg=reqsize;
        }
        isizesq=identitysize*identitysize;
        isizesqminusone=isizesq-1;
        osizesqminusone=256-1;
        for(y=0,ptr=(unsigned char *)(*bufimg);y<l;y++) {
                for(x=0;x<l;x++,ptr+=4) {
                        r=x%isizesq;
                        g=(y%identitysize)*identitysize+(x/isizesq);
                        b=y/identitysize;
                        ptr[0]=r*osizesqminusone/isizesqminusone;
                        ptr[1]=g*osizesqminusone/isizesqminusone;
                        ptr[2]=b*osizesqminusone/isizesqminusone;
                        ptr[3]=0xff;
                }
        }
        return(0);
}

int
luttool_rgba2lut4096(char *rgba, int width, int height, char **lut4096, int *sizelut4096)
{
        int n;
        int reqsize;
        int x,y;
        int r,g,b;
        int isize,isizesq,isizesqminusone;
        int osize,osizesq,osizesqminusone;
        unsigned char *orig,*dest;
        if(rgba==NULL || width!=height || lut4096==NULL || sizelut4096==NULL)
                return(-1); /* sanity check error */
        /* check if size is a supported size */
        for(n=2;n<=16 && width!=(n*n*n);n++)
                ;
        if(n>16)
                return(-1); /* invalid size */
        /* expand lut4096 as necessary */
        reqsize=4096*4096*3;
        if(*sizelut4096<reqsize) {
                char *newlut4096;
                if((newlut4096=realloc(*lut4096,reqsize))==NULL)
                      return(-1); /* insuf. mem. for resulting lut */
                *lut4096=newlut4096;
                *sizelut4096=reqsize;
        }
        /* do the conversion */
        isize=n;
        isizesq=n*n;
        isizesqminusone=isizesq-1;
        osize=16;
        osizesq=osize*osize;
        osizesqminusone=osizesq-1;
        for(r=0;r<osizesq;r++) {
                for(g=0;g<osizesq;g++) {
                        for(b=0;b<osizesq;b++) {
                                x=((isizesqminusone*r)/osizesqminusone)+(((isizesqminusone*g)/osizesqminusone)%isize)*isizesq;
                                y=((isizesqminusone*b)/osizesqminusone)*isize+(((isizesqminusone*g)/osizesqminusone)/isize);
                                orig=(unsigned char *) (rgba+((y*width+x)<<2));
                                dest=(unsigned char *) ((*lut4096)+((r<<16)+(g<<8)+b)*3);
                                dest[0]=orig[0];
                                dest[1]=orig[1];
                                dest[2]=orig[2];
                        }
                }
        }
        return(0);
}

int
luttool_applylut4096(char *lut4096,char *imgbuf,int imgwidth,int imgheight)
{
        int x,y;
        unsigned char *orig;
        unsigned char *lut;
        if(lut4096==NULL || imgbuf==NULL || imgwidth<=0 || imgheight<=0)
                return(-1);
        for(y=0,orig=(unsigned char *)imgbuf;y<imgheight ;y++) {
                for(x=0;x<imgwidth;x++,orig+=4) {
                        lut=((unsigned char *)lut4096)+((((int)(orig[0]))<<16)+(((int)(orig[1]))<<8)+((int)(orig[2])))*3;
                        orig[0]=lut[0];
                        orig[1]=lut[1];
                        orig[2]=lut[2];
                }
        }
        return(0);
}