/*
 * W2000A Screenshot Utility v0.4
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <strings.h>
#ifndef WIN32
#include <fcntl.h>
#include <err.h>
#include <termios.h>
#include <time.h>
#else
#include <windows.h>
#include "ComTools.h"
#endif

#define MAXPLANES	16

static void usage(void);
static void set_serial(int);
static int open_serial(char *);
static int rdchar(int);
static int wrchar(int, char);
static void csv_hdr(FILE *);
static void csv_data(FILE *, unsigned char [4]);
static void csv_ftr(FILE *);
static void ascii_hdr(FILE *);
static void ascii_data(FILE *, unsigned char [4]);
static void ascii_ftr(FILE *);
static void receive_trace(int, int, const char *, void (*)(FILE *), void (*)(FILE *, unsigned char [4]), void (*)(FILE *));
static void receive_screen(int, int, int);
static void retrieve(int, int, int);
static void combine_screen(int);
static time_t gettime(void);

/*
 * XXX: Please note:
 *      This is _not_ 100% identical to the actual
 *      output seen on the DSO.  For example:
 *      The math channel is drawn _over_
 *      the browse bar.  Put the Math channel
 *      behind UI in the planeorder to witness
 *      the actual output.
 */
static const int planeorder[] = {
	1,			/* Grid */
	7, 8, 9, 10, 		/* Channels 1-4 */
				/* XXX: Red Arrow and Stop-Sign in UI on 10/Channel4? */
	11,			/* Math Channel */
	15, 16,			/* Cursors and Trigger Mark */
	4, 5, 6, 2, 3,		/* UI */
	12, 13, 14,		/* Not yet determined */
	0
};

static const int planecolours[][3] = {
			/* R     G     B */
/* Grid_Plane */	{ 0x3F, 0x3F, 0x3F }, /* dark gray */
/* UI_Plane1 */		{ 0xFB, 0xFB, 0xFB }, /* dark white */
/* UI_Plane2 */		{ 0x00, 0x00, 0x00 }, /* black */
/* UI_Plane3 */		{ 0xFF, 0xFD, 0xEE }, /* yellow tinted white */
/* UI_Plane4 */		{ 0xE6, 0xE6, 0x83 }, /* beige-ish yellow */
/* UI_Plane5 */		{ 0x7E, 0x7E, 0x7E }, /* gray */
/* Channel_Plane1 */	{ 0xFF, 0xFF, 0x00 }, /* yellow */
/* Channel_Plane2 */	{ 0x90, 0xEE, 0x90 }, /* light green */
/* Channel_Plane3 */	{ 0x32, 0xAB, 0xFF }, /* light blue */
/* Channel_Plane4 */	{ 0xFF, 0x00, 0x00 }, /* red */
/* Channel_Math_Plane */{ 0xFF, 0x96, 0xF6 }, /* pink */
/* Memory_Plane1 */	{ 0x00, 0x00, 0x00 }, /* n.y.d. */
/* Memory_Plane2 */	{ 0x00, 0x00, 0x00 }, /* n.y.d. */
/* Memory_Plane3 */	{ 0x00, 0x00, 0x00 }, /* n.y.d. */
/* Marker_Plane1 */	{ 0xFF, 0x67, 0x00 }, /* orange-red */
/* Marker_Plane2 */	{ 0xFF, 0x67, 0x00 }, /* orange-red */
/* Background */	{ 0x00, 0x00, 0x00 }, /* black; don't remove this! */
/* B/W Screenshot */	{ 0xFF, 0xFF, 0xFF }, /* white; don't remove this! */
};

static unsigned char pixmaps[MAXPLANES][640*480];
static char fnamprefix[256];
static int fnum;
static int rdplanes;
static time_t last_t, now_t;
static int iflag;
static char tracetype[16];

#ifdef WIN32
static void
err(int e, const char *fmt, ...)
{

	/*
	 * Don't even know if win32 has sth. like vsnprintf(),
	 * so we don't care.
	 */
	(void)fprintf(stderr, "Error: %s, exiting.\n", fmt);
	exit(e);
}

static void
errx(int e, const char *fmt, ...)
{

	err(e, fmt);
}
#endif

