/* loadavg.png.c   by Michael Thorpe   2022-01-14 */

/* Largest allowed scale (no longer strictly necessary) */
#define MAX_SCALE 100000
/* Smallest and largest allowed dimentions */
#define MIN_DIM 10
#define MAX_DIM 2000
/* Number of seconds between samples of the loadavg data */
#define DATA_INTERVAL 5
/* Also, see below for data source definitions */

#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <glob.h>
/* Needed for open */
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
/* Needed for isatty and chdir */
#include <unistd.h>
/* unistd and this needed for mmap */
#include <sys/mman.h>
/* Needed for opendir(), etc. */
#include <dirent.h>
#include <sys/types.h>
/* Both needed for png.h :( */
#include <libpng/png.h>
#include <setjmp.h>
/* Needed for INT_MAX */
#include <limits.h>

struct datasrc {
	char *name;
	char *dir;
};

struct datasrc datasrcs[]={
	{"","/var/log/loadavg"},
	{"mysqlthreads","/var/log/mysqlthreads"},
	{0,0}
};

#define WHITE  0
#define BLACK  1
#define RED    3
#define GREEN  2
#define BLUE   4
#define YELLOW 5
#define PURPLE 6
#define TEAL   7
#define COLORS 8
#define LOG2COLORS 4 /* libpng isn't happy with 3 for some reason */

#define BACK       WHITE
#define EDGE       BLACK
#define RULE       GREEN
#define LOAD       RED
#define RESTART    PURPLE
#define TIMECHANGE TEAL
#define NODATA     YELLOW

#define AUX_RESTART    1
#define AUX_TIMECHANGE 2

#define DATESIZE 32  /* This contains an extra two bytes "for luck" ;) */

struct rawdata {
	char **image;
	int *mindata,*maxdata,*auxdata,*nodata;
	int width,height;
	int togo,samplestogo;
	int max,scale,horizlines,cap;
	int oldsample;
	char *rawimage;
};

/* Returns false if processing should stop */
static int addsample(struct rawdata *srd,char *s) {
	int auxdata,halfway,i,nodata,sample=0;
	unsigned long l;

	if(isdigit(*s)) {
		l=strtoul(s,&s,10);
		if(l+1>INT_MAX/100) {
			sample=INT_MAX;
			if(*s=='.' && isdigit(*++s) && isdigit(*++s))
				s++;
		} else { /* l*100+99 <= INT_MAX */
			sample=l*100;
			if(*s=='.' && isdigit(*++s)) {
				sample+=10*(*s++-'0');
				if(isdigit(*s))
					sample+=*s++-'0';
			}
		}
		if(s[0] != '\n')
			return(0);
		if(srd->cap && sample>srd->cap)
			sample=srd->cap;
		if(srd->max<sample)
			srd->max=sample;
got_negative_sample:
		if(srd->samplestogo) {
			if(srd->togo<srd->width) {
				if(sample<srd->mindata[srd->togo])
					srd->mindata[srd->togo]=sample;
				else if(srd->maxdata[srd->togo]<sample)
					srd->maxdata[srd->togo]=sample;
			}
			srd->samplestogo--;
		} else {
			if(srd->togo<srd->width && -1 != srd->oldsample) {
				halfway=srd->oldsample-(srd->oldsample-sample)/2;
				srd->mindata[srd->togo]=srd->mindata[srd->togo]<halfway?srd->mindata[srd->togo]:halfway;
				srd->maxdata[srd->togo]=srd->maxdata[srd->togo]<halfway?halfway:srd->maxdata[srd->togo];
			}
			srd->togo--;
			srd->mindata[srd->togo]=sample;
			srd->maxdata[srd->togo]=sample;
			if(-1 != srd->oldsample) {
				halfway=sample+(srd->oldsample-sample)/2;
				if(sample>halfway)
					srd->mindata[srd->togo]=halfway;
				if(sample<halfway)
					srd->maxdata[srd->togo]=halfway;
			}
			srd->samplestogo=srd->scale-1;
		}
		srd->oldsample=sample;
		return(srd->togo || srd->samplestogo);
	} else if(!strncmp("nodata\n",s,7) || !strncmp("none\n",s,5)) {
		auxdata=0;
		nodata=1;
		goto got_fake_sample;
	} else if(!strncmp("restart\n",s,8)) {
		auxdata=AUX_RESTART;
		nodata=0;
		goto got_fake_sample;
	} else if(!strncmp("timechange\n",s,11)) {
		auxdata=AUX_TIMECHANGE;
		nodata=0;
		goto got_fake_sample;
	} else if('-'==s[0]) {
		while(isdigit(*++s))
			;
		if(*s=='.' && isdigit(*++s) && isdigit(*++s))
			s++;
		if(*s != '\n') {
			fprintf(stderr,"loadavg.png: Invalid negative sample\n");
			return(0);
		}
		sample=0;
		goto got_negative_sample;
	}
	for(i=0;s[i] && s[i] != '\n';s++)
		;
	fprintf(stderr,"loadavg.png: Invalid sample: \"%.*s\"\n",i,s);
	return(0);
got_fake_sample:
	if(srd->samplestogo) {
		srd->samplestogo--;
	} else {
		srd->togo--;
		srd->samplestogo=srd->scale-1;
	}
	srd->auxdata[srd->togo]|=auxdata;
	srd->nodata[srd->togo]+=nodata;
	srd->oldsample=-1;
	return(srd->togo || srd->samplestogo);
}

