/*
 * capture.c: capture data from stdout and stderr
 *
 *
 * Copyright (C) 2007-2012 Michael Brunner <mibru@gmx.de>
 *
 * 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
 * General Public License for more details.
 */

#include "logapp.h"

#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <termios.h>

#include <sys/socket.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <signal.h>
#include <time.h>
#include <pthread.h>

#include "configuration.h"
#include "logfile.h"
#include "capture.h"

int handle_cbstream(pipe_t* pipe);

void reset_console(void)
{
	int err = 0;;

	if (!config.dumbterm) {
		err = write(app.pstdout->cfhno, app.pstdout->escreset,
		      app.pstdout->escresetlen);
		err = write(app.pstderr->cfhno, app.pstderr->escreset,
		      app.pstderr->escresetlen);
	}

	if (app.ptytermios_bak != NULL)
		tcsetattr(STDIN_FILENO, TCSADRAIN, app.ptytermios_bak);

	tcdrain(STDOUT_FILENO);
	tcdrain(STDERR_FILENO);

	if (err == -1) {
		warning("Problem reseting console...");
	}
}

static int recalc_cliplen(pipe_t *pipe)
{
	int i;
	int length=0;

	for (i=0; i<pipe->bfill; i++) {
		switch(pipe->buf[i]) {
			case '\t':	length += (8-length%8);
					break;
			default:	length++;
					break;
		}
		if (length>pipe->clip)
			return i--;
	}

	return pipe->clip;

}

int print_stream(pipe_t *pipe)
{
	static pipe_t *pipe_old;
	static unsigned	offset;
	char* buf;
	unsigned ccount;
	unsigned esclength = 0;
	int regmatch = 0;

	if (pipe->bfill < 1)
		return 1;

	if (!config.dumbterm) {
		/* Check for regular expression match */
		if (pipe->regexp) {
			if (!(regexec(&pipe->preg, pipe->buf, 0, NULL, 0))) {
				regmatch = 1;
			}
		}

		/* Change font style if we are switching the stream */
		if (((void*)pipe_old != (void*)pipe)||(regmatch)) {
			offset = 0;
			buf = pipe->esccolor;
			esclength = pipe->buf - pipe->esccolor;
			if (pipe->bgesccolor) {
				if (!regmatch) {
					*pipe->bgesccolor = pipe->bgcol + '0';
					pipe_old = pipe;
				} else {
					*pipe->bgesccolor = pipe->regbgcol
						+ '0';
					pipe_old = NULL;
				}
			} else {
				pipe_old = pipe;
			}
		} else {
			buf = pipe->buf;
		}
	} else {
		buf = pipe->buf;
	}

	/* Clip output if configured */
	if (pipe->eclip) {
		unsigned int clip; /* clip len  should  be unsigned by now */

		/* recalculate clipping if there are tabs in the stream */
		if (pipe->recalclen) {
			clip = recalc_cliplen(pipe);
		} else {
			clip = pipe->clip;
		}
		
		if (offset >= clip)
			return 0;

		if ((pipe->bfill + offset) > clip) {
			ccount = clip - offset;
			pipe->buf[ccount] = '\r';
			ccount++;
			pipe->buf[ccount] = '\n';
			ccount++;
			offset = 0;
		} else {
			ccount = pipe->bfill;
			if (pipe->buf[ccount - 1] == '\n') {
				offset = 0;
			} else {
				offset += ccount;
			}
		}
	} else {
		ccount = pipe->bfill;
	}

	if (regmatch && !config.dumbterm) {
		/* Prevent colorizing the next line if regexp matches */
		if (pipe->buf[ccount-1] == '\n') {
			int rlen = pipe->escresetlen + 1;
			strcpy(&pipe->buf[ccount-1],  pipe->escreset);
			strcpy(&pipe->buf[ccount-1 + rlen],  "\n");
			esclength+=rlen;
			regmatch = 1;
		} else {
			regmatch = 0;
		}
	}

	if (write(pipe->cfhno, buf, ccount + esclength) == -1)
		return -1;

	if (regmatch && !config.dumbterm) {
		/* Restore \n */
		if (regmatch) {
			pipe->buf[ccount-1] = '\n';
		}
	}

	pipe->recalclen = 0;

	return 1;
}