static void
usage(void)
{

	(void)fprintf(stderr,
			"Usage:\n"
#ifdef WIN32
			"w2000a-screenshot.exe [-h] [-c n] [-f prefix] [-n num] [-b] [-smd] [-o] [-t type]\n"
#else
			"w2000a-screenshot [-h] [-c n] [-f prefix] [-n num] [-b] [-smd] [-o] [-t type]\n"
#endif
			" -h\t\tthis help\n"
#ifdef WIN32
			" -c n\t\tCOM port (0 < n < 10)\n"
#else
			" -c n\t\tSerial device, e.g. /dev/ttyS0\n"
#endif
			" -f prefix\toutput file\n"
			" -n num\t\tdefine start of file numbering\n"
			" -b\t\twrite as BMP file, not P(P|G)M\n"
			" -s\t\tinitiate standard screenshot (colour dump)\n"
			" -m\t\tinitiate monochrome dump\n"
			" -d\t\tinitiate trace dump\n"
			" -o\t\tOne-Shot, do not wait for more dumps (set on [-smd])\n"
			" -i\t\tInvert colours (i.e. for monochromatic screenshot)\n"
			" -t type\tspecify trace file type, default: csv, valid: csv, ascii\n"
			"\n"
#ifdef WIN32
			"Input is read from specified COM port,\n"
			"the default is COM1.\n"
#else
			"Input is read from serial port, diagnostics and\n"
			"status messages are directed to stderr.\n"
			"Default is /dev/ttyS0\n"
#endif
			"\n"
			"If P(P|G)M (default) output is in use, planes\n"
			"are written to \"[prefix]-[num]_dd.pgm\", with\n"
			"\"dd\" as the two digit plane sequence number.\n"
			"Prefix defaults to \"screenshot\".\n"
			"\n"
			"Planes are coloured and combined in \"[prefix]-[num].[suf]\",\n"
			"with suffix ppm or bmp.\n"
			"\n"
			"Trace data is written to [prefix]-[num].[typesuf].\n");
	exit(1);
}

static void
set_serial(int com)
{

#ifdef WIN32
	(void)com; /* unused */
#else
	struct termios termios;

	if (tcgetattr(com, &termios) == -1)
		err(1, "tcgetattr() for serial port failed");
	if (cfsetspeed(&termios,B115200) == -1)
		err(1, "cfsetspeed() failed, unable to set desired speed (115200 Baud)");

	cfmakeraw(&termios);
	termios.c_cc[VTIME]=100;
	termios.c_cc[VMIN]=0;
	if (tcsetattr(com, TCSANOW, &termios) == -1)
		err(1, "tcsetattr() failed, unable to set line attributes");
#endif
}

static int
open_serial(char *device)
{
	int com;

#ifdef WIN32
	if (device == NULL)
		com = 1;
	else
		com = *device - '0';
	if (!ComOpen(com - 1, 115200, P_NONE, S_1BIT, D_8BIT))
		err(1, "Unable to open COM port");
#else
	if (device == NULL)
		device = (char *)"/dev/ttyS0";
	if ((com = open(device, O_RDWR)) == -1)
		err(1, "open(%s) failed", device);
	set_serial(com);
#endif

	return com;
}

static int
rdchar(int com)
{
#ifdef WIN32
	while (ComGetReadCount(com - 1) < 1)
		(int)usleep(1000);
	return ComRead(com - 1);
#else
	unsigned char c;

	if (read(com, &c, 1) > 0)
		return c;
	return -1;
#endif
}

static int
wrchar(int com, char c)
{

#ifdef WIN32
	return ComWrite(com - 1, c);
#else
	return write(com, &c, 1);
#endif
}

static void
csv_hdr(FILE *fp)
{

	(void)fp; /* unused */
}

static void
csv_data(FILE *fp, unsigned char c[4])
{

	(void)fprintf(fp, "%d;%d;%d;%d\n", c[0], c[1], c[2], c[3]);
}

static void
csv_ftr(FILE *fp)
{

	(void)fp; /* unused */
}

static void
ascii_hdr(FILE *fp)
{

	(void)fprintf(fp, "Chan1 Chan2 Chan3 Chan4\n");
}

static void
ascii_data(FILE *fp, unsigned char c[4])
{

	(void)fprintf(fp, "%d %d %d %d\n", c[0], c[1], c[2], c[3]);
}

static void
ascii_ftr(FILE *fp)
{

	(void)fp; /* unused */
}