/* Returns 0 normally, 1 on stop-processing, -1 on error */
static int loaddata(struct rawdata *srd,char *file) {
	int f,retval=1;
	char *buf;
	off_t l,size;

	if(0==srd->togo && 0==srd->samplestogo)
		return(0);
	f=open(file,O_RDONLY);
	if(f==-1)
		return(-1);
	size=lseek(f,0,SEEK_END);
	if(size==-1 || size<2) {
		close(f);
		return(size); /* Strangely enough, this works in all cases */
	}
	buf=mmap(0,size,PROT_READ,MAP_SHARED,f,0);
	l=size-1;
	while(1) {
		if(l==0)
			goto nomore;
		if(buf[l]=='\n')
			break;
		l--;
	}
	while(1) {
		if(l==0 || buf[l-1]=='\n') {
			if(!addsample(srd,buf+l))
				goto nomore;
			if(l==0)
				break;
		}
		l--;
	}
	retval=0;
nomore:
	munmap(buf,size);
	if(close(f))
		return(-1);
	return(retval);
}

static void fillimage(struct rawdata *srd) {
	int i,j,m,M,n,d;

	memset(srd->rawimage,BACK,srd->height*srd->width);
	if(srd->horizlines)
		srd->max=((srd->max+srd->horizlines-1)/srd->horizlines)*srd->horizlines;
	n=srd->height-1;
	d=srd->max;
	for(i=1;i<srd->width-1;i++) {
		if(srd->auxdata[i]) {
			if(srd->auxdata[i]&AUX_RESTART)
				m=RESTART;
			else if(srd->auxdata[i]&AUX_TIMECHANGE)
				m=TIMECHANGE;
			else
				m=RULE;
			for(j=0;j<srd->height;j++)
				srd->image[j][i]=m;
		}
		if(srd->nodata[i] && (srd->nodata[i]<srd->scale || !srd->auxdata[i])) {
			for(j=srd->height*(srd->scale-srd->nodata[i])/srd->scale;j<srd->height;j++)
				srd->image[j][i]=NODATA;
		}
	}
	if(srd->horizlines) {
		for(i=1;;i++) {
			m=srd->height-1-i*srd->horizlines*n/d;
			if(m<=0)
				break;
			for(j=1;j<srd->width-1;j++)
				srd->image[m][j]=RULE;
		}
	}
	for(i=1;i<srd->width-1;i++) {
		if(srd->maxdata[i]<srd->mindata[i])
			continue;
		m=srd->height-1-srd->mindata[i]*n/d;
		M=srd->height-1-srd->maxdata[i]*n/d;
		do
			srd->image[M++][i]=LOAD;
		while(M<=m);
	}
	for(i=0;i<srd->width;i++) {
		srd->image[0][i]=EDGE;
		srd->image[srd->height-1][i]=EDGE;
	}
	for(i=1;i<srd->height-1;i++) {
		srd->image[i][0]=EDGE;
		srd->image[i][srd->width-1]=EDGE;
	}
}