void print_summary(void)
{
	time_t sec;
	unsigned int executiontime = 0;

	time(&sec); 
	
	if ((app.starttime > -1)&&(app.starttime <= sec)) {
		executiontime = (unsigned int)(sec - app.starttime);
	} else {
		error("Error calculating execution time\n");
	}
		
	fprintf(stderr, "\nLogapp exited after %us; ", executiontime);
	
	if (!app.pstdout->charbased) {
		fprintf(stderr, "%u", app.pstdout->linecount);
	} else {
		fprintf(stderr, "?");
	}

	fprintf(stderr, " stdout, ");

	if (!app.pstderr->charbased) {
		fprintf(stderr, "%u", app.pstderr->linecount);
	} else {
		fprintf(stderr, "?");
	}

	fprintf(stderr, " stderr lines;");

	if (WIFEXITED(app.exit_state)) {
		fprintf(stderr, " exit state %u\n",
			WEXITSTATUS(app.exit_state));
	} else if (WIFSIGNALED(app.exit_state)) {
		fprintf(stderr, " signal %d\n",
			WTERMSIG(app.exit_state));
	} else {
		fprintf(stderr, " unknown exit state\n");
	}

	/* Show a message if we changed the logfile name */
	if ((logfile.appendnr && config.warnlogfilelock)
	    || config.printlogname)
		fprintf(stderr, "Logapp logfile is: %s\n", logfile.name);
}

int is_rdbuf_empty(pipe_t *pipe)
{
	return (pipe->rdbuf.read >= pipe->rdbuf.end);
}

int buf_get_char(pipe_t *pipe, char* buf)
{
	int i;

	/* Fill buffer if it is empty */
	if (is_rdbuf_empty(pipe)) {
		pipe->rdbuf.read = pipe->rdbuf.start;

		if (pipe->switchtocbmode) {
			pipe->charbased = 1;
			pipe->handler = (int(*)(void*)) handle_cbstream;
			return 0;
		}

		i = read(pipe->fh, pipe->rdbuf.start, pipe->rdbuf.len);
		if (i < 1) {
			pipe->rdbuf.end = pipe->rdbuf.start;
			return i;
		}

		pipe->rdbuf.end = pipe->rdbuf.start + i;
	}

	/* Copy one char into the provided buffer and increase read pointer */
	*buf = *pipe->rdbuf.read;

	pipe->rdbuf.read++;

	return 1;
}

int read_pipe(pipe_t *pipe)
{
	int i = 0;
	
	if (pipe->state > 0) {
		pipe->recalclen = 0;
		pipe->bfill = 0;
	}

	/* If a CR is left from the last read, put it into the buffer again */
	if (CONFIG_SUPPORT_PTY && pipe->memcr) {
		*pipe->buf = '\r';
		pipe->bfill++;
		pipe->memcr = 0;
	}

	while(1) {
		i = buf_get_char(pipe, &pipe->buf[pipe->bfill]);
		if (i < 1)
			return(pipe->bfill);

		if (pipe->buf[pipe->bfill] == '\t')
			pipe->recalclen = 1;

		if (pipe->buf[pipe->bfill] == '\n'
		    || ( i + pipe->bfill >= pipe->blen)
		    || ((pipe->buf[pipe->bfill] == 27) && pipe->detectescape)
		    ) {
			if ((pipe->buf[pipe->bfill] == '\n') && (i > 0)) {
				/* Change CR-LF into LF if configured */
				if (CONFIG_SUPPORT_PTY && pipe->ptyremovecr
				    && (pipe->bfill > 0)) {
					if (pipe->buf[pipe->bfill-1] =='\r'){
						pipe->bfill--;
						pipe->buf[pipe->bfill]='\n';
					}
				}
				pipe->bfill++;
				pipe->linecount++;
			} else if ((pipe->buf[pipe->bfill] == 27)
				   && (pipe->detectescape) && (i > 0)) {
				pipe->bfill++;
				/* Switch to charbased stream handling as we
				 * just detected an escape sequence */
				pipe->switchtocbmode = 1;
			} else if ( i + pipe->bfill >= pipe->blen) {
				pipe->bfill++;
			}

			/* If there is a lonely CR at the end of the buffer,
			 * ignore it this time until the next read */
			if (pipe->buf[pipe->bfill-1] == '\r') {
				pipe->memcr = 1;	
				pipe->bfill--;
			}

			/* NULL terminate string, as needed for regexec */
			pipe->buf[pipe->bfill] = '\0';

			return(pipe->bfill);
		};
		if (pipe->bfill < pipe->blen)
			pipe->bfill++;
	}

	return 0;
}