static void
receive_trace(int com, int type, const char *ftype, void fun_hdr(FILE *), void fun_data(FILE *, unsigned char [4]), void fun_ftr(FILE *))
{
	FILE *fp;
	unsigned char c[4];
	char fnam[256];
	int r, chan, idx;

	(void)type; /* unused */

	(void)snprintf(fnam, sizeof(fnam), "%s-%.4d.%s", fnamprefix, fnum, ftype);
	(void)fprintf(stderr, "  - Receiving trace data (file: %s)\n", fnam);
	if ((fp = fopen(fnam, "w+")) == NULL)
		err(1, "fopen(%s, w+) failed", fnam);
	fnum++;

	fun_hdr(fp);

	for (idx=0;idx<16384;idx++) {
		for (chan=0;chan<4;chan++) {
			r = rdchar(com);
			c[chan] = (unsigned char)r;
		}
		fun_data(fp, c);
	}

	fun_ftr(fp);

	if (fclose(fp) != 0)
		err(1, "fclose(%s) failed", fnam);
}

static void
receive_screen(int com, int all, int type)
{
	FILE *fp;
	char fnam[256];
	unsigned char *buffer, *bufptr, *bufend;
	size_t bufsize;
	unsigned int llength;			/* RLE */
	unsigned char lastbyte, esc;
	int outcnt, incnt;
	int r, pnum, rd, total, ppos;
	unsigned char c;
	unsigned int i, j;

	switch (type) {
	case 0:
		bufsize=640*480*MAXPLANES/8;
		(void)fprintf(stderr, "  - Receiving a colour screenshot\n");
		break;
	case 1:
		bufsize=640*480/8;
		(void)fprintf(stderr, "  - Receiving a monochrome screenshot\n");
		break;
	default:
		errx(1, "Internal Error:  Unhandled Dumptype in receive_screen(): %d", type);
	}

	/*********************************************************************
	* Allocate enough space for all planes
	*********************************************************************/
	if ((buffer = (unsigned char *)malloc(bufsize)) == NULL)
		err(1, "malloc %ld failed", bufsize);
	bufptr = buffer;

	/*********************************************************************
	* Read all from stdin/COMx, decompress and fill buffer
	*********************************************************************/
	incnt = outcnt = 0;
	llength = 0;
	lastbyte = esc = 0;
	while (((r = rdchar(com)) >= 0) && (bufptr<buffer+bufsize)) {
		c=(unsigned char)r;
		incnt++;
		if (incnt % 100 == 0)
			(void)fprintf(stderr, "\r    %d bytes...", incnt);

		if (esc==2) {
			llength|=(c);					//2nb Byte after ESC is MSB
			for (i=0;i<llength;i++) {
				*bufptr++=lastbyte;
				outcnt++;
			}
			llength=0;
			esc=0; }
		else if (esc==1) {					//2nd escape
			if (c==0xFF) {					//escaped 0xFF
				*bufptr++=c;
				outcnt++;
				lastbyte=c;
				llength=0;
				esc=0;					//Escape-Sequence finished
			} else {					//1st Byte after ESC is MSB
				llength=c<<8;
				esc++; }
			}
		else if (c != 0xFF) {
			*bufptr++=c;
			outcnt++;
			lastbyte=c;
		} else
			esc=1;
	}
	(void)fprintf(stderr, "\n");
	bufend=bufptr;
	bufptr=buffer;

	(void)fprintf(stderr, "* Total bytes transferred: %d(%d) ", outcnt, incnt);
	now_t = gettime();
	(void)fprintf(stderr, "in %lds\n", now_t-last_t);
	last_t = now_t;

	total = 0;
	rdplanes = 0;
	fp = NULL;
	for (pnum = 0; pnum < MAXPLANES; pnum++) {
		(void)fprintf(stderr, "  - Processing Plane #%.2d...\n", pnum+1);

		if (all) {
			(void)snprintf(fnam, sizeof(fnam), "%s-%.4d_%.2d.pgm", fnamprefix, fnum, pnum+1);
			if ((fp = fopen(fnam, "w+")) == NULL)
				err(1, "fopen(%s, w+) failed", fnam);

			(void)fprintf(fp, "P5\n640 480\n255 ");
		}

		/*********************************************************************
		* get data from buffers
		*********************************************************************/
		rd = 0;
		ppos = 0;
		while (bufptr<bufend) {
			c = *bufptr++;;
			rd++;
			for (j=0;j<8;j++) {
				pixmaps[pnum][ppos++] = ((c & (1<<j)) ? 255 : 0);
				if (fp) (void)fputc(((c & (1<<j)) ? 255 : 0), fp);
			}

			if (rd==640*480/8)
				break;
		}
		total += rd;
		rdplanes++;

		(void)fprintf(stderr, "\r    Read %d bytes.\n", rd);

		if (all && fclose(fp) != 0)
			err(1, "fclose(%s) failed", fnam);

		if (bufptr<=bufend) {
			c=*bufptr++;
			if (c == 'E') {
				(void)fprintf(stderr, "* End of Transmission\n");
				break;
			} else if (c == 0xFF)
				(void)fprintf(stderr, "    (another plane follows)\n");
		} else {
			(void)fprintf(stderr, "* Short read, ignoring\n");
			break;
		}
	}

	free(buffer);
}