static int writepng(struct rawdata *srd,FILE *out,char *date) {
	png_structp png_ptr;
	png_infop info_ptr;
	png_color *palette;
	png_text *text_ptr;

	png_ptr=png_create_write_struct(PNG_LIBPNG_VER_STRING,0,0,0);
	if(!png_ptr)
		return(-1);
	info_ptr=png_create_info_struct(png_ptr);
	if(!info_ptr) {
		png_destroy_write_struct(&png_ptr,0);
		return(-1);
	}
	if(setjmp(png_jmpbuf(png_ptr))) {
		png_destroy_write_struct(&png_ptr,&info_ptr);
		return(-1);
	}
	png_init_io(png_ptr,out);
	png_set_IHDR(png_ptr,info_ptr,srd->width,srd->height,LOG2COLORS,PNG_COLOR_TYPE_PALETTE,PNG_INTERLACE_NONE,PNG_COMPRESSION_TYPE_BASE,PNG_FILTER_TYPE_BASE);
	text_ptr=(png_text *)calloc(1,sizeof(png_text));
	if(!text_ptr) {
		png_destroy_write_struct(&png_ptr,&info_ptr);
		return(-1);
	}
	text_ptr[0].compression=PNG_TEXT_COMPRESSION_NONE;
	text_ptr[0].key="Creation Time";
	text_ptr[0].text=date;
	png_set_text(png_ptr,info_ptr,text_ptr,1);
	palette=(png_color *)malloc(COLORS*sizeof(png_color));
	if(!palette) {
		free(text_ptr);
		png_destroy_write_struct(&png_ptr,&info_ptr);
		return(-1);
	}
	palette[WHITE].red=-1;
	palette[WHITE].green=-1;
	palette[WHITE].blue=-1;
	palette[BLACK].red=0;
	palette[BLACK].green=0;
	palette[BLACK].blue=0;
	palette[RED].red=-1;
	palette[RED].green=0;
	palette[RED].blue=0;
	palette[GREEN].red=0;
	palette[GREEN].green=-1;
	palette[GREEN].blue=0;
	palette[BLUE].red=0;
	palette[BLUE].green=0;
	palette[BLUE].blue=-1;
	palette[PURPLE].red=-1;
	palette[PURPLE].green=0;
	palette[PURPLE].blue=-1;
	palette[YELLOW].red=-1;
	palette[YELLOW].green=-1;
	palette[YELLOW].blue=0;
	palette[TEAL].red=0;
	palette[TEAL].green=-1;
	palette[TEAL].blue=-1;
	png_set_PLTE(png_ptr,info_ptr,palette,COLORS);
	png_write_info(png_ptr,info_ptr);
	png_set_packing(png_ptr);
	png_write_image(png_ptr,(png_byte **)srd->image);
	png_write_end(png_ptr,info_ptr);
	free(palette);
	free(text_ptr);
	png_destroy_write_struct(&png_ptr,&info_ptr);
	return(0);
}