int prepare_buffer(pipe_t* pipe)
{
	if (pipe->blen) {
		/* length = buffer length + escape prefix + escape postfix +
		 * CR/LF */
		int len = pipe->blen + pipe->esccolorlen + 2;

		if (!(pipe->dbuf = (char*)malloc(len))) {
			error_outofmemory();
			error("trying to allocate %d byte "
			      "%s buffer\n", len, pipe->name);
			return -1;
		}

		/* in line-based mode an intermediate ring-buffer is used for
		 * faster data processing */
		if (!pipe->charbased) {
			struct stat statbuf;
			/* get pipe stats to determine optimal buffer size*/
			if (fstat(pipe->fh, &statbuf)) {
				error("unable to get io buffer stats for %s\n",
				      pipe->name);
				free(pipe->dbuf);
				pipe->dbuf = NULL;
				return -1;
			}
			pipe->rdbuf.len = statbuf.st_blksize;

			pipe->rdbuf.start = (char*)malloc(pipe->rdbuf.len);
			if (!pipe->rdbuf.start) {
				error_outofmemory();
				error("trying to allocate %d "
				      "byte %s read buffer\n",
				      pipe->rdbuf.len, pipe->name);
				free(pipe->buf);
				pipe->buf = NULL;
				return -1;
			}
			pipe->rdbuf.read = pipe->rdbuf.start;
			pipe->rdbuf.end = pipe->rdbuf.start;
		}

		/* We put the color escape sequence directly before the line
		 * buffer memory to be able to squeeze it all out at once.
		 * The reserved space at the end is for a console reset
		 * sequence to end the regexp match background color */
		strcpy(pipe->dbuf, pipe->esccolor);
		free(pipe->esccolor);
		if (pipe->bgesccolor) {
			pipe->bgesccolor = pipe->dbuf + (pipe->bgesccolor
				- pipe->esccolor); /* fixup bg color pointer */
		}
		pipe->esccolor = pipe->dbuf;
		pipe->buf = pipe->dbuf + pipe->esccolorlen;
	};

	return 0;
}

void free_buffer(pipe_t* pipe)
{
	if (pipe->dbuf) {
		free(pipe->dbuf);
		if (pipe->dbuf == pipe->esccolor)
			pipe->esccolor = NULL;
		pipe->dbuf = NULL;
		pipe->buf = NULL;
	}

	if (pipe->rdbuf.start) {
		free(pipe->rdbuf.start);
		pipe->rdbuf.start = NULL;
		pipe->rdbuf.read = NULL;
		pipe->rdbuf.end = NULL;
	}

}