static void
retrieve(int all, int com, int dumptype_req)
{
	int dumptype, r;
	unsigned char c, l;

	if (dumptype_req != -1) {
		(void)fprintf(stderr, "* Triggering Dump...\n");

		c = 0x02;
		if (wrchar(com, c) != 1)
			err(1, "Unable to write to COM port (at 1)");
		#ifdef WIN32
			Sleep(1000);
		#else
			(void)sleep(1);
		#endif

		if (dumptype_req == 0)
			c = 0x01;
		else if (dumptype_req == 1)
			c = 0x02;
		else if (dumptype_req == 2)
			c = 0x03;
		else if (dumptype_req != -1)
			errx(1, "Internal Error:  Unhandled Dumptype in retrieve() (at 1): %d", dumptype_req);

		if (wrchar(com, c) != 1)
			err(1, "Unable to write to COM port (at 2)");
	}

	(void)fprintf(stderr, "* Waiting for Screenshot or Trace...\n");

	last_t = gettime();
	l = 0;
	while ((r = rdchar(com)) > 0) {
		c = (unsigned char)r;
		if (c == 255 && l == 'S')
			break;
		l = c;
	}
	if (r <= 0)
		exit(1);

	(void)fprintf(stderr, "* Found Data Start Marker");
	now_t = gettime();
	(void)fprintf(stderr, " after %lds\n", now_t-last_t);
	last_t = now_t;

	/*********************************************************************
	* Next char determines type of data
	*********************************************************************/
	dumptype = (int)rdchar(com);

	switch (dumptype) {
	case 0: /* FALLTHRU */
	case 1:
		receive_screen(com, all, dumptype);
		combine_screen(!all);
		break;
	case 2:
		if (strcasecmp(tracetype, "csv") == 0)
			receive_trace(com, dumptype, "csv", csv_hdr, csv_data, csv_ftr);
		else if (strcasecmp(tracetype, "ascii") == 0)
			receive_trace(com, dumptype, "asc", ascii_hdr, ascii_data, ascii_ftr);
		else
			errx(1, "Internal Error:  Unhandled Trace Filetype in retrieve() (at 2): %s", tracetype);
		break;
	default:
		errx(1, "Internal Error:  Unhandled Dumptype in retrieve() (at 2): %d", dumptype);
	}
}