static char weekdayname[7][4]={"Sun","Mon","Tue","Wed","Thu","Fri","Sat"};
static char monthname[12][4]={"Jan","Feb","Mar","Apr","May","Jun","Jul","Aug","Sep","Oct","Nov","Dec"};
static int curgmtime(char **cur,char **next) {
	time_t t;
	struct tm *tm;

	t=time(0);
	*cur=(char *)malloc(DATESIZE);
	if(!*cur)
		return(-1);
	*next=(char *)malloc(DATESIZE);
	if(!*next)
		return(-1);
	tm=gmtime(&t);
	snprintf(*cur,DATESIZE,"%.3s, %2.2d %.3s %.4d %2.2d:%2.2d:%2.2d GMT",weekdayname[tm->tm_wday],tm->tm_mday,monthname[tm->tm_mon],1900+tm->tm_year,tm->tm_hour,tm->tm_min,tm->tm_sec);
	t+=DATA_INTERVAL;
	tm=gmtime(&t);
	snprintf(*next,DATESIZE,"%.3s, %2.2d %.3s %.4d %2.2d:%2.2d:%2.2d GMT",weekdayname[tm->tm_wday],tm->tm_mday,monthname[tm->tm_mon],1900+tm->tm_year,tm->tm_hour,tm->tm_min,tm->tm_sec);
	return(0);
}

static struct rawdata *newsrd(int width,int height,int scale,int horizlines,int cap) {
	struct rawdata *srd;
	int i;

	srd=(struct rawdata *)malloc(sizeof(struct rawdata));
	if(srd) {
		srd->width=width;
		srd->height=height;
		srd->togo=width;
		srd->samplestogo=scale-1;
		srd->max=horizlines?horizlines:100;
		srd->scale=scale;
		srd->horizlines=horizlines;
		srd->cap=cap?cap:(srd->horizlines*(srd->height-1));
		srd->oldsample=-1;
		srd->rawimage=(char *)malloc(width*height);
		if(cap)
			srd->max=cap;
		if(!srd->rawimage)
			goto fail1;
		srd->image=(char **)malloc(height*sizeof(char *));
		if(!srd->image)
			goto fail2;
		for(i=0;i<height;i++)
			srd->image[i]=srd->rawimage+i*width;
		srd->mindata=(int *)malloc(width*sizeof(int));
		if(!srd->mindata)
			goto fail3;
		for(i=0;i<width;i++)
			srd->mindata[i]=1;
		srd->maxdata=(int *)calloc(width,sizeof(int));
		if(!srd->maxdata)
			goto fail4;
		srd->auxdata=(int *)calloc(width,sizeof(int));
		if(!srd->auxdata)
			goto fail5;
		srd->nodata=(int *)calloc(width,sizeof(int));
		if(!srd->nodata)
			goto fail6;
		return(srd);
fail6:
		free(srd->auxdata);
fail5:
		free(srd->maxdata);
fail4:
		free(srd->mindata);
fail3:
		free(srd->image);
fail2:
		free(srd->rawimage);
fail1:
		free(srd);
	}
	return(0);
}

static void killsrd(struct rawdata *srd) {
	free(srd->nodata);
	free(srd->auxdata);
	free(srd->maxdata);
	free(srd->mindata);
	free(srd->rawimage);
	free(srd->image);
	free(srd);
}

static int loadfiles(struct rawdata *srd) {
	glob_t gl;
	int i,j;

	if(glob("*",GLOB_ERR,0,&gl))
		return(-1);
	for(i=gl.gl_pathc-1;i>=0 && (srd->togo || srd->samplestogo);i--) {
		j=loaddata(srd,gl.gl_pathv[i]);
		if(j<0)
			return(-1);
		if(j>0)
			break;
	}
	globfree(&gl);
	return(0);
}