int handle_cbstream(pipe_t* pipe)
{
	/* If a CR is left from the last read, put it into the buffer again */
	if (CONFIG_SUPPORT_PTY && pipe->memcr) {
		*pipe->buf = '\r';
	}

	if ((pipe->bfill = read(pipe->fh, pipe->buf+pipe->memcr,
				pipe->blen-pipe->memcr)) > 0) {

		pipe->bfill += pipe->memcr;

		/* Change CR-LF into LF if configured */
		if (CONFIG_SUPPORT_PTY && pipe->ptyremovecr) {
			int i;
			int lbo = 0;
			char* bufa = pipe->buf;
			char* bufb = pipe->buf;

			for (i=0; i<(pipe->bfill-1); i++) {
				*bufb = *bufa;
				bufa++;
				bufb++;
				if (*bufa == '\n') {
					if (*(bufa-1) == '\r')
						bufb--;
					if ((config.alignlinebreaks)
					    && (lbo<LINEBREAKBUFSIZE)) {
						pipe->linebreakpos[lbo++] = 
							bufb - pipe->buf + 1;
					}
				}
			}
			pipe->bfill -= (bufa-bufb);
			*bufb = *bufa;
			/* linebreakpos buffer is defined size + 1 for the 0 */
			pipe->linebreakpos[lbo] = 0;

			/* if there is no linebreak detected so far, add
			 * length of the packet to the array */
			if (!pipe->linebreakpos[0]) {
				pipe->linebreakpos[0] = pipe->bfill;
				pipe->linebreakpos[1] = 0;
			}

			/* If there is a lonely CR at the end of the buffer,
			 * ignore it this time until the next read */
			if (pipe->buf[pipe->bfill-1] == '\r') {
				pipe->memcr = 1;	
				pipe->bfill--;

				if (pipe->bfill <= 0) {
					return 0;
				}
			} else {
				pipe->memcr = 0;
			}
		}

		/* Write to console */
		if (write(pipe->cfhno, pipe->buf, pipe->bfill) == -1) {
			error("Problem writing data to screen for %s\n",
			      pipe->name);
			return -1;
		}

		/* Write to logfile */
		if (logfile_write(app.logfile, pipe)) {
			error("error writing logfile\n");
			return -1;
		}
	} else {

		/* Pipe didn't provide data - exit if
		 * application isn't active anymore */
		if (!app.active) {
			/* Print remembered CR before exiting the
			 * application */
			if (CONFIG_SUPPORT_PTY && pipe->memcr) {
				pipe->bfill = 1;

				/* Write to console */
				if (write(pipe->cfhno, pipe->buf, pipe->bfill)
				    == -1) {
					error("Problem writing data to screen "
					      "for %s\n", pipe->name);
					return -1;
		}

				/* Write to logfile */
				if (logfile_write(app.logfile, pipe)) {
					error("error writing logfile\n");
					return -1;
				}
			}
			return 1;
		}

		return 2;
	}

	return 0;
}

int handle_lbstream(pipe_t* pipe)
{
	if ((pipe->state = read_pipe(pipe)) > 0) {
		int clearbuf;

		if (logfile_write(app.logfile, pipe)) {
			error("error writing logfile\n");
			return -1;
		}

		if ((clearbuf = print_stream(pipe)) < 0) {
			error("error printing %s output\n", pipe->name);
			return -1;
		}

		/* Reset buffer len after printing and logging one line */
		if (clearbuf)
			pipe->bfill = 0;

		/* Check if we should execute a command */
		if (pipe->execregexp) {
			if (!regexec(&pipe->pexecreg, pipe->buf, pipe->bfill,
				     NULL, 0)) {
				int ret = execcmd(pipe->execcommand);
				if (ret) {
					print(config.exitonexecfail
					      ?ERROR:WARNING,
					      "command \"%s\" returned 0x%x\n",
					      pipe->execcommand, ret);
					if (config.exitonexecfail) {
						exit(EXIT_FAILURE);
					}
				}
			}
		}
	} else {
		/* Don't exit if there is data still in the buffer */
		if (!is_rdbuf_empty(pipe))
			return 0;

		/* Pipe didn't provide data - exit if
		 * application isn't active anymore */
		if (!app.active)
			return 1;
		return 2;
	}

	return 0;
}

