Add ffmpeg capture to swaygrab (with limitations)

This needs to be multithreaded to have any sort of realistic expectation
of performance, due to issues with syncronous I/O.

3 files changed, 111 insertions, 10 deletions

diff --git a/swaygrab.1.txt b/swaygrab.1.txt
index 54a1c37a..44152430 100644
--- a/swaygrab.1.txt
+++ b/swaygrab.1.txt
@@ -31,6 +31,13 @@ Options
 	Use the specified socket path. Otherwise, swaymsg will ask sway where the
 	socket is (which is the value of $SWAYSOCK, then of $I3SOCK).
 
+*-r, --rate* <rate>::
+	Specify a framerate (in frames per second). Used in combination with -c.
+	Default is 30. Must be an integer.
+
+*--raw*::
+	Instead of invoking ImageMagick or ffmpeg, dump raw rgba data to stdout.
+
 Examples
 --------
 
diff --git a/swaygrab/CMakeLists.txt b/swaygrab/CMakeLists.txt
index 5b47a694..8bc8ed8b 100644
--- a/swaygrab/CMakeLists.txt
+++ b/swaygrab/CMakeLists.txt
@@ -10,6 +10,8 @@ add_executable(swaygrab
   ${common}
 )
 
+TARGET_LINK_LIBRARIES(swaygrab rt)
+
 install(
   TARGETS   swaygrab
   RUNTIME   DESTINATION bin
diff --git a/swaygrab/main.c b/swaygrab/main.c
index c05f62cd..2c169017 100644
--- a/swaygrab/main.c
+++ b/swaygrab/main.c
@@ -4,6 +4,7 @@
 #include <getopt.h>
 #include <unistd.h>
 #include <math.h>
+#include <time.h>
 #include "log.h"
 #include "ipc-client.h"
 
@@ -21,19 +22,28 @@ int numlen(int n) {
 	return 1;
 }
 
-void grab_and_apply_magick(const char *file, const char *output, int socketfd) {
+void grab_and_apply_magick(const char *file, const char *output,
+		int socketfd, int raw) {
 	uint32_t len = strlen(output);
 	char *pixels = ipc_single_command(socketfd,
 			IPC_SWAY_GET_PIXELS, output, &len);
 	uint32_t *u32pixels = (uint32_t *)(pixels + 1);
 	uint32_t width = u32pixels[0];
 	uint32_t height = u32pixels[1];
+	len -= 9;
 	pixels += 9;
 
 	if (width == 0 || height == 0) {
 		sway_abort("Unknown output %s.", output);
 	}
 
+	if (raw) {
+		fwrite(pixels, 1, len, stdout);
+		fflush(stdout);
+		free(pixels - 9);
+		return;
+	}
+
 	const char *fmt = "convert -depth 8 -size %dx%d+0 rgba:- -flip %s";
 	char *cmd = malloc(strlen(fmt) - 6 /*args*/
 			+ numlen(width) + numlen(height) + strlen(file) + 1);
@@ -43,13 +53,76 @@ void grab_and_apply_magick(const char *file, const char *output, int socketfd) {
 	fwrite(pixels, 1, len, f);
 	fflush(f);
 	fclose(f);
-	free(pixels);
+	free(pixels - 9);
+	free(cmd);
+}
+
+void grab_and_apply_movie_magic(const char *file, const char *output,
+		int socketfd, int raw, int framerate) {
+	if (raw) {
+		sway_log(L_ERROR, "Raw capture data is not yet supported. Proceeding with ffmpeg normally.");
+	}
+
+	uint32_t len = strlen(output);
+	char *pixels = ipc_single_command(socketfd,
+			IPC_SWAY_GET_PIXELS, output, &len);
+	uint32_t *u32pixels = (uint32_t *)(pixels + 1);
+	uint32_t width = u32pixels[0];
+	uint32_t height = u32pixels[1];
+	pixels += 9;
+
+	if (width == 0 || height == 0) {
+		sway_abort("Unknown output %s.", output);
+	}
+
+	const char *fmt = "ffmpeg -f rawvideo -framerate %d "
+		"-video_size %dx%d -pixel_format argb "
+		"-i pipe:0 -r %d -vf vflip %s";
+	char *cmd = malloc(strlen(fmt) - 8 /*args*/
+			+ numlen(width) + numlen(height) + numlen(framerate) * 2 
+			+ strlen(file) + 1);
+	sprintf(cmd, fmt, framerate, width, height, framerate, file);
+
+	long ns = (long)(1000000000 * (1.0 / framerate));
+	struct timespec start, finish, ts;
+	ts.tv_sec = 0;
+
+	FILE *f = popen(cmd, "w");
+	fwrite(pixels, 1, len, f);
+	free(pixels - 9);
+	int sleep = 0;
+	while (sleep != -1) {
+		clock_gettime(CLOCK_MONOTONIC, &start);
+		len = strlen(output);
+		pixels = ipc_single_command(socketfd,
+				IPC_SWAY_GET_PIXELS, output, &len);
+		pixels += 9;
+		len -= 9;
+
+		fwrite(pixels, 1, len, f);
+
+		clock_gettime(CLOCK_MONOTONIC, &finish);
+		ts.tv_nsec = ns;
+		double fts = (double)finish.tv_sec + 1.0e-9*finish.tv_nsec;
+		double sts = (double)start.tv_sec + 1.0e-9*start.tv_nsec;
+		long diff = (fts - sts) * 1000000000;
+		sway_log(L_INFO, "%f %f %ld", sts, fts, diff);
+		ts.tv_nsec = ns - diff;
+		if (ts.tv_nsec < 0) {
+			ts.tv_nsec = 0;
+		}
+		sleep = nanosleep(&ts, NULL);
+	}
+	fflush(f);
+
+	fclose(f);
 	free(cmd);
 }
 
 int main(int argc, char **argv) {
-	static int capture = 0;
+	static int capture = 0, raw = 0;
 	char *socket_path = NULL;
+	int framerate = 30;
 
 	init_log(L_INFO);
 
@@ -57,13 +130,15 @@ int main(int argc, char **argv) {
 		{"capture", no_argument, &capture, 'c'},
 		{"version", no_argument, NULL, 'v'},
 		{"socket", required_argument, NULL, 's'},
+		{"raw", no_argument, &raw, 'r'},
+		{"rate", required_argument, NULL, 'R'},
 		{0, 0, 0, 0}
 	};
 
 	int c;
 	while (1) {
 		int option_index = 0;
-		c = getopt_long(argc, argv, "cvs:", long_options, &option_index);
+		c = getopt_long(argc, argv, "cvs:r", long_options, &option_index);
 		if (c == -1) {
 			break;
 		}
@@ -73,6 +148,15 @@ int main(int argc, char **argv) {
 		case 's': // Socket
 			socket_path = strdup(optarg);
 			break;
+		case 'r':
+			raw = 1;
+			break;
+		case 'c':
+			capture = 1;
+			break;
+		case 'R': // Frame rate
+			framerate = atoi(optarg);
+			break;
 		case 'v':
 #if defined SWAY_GIT_VERSION && defined SWAY_GIT_BRANCH && defined SWAY_VERSION_DATE
 			fprintf(stdout, "sway version %s (%s, branch \"%s\")\n", SWAY_GIT_VERSION, SWAY_VERSION_DATE, SWAY_GIT_BRANCH);
@@ -91,19 +175,27 @@ int main(int argc, char **argv) {
 		}
 	}
 
-	if (optind >= argc - 1) {
-		sway_abort("Expected output and file on command line. See `man swaygrab`");
+	char *file, *output;
+	if (raw) {
+		if (optind >= argc) {
+			sway_abort("Invalid usage. See `man swaygrab` %d %d", argc, optind);
+		}
+		output = argv[optind];
+	} else {
+		if (optind >= argc - 1) {
+			sway_abort("Invalid usage. See `man swaygrab`");
+		}
+		file = argv[optind + 1];
+		output = argv[optind];
 	}
 
-	char *file = argv[optind + 1];
-	char *output = argv[optind];
 	int socketfd = ipc_open_socket(socket_path);
 	free(socket_path);
 
 	if (!capture) {
-		grab_and_apply_magick(file, output, socketfd);
+		grab_and_apply_magick(file, output, socketfd, raw);
 	} else {
-		sway_abort("Capture is not yet supported");
+		grab_and_apply_movie_magic(file, output, socketfd, raw, framerate);
 	}
 
 	close(socketfd);