static int parseargs(int *width,int *height,int *scale,int *horizlines,int *cap,int *refresh,char **dir) {
	char *s,*t;
	struct datasrc *sds;

	*width=600;
	*height=101;
	*scale=1;
	*horizlines=100;
	*cap=0;
	*refresh=-1;
	*dir=datasrcs[0].dir;
	s=getenv("QUERY_STRING");
	if(!s)
		return(0);
	while(*s) {
		if(!strncmp(s,"width=",6)) {
			s+=6;
			*width=0;
			if(isdigit(*s))
				*width=strtoul(s,&s,0);
		} else if(!strncmp(s,"height=",7)) {
			s+=7;
			*height=0;
			if(isdigit(*s))
				*height=strtoul(s,&s,0);
		} else if(!strncmp(s,"scale=",6)) {
			s+=6;
			*scale=0;
			if(isdigit(*s))
				*scale=strtoul(s,&s,0);
		} else if(!strncmp(s,"horizlines=",11)) {
			s+=11;
			*horizlines=-1;
			if(isdigit(*s))
				*horizlines=strtoul(s,&s,0);
		} else if(!strncmp(s,"cap=",4)) {
			s+=4;
			*cap=0;
			if(isdigit(*s))
				*cap=strtoul(s,&s,0);
		} else if(!strncmp(s,"refresh=",8)) {
			s+=8;
			*refresh=-1;
			if(isdigit(*s))
				*refresh=strtoul(s,&s,0);
		} else if(!strncmp(s,"datasrc=",8)) {
			s+=8;
			*dir=datasrcs[0].dir;
			t=s;
			while(isalnum(*s))
				s++;
			for(sds=datasrcs;sds->name;sds++)
				if(!strncmp(t,sds->name,s-t))
					break;
			if(sds->name)
				*dir=sds->dir;
		} else
			break;
		if(*s != '&' && *s != ';')
			break;
		s++;
	}
	if(*horizlines<0)
		*horizlines=0;
	if(*scale<1)
		*scale=1;
	if(*scale>MAX_SCALE)
		*scale=MAX_SCALE;
	if(*width<MIN_DIM)
		*width=MIN_DIM;
	if(*height<MIN_DIM)
		*height=MIN_DIM;
	if(*width>MAX_DIM)
		*width=MAX_DIM;
	if(*height>MAX_DIM)
		*height=MAX_DIM;
	return(0);
}

int main(int argc,char **argv) {
	char *curtime,*nexttime,*dir;
	int width,height,scale,horizlines,cap,refresh;
	FILE *f=0;
	struct rawdata *srd;

	if(curgmtime(&curtime,&nexttime)) {
		perror("curgmtime");
		return(-1);
	}
	if(parseargs(&width,&height,&scale,&horizlines,&cap,&refresh,&dir)) {
		fprintf(stderr,"parseargs failed\n");
		return(1);
	}
	if(isatty(2)) {
		if(argc != 2 || (argv[1][0]=='-' && argv[1][1])) {
			fprintf(stderr,"usage: %s <filename>\n",argv[0]);
			return(1);
		}
		if(strcmp(argv[1],"-")) {
			f=fopen(argv[1],"wb");
			if(!f) {
				perror("fopen");
				return(-1);
			}
		} else
			f=stdout;
	} else {
		f=stdout;
		fputs("Content-Type: image/png\r\n",f);
		fprintf(f,"Last-Modified: %s\r\n",curtime);
		fprintf(f,"Expires: %s\r\n",nexttime);
		if(refresh>=0)
			fprintf(f,"Refresh: %d\r\n",refresh);
		fputs("\r\n",f);
	}
	srd=newsrd(width,height,scale,horizlines,cap);
	if(!srd) {
		perror("newsrd");
		return(-1);
	}
	if(chdir(dir)) {
		perror("chdir");
		return(-1);
	}
	if(loadfiles(srd)) {
		perror("loadfiles");
		return(-1);
	}
	fillimage(srd);
	if(writepng(srd,f,curtime)) {
		perror("writepng");
		return(-1);
	}
	killsrd(srd);
	fclose(f);
	free(curtime);
	free(nexttime);
	return(0);
}