void* capture_thread(void* do_pipe)
{
	fd_set fd_pipe;
	pipe_t* pipe = (pipe_t*) do_pipe;


	FD_ZERO(&fd_pipe);
	FD_SET(pipe->fh, &fd_pipe);

	if (pipe->charbased) {
		pipe->handler = (int(*)(void*)) handle_cbstream;
	} else {
		pipe->handler = (int(*)(void*)) handle_lbstream;
	}

	while(1) {
		if (select(pipe->fh + 1, &fd_pipe, NULL, NULL, NULL) == 1) {
			switch (pipe->handler(pipe)) {
				case -1:	reset_console();
					 	exit(EXIT_FAILURE);
					 	break;
				case 1:		if (CONFIG_USE_THREADS)
							pthread_exit(NULL);
						break;
				case 2:		usleep(10000);
						break;
				default:	break;
			}
		} else {
			if (CONFIG_USE_THREADS)
				pthread_exit(NULL);
		}
	}
}

int capture_loop(pipe_t* pipe_stdout, pipe_t* pipe_stderr)
{
	fd_set	rdfs;
	long	flags = 0;
	int	do_stdout_loop = 1;
	int	do_stderr_loop = 1;

	if (pipe_stdout->charbased) {
		pipe_stdout->handler = (int(*)(void*)) handle_cbstream;
	} else {
		pipe_stdout->handler = (int(*)(void*)) handle_lbstream;
	}

	if (pipe_stderr->charbased) {
		pipe_stderr->handler = (int(*)(void*)) handle_cbstream;
	} else {
		pipe_stderr->handler = (int(*)(void*)) handle_lbstream;
	}

	fcntl(pipe_stdout->fh, F_GETFL, flags);
	fcntl(pipe_stdout->fh, F_SETFL, flags | O_NONBLOCK);
	fcntl(pipe_stderr->fh, F_GETFL, flags);
	fcntl(pipe_stderr->fh, F_SETFL, flags | O_NONBLOCK);

	while (do_stdout_loop || do_stderr_loop)
	{
		FD_ZERO(&rdfs);
		FD_SET(pipe_stdout->fh, &rdfs);
		FD_SET(pipe_stderr->fh, &rdfs);

		if (select(FD_SETSIZE, &rdfs, NULL, NULL, NULL) > 0)
		{
			if (FD_ISSET(pipe_stdout->fh, &rdfs))
			{
				switch (pipe_stdout->handler(pipe_stdout)) {
					case -1:	reset_console();
							do_stdout_loop = 0;
							break;
					case 1:		do_stdout_loop = 0;
							break;
					case 2:		
					default:	break;
				}
			}
			if (FD_ISSET(pipe_stderr->fh, &rdfs))
			{
				switch (pipe_stdout->handler(pipe_stderr)) {
					case -1:	reset_console();
							do_stderr_loop = 0;
							break;
					case 1:		do_stderr_loop = 0;
							break;
					case 2:		
					default:	break;
				}
			}
		}
		else
		{
			do_stdout_loop = 0;
			do_stderr_loop = 0;
		}
	}


	return 0;
}

int capture_start(void)
{
	if (prepare_buffer(app.pstdout)) {
		exit(EXIT_FAILURE);
	}
	if (prepare_buffer(app.pstderr)) {
		exit(EXIT_FAILURE);
	}

	if (CONFIG_USE_THREADS) {
		pthread_create(&app.pstderr->ct, NULL, capture_thread,
			       app.pstderr);
		pthread_create(&app.pstdout->ct, NULL, capture_thread,
			       app.pstdout);
	} else {
		return capture_loop(app.pstdout, app.pstderr);
	}

	return 0;
}

int capture_end(void)
{
	if (CONFIG_USE_THREADS) {
		/* wait for capture threads to exit */
		pthread_join (app.pstderr->ct, NULL);
		pthread_join (app.pstdout->ct, NULL);
	}

	free_buffer(app.pstdout);
	free_buffer(app.pstderr);

	return 0;
}