static void
combine_screen(int bmp)
{
	FILE *fp;
	char fnam[256], suf[4];
	int ppos, i, plane, colour;

	(void)snprintf(suf, sizeof(suf), "ppm");
	if (bmp)
		(void)snprintf(suf, sizeof(suf), "bmp");

	(void)snprintf(fnam, sizeof(fnam), "%s-%.4d.%s", fnamprefix, fnum, suf);
	(void)fprintf(stderr, "* Combining (file: %s)...\n", fnam);
	if ((fp = fopen(fnam, "w+")) == NULL)
		err(1, "fopen(%s, w+) failed", fnam);
	fnum++;

	if (bmp) {
		/*
		 * http://en.wikipedia.org/wiki/BMP_file_format
		 */
		/* 0x00 */ (void)fprintf(fp, "BM");
		/* 0x02 */ (void)fprintf(fp, "%c%c%c%c%c%c%c%c", 0, 0, 0, 0, 0, 0, 0, 0);
		/* 0x0A */ (void)fprintf(fp, "%c%c%c%c", 54, 0, 0, 0);		/* 40 (header size) + 0x0A + 4 */
		/* 0x0E */ (void)fprintf(fp, "%c%c%c%c", 40, 0, 0, 0);		/* header size */
		/* 0x12 */ (void)fprintf(fp, "%c%c%c%c", 128, 2, 0, 0);		/* 640px */
		/* 0x16 */ (void)fprintf(fp, "%c%c%c%c", 32, 254, 255, 255);	/* -480px */
		/* 0x1A */ (void)fprintf(fp, "%c%c", 1, 0);
		/* 0x1C */ (void)fprintf(fp, "%c%c", 24, 0);			/* bits per pixel */
		/* 0x1E */ (void)fprintf(fp, "%c%c%c%c", 0, 0, 0, 0);		/* no compression */
		/* 0x22 */ (void)fprintf(fp, "%c%c%c%c", 0, 16, 14, 0);		/* bitmap data size: 640*480*3=921600 Bytes */
		/* 0x26 */ (void)fprintf(fp, "%c%c%c%c", 0, 0, 0, 0);
		/* 0x2A */ (void)fprintf(fp, "%c%c%c%c", 0, 0, 0, 0);
		/* 0x2E */ (void)fprintf(fp, "%c%c%c%c", 0, 0, 0, 0);
		/* 0x32 */ (void)fprintf(fp, "%c%c%c%c", 0, 0, 0, 0);
	} else {
		/*
		 * http://netpbm.sourceforge.net/doc/ppm.html
		 */
		(void)fprintf(fp, "P6\n640 480\n255 ");
	}

	#define INVP(colourcomp)	iflag ? (255 - (colourcomp)) : (colourcomp)

	for (ppos = 0; ppos < 640*480; ppos++) {
		colour = MAXPLANES; /* default colour, black */
		for (i = 0; i < MAXPLANES && i < rdplanes; i++) {
			plane = planeorder[i] - 1;
			if (pixmaps[plane][ppos] == 0xFF) {
				colour = plane;
				if (rdplanes <= 2)
					colour = MAXPLANES + 1;
			}
		}
		if (bmp) {
			/* colour order BGR */
			(void)fputc(INVP(planecolours[colour][2]), fp);
			(void)fputc(INVP(planecolours[colour][1]), fp);
			(void)fputc(INVP(planecolours[colour][0]), fp);
		} else {
			/* colour order RGB */
			(void)fputc(INVP(planecolours[colour][0]), fp);
			(void)fputc(INVP(planecolours[colour][1]), fp);
			(void)fputc(INVP(planecolours[colour][2]), fp);
		}
	}

	if (fclose(fp) != 0)
		err(1, "fclose(%s) failed", fnam);

	(void)fprintf(stderr, "* Done\n");
}

static time_t
gettime(void)
{
#ifdef WIN32
	/*
	 * Yes, this is crude.  But I won't implement
	 * a decent API just for Windows.
	 */
	SYSTEMTIME st;

	GetSystemTime(&st);
	return st.wDay * 24*3600 + st.wHour * 3600 + st.wMinute * 60 + st.wSecond;
#else
	return time(NULL);
#endif
}

int
main(int argc, char **argv)
{
	int ch, bflag, oneshot, dumptype;
	char *device;
	int com;

	device = NULL;
	bflag = iflag = 0;
	oneshot = 0;
	dumptype = -1;
	fnum = 0;
	(void)snprintf(tracetype, sizeof(tracetype), "csv");
	(void)snprintf(fnamprefix, sizeof(fnamprefix), "screenshot");
	while ((ch = getopt(argc, argv, "c:f:bhsmdon:it:")) != -1) {
		switch (ch) {
		case 'c':
#ifdef WIN32
			device = optarg;
			if (*device - '0' < 1 || *device - '0' > 9 || strlen(device) > 1)
				errx(1, "invalid COM port given, use -c n, where 0 < n < 10");
#else
			device = optarg;
#endif
			break;
		case 'f':
			(void)snprintf(fnamprefix, sizeof(fnamprefix), "%s", optarg);
			break;
		case 'd':
			dumptype = 2;
			break;
		case 'm':
			dumptype = 1;
			break;
		case 's':
			dumptype = 0;
			break;
		case 'b':
			bflag = 1;
			break;
		case 'o':
			oneshot = 1;
			break;
		case 'n':
			fnum = atoi(optarg);
			if (fnum < 0 && fnum > 9999)
				errx(1, "invalid start of file numbering given, use -n num, where 0 <= num <= 9999");
			break;
		case 'i':
			iflag = 1;
			break;
		case 't':
			(void)snprintf(tracetype, sizeof(tracetype), "%s", optarg);
			if (strcasecmp(tracetype, "csv") != 0 && strcasecmp(tracetype, "ascii") != 0)
				errx(1, "invalid type given, valid: -t ascii and -t csv");
			break;
		case 'h':
		default:
			usage();
			/* NOTREACHED */
		}
	}

	if (dumptype != -1)
		oneshot = 1;

	com = open_serial(device);
	set_serial(com);

	do {
		retrieve(!bflag, com, dumptype);
	} while (!oneshot);

	return 0;
}

 	  	 
