/*
 * Copyright (c) 2010 Brandyn White (bwhite@dappervision.com)
 * Copyright (c) 2016 Luke Shumaker (lukeshu@sbcglobal.net)
 *
 * This code is licensed to you under the terms of the Apache License, version
 * 2.0, or, at your option, the terms of the GNU General Public License,
 * version 2.0. See the APACHE20 and GPL2 files for the text of the licenses,
 * or the following URLs:
 * http://www.apache.org/licenses/LICENSE-2.0
 * http://www.gnu.org/licenses/gpl-2.0.txt
 *
 * If you redistribute this file in source form, modified or unmodified, you
 * may:
 *   1) Leave this header intact and distribute it under the same terms,
 *      accompanying it with the APACHE20 and GPL20 files, or
 *   2) Delete the Apache 2.0 clause and accompany it with the GPL2 file, or
 *   3) Delete the GPL v2 clause and accompany it with the APACHE20 file
 * In all cases you must keep the copyright notice intact and include a copy
 * of the CONTRIB file.
 *
 * Binary distributions must follow the binary distribution requirements of
 * either License.
 */
#include <errno.h>
#include <error.h>
#include <inttypes.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h> /* atexit, EXIT_FAILURE */
#include <unistd.h> /* dup2 */

#include <systemd/sd-daemon.h>

#include <libfreenect/libfreenect.h>
#include <libusb-1.0/libusb.h>

#include <jpeglib.h>

#include "util.h"

void stop(int sig);
#define threaderror(stat, errnum, ...) do { \
		error(0, errnum, __VA_ARGS__); \
		exitcode = stat; \
		stop(0); \
	} while(0)

struct mpjpg_encoder {
	struct jpeg_error_mgr jpeg_encoder_err;
	struct jpeg_compress_struct jpeg_encoder;
	unsigned char *jpeg_buf;
	unsigned long jpeg_len;
};

int exitcode = 0;
bool running = true;
FILE *depth_stream = NULL;
FILE *video_stream = NULL;
FILE *accel_stream = NULL;
freenect_context *ctx = NULL;
freenect_device *dev = NULL;
struct mpjpg_encoder depth_encoder;
struct mpjpg_encoder video_encoder;

void mpjpg_init(struct mpjpg_encoder *e) {
	e->jpeg_encoder.err = jpeg_std_error(&e->jpeg_encoder_err);
	jpeg_create_compress(&e->jpeg_encoder);
	e->jpeg_encoder.image_width = 640;
	e->jpeg_encoder.image_height = 480;
	e->jpeg_encoder.input_components = 3;
	e->jpeg_encoder.in_color_space = JCS_RGB;
	jpeg_set_defaults(&e->jpeg_encoder);
	e->jpeg_buf = NULL;
	e->jpeg_len = 0;
	jpeg_mem_dest(&e->jpeg_encoder, &e->jpeg_buf, &e->jpeg_len);
}

void mpjpg_destroy(struct mpjpg_encoder *e) {
	jpeg_destroy_compress(&e->jpeg_encoder);
	free(e->jpeg_buf);
}

void write_imageframe(const char *name, FILE *stream, struct mpjpg_encoder *e, JSAMPLE *data) {
	JSAMPLE *row_pointers[480];
	for (size_t i = 0; i < 480; i++)
		row_pointers[i] = &data[i*640*3];

	jpeg_start_compress(&e->jpeg_encoder, TRUE);
	while (e->jpeg_encoder.next_scanline < e->jpeg_encoder.image_height)
		jpeg_write_scanlines(&e->jpeg_encoder, &row_pointers[e->jpeg_encoder.next_scanline], 480-e->jpeg_encoder.next_scanline);
	jpeg_finish_compress(&e->jpeg_encoder);

	if (fprintf(stream,
	            "--ffserver\r\n"
	            "Content-type: image/jpeg\r\n"
	            "Content-length: %lu\r\n"
	            "\r\n",
	            e->jpeg_len) < 1)
		threaderror(EXIT_FAILURE, ferror(stream), "%s: writing header", name);
	if (fwrite(e->jpeg_buf, e->jpeg_len, 1, stream) < 1)
		threaderror(EXIT_FAILURE, ferror(stream), "%s: writing body", name);
	if (fwrite("\r\n", 2, 1, stream) < 1)
		threaderror(EXIT_FAILURE, ferror(stream), "%s: writing footer", name);
	fflush(stream);
}

void handle_accel(freenect_device *dev UNUSED, freenect_raw_tilt_state* data) {
	double x, y, z, angle;
	freenect_get_mks_accel(data, &x, &y, &z);
	angle = freenect_get_tilt_degs(data);

	const char *motor;
	switch (data->tilt_status) {
	case TILT_STATUS_STOPPED:
		motor = "stopped";
		break;
	case TILT_STATUS_LIMIT:
		motor = "limit";
		break;
	case TILT_STATUS_MOVING:
		motor = "moving";
		break;
	default:
		motor = "unknown";
	}

	char *json = NULL;
	int len = asprintf(&json,
	                 "{	\"x\": %f,\n"
	                 "	\"y\": %f,\n"
	                 "	\"z\": %f,\n"
	                 "	\"angle\": %f,\n"
	                 "	\"motor\": \"%s\" }\n",
	                 x, y, z, angle, motor);
	if (len < 0)
		threaderror(EXIT_FAILURE, errno, "accel.mjson: asprintf");
	if (fprintf(accel_stream,
	        "--ffserver\r\n"
	        "Content-type: application/json\r\n"
	        "Content-length: %d\r\n"
	        "\r\n"
	        "%s\r\n",
	            len, json) < 0)
		threaderror(EXIT_FAILURE, ferror(accel_stream), "accel.mjson: write");
	free(json);
	fflush(accel_stream);
}

uint8_t depth_rgb[640*480*3];
void handle_depth(freenect_device *dev, void *depth, uint32_t timestamp UNUSED) {
	uint32_t size = freenect_get_current_depth_mode(dev).bytes;
	if (size != 640*480*2)
		threaderror(EXIT_FAILURE, 0, "handle_depth: expected 640*480*2 byte frame, but got %"PRId32, size);
	/* scale the 11-bit values into 8-bit values */
	uint16_t *depth_grey = depth;
	for (size_t i = 0; i < 640*480; i++)
		depth_rgb[i*3+0] = depth_rgb[i*3+1] = depth_rgb[i*3+2] = (uint8_t)(depth_grey[i]*8.0/11.0);
	/* write the image */
	write_imageframe("depth.mjpg", depth_stream, &depth_encoder, depth_rgb);
}

void handle_video(freenect_device *dev, void *rgb, uint32_t timestamp UNUSED) {
	uint32_t size = freenect_get_current_video_mode(dev).bytes;
	if (size != 640*480*3)
		threaderror(EXIT_FAILURE, 0, "handle_video: expected 640*480*3 byte frame, but got %"PRId32, size);
	/* write the image */
	write_imageframe("video.mjpg", video_stream, &video_encoder, rgb);
}

void stop(int sig) {
	if (sig != 0)
		log("Caught %d", sig);
	sd_notify(0, "STOPPING=1");
	running = false;
}

FILE *xfopen(const char *path, const char *mode) {
	FILE *stream = fopen(path, mode);
	if (stream == NULL)
		error(EXIT_FAILURE, errno, "fopen: %s", path);
	return stream;
}

FILE *xfdopen(const char *path, const char *mode) {
	int fd = get_fd(path);
	if (fd < 0)
		error(EXIT_FAILURE, -fd, "get_fd: %s", path);
	FILE *stream = fdopen(fd, mode);
	if (stream == NULL)
		error(EXIT_FAILURE, errno, "fdopen: %s (%d)", path, fd);
	return stream;
}

void usage() {
	printf("Usage: %s [-h] [-v video.mjpg|-V video_fd] [-d depth.mjpg|-D depth_fd] [-a accel.mjson|-A accel_fd]\n", program_invocation_name);
}

void cleanup(void);

int main(int argc, char *argv[]) {
	int res = 0;

	mpjpg_init(&video_encoder);
	mpjpg_init(&depth_encoder);
	atexit(cleanup);

	for (int i = 1; i < argc; i++) {
		if (argv[i][0] == '-' && argv[i][1] != '\0' && argv[i][2] == '\0') {
			switch (argv[i][1]) {
			case 'h':
				usage();
				return EXIT_SUCCESS;
			case 'v':
				i++;
				video_stream = xfopen(argv[i], "w");
				break;
			case 'd':
				i++;
				depth_stream = xfopen(argv[i], "w");
				break;
			case 'a':
				i++;
				accel_stream = xfopen(argv[i], "w");
				break;
			case 'V':
				i++;
				video_stream = xfdopen(argv[i], "w");
				break;
			case 'D':
				i++;
				depth_stream = xfdopen(argv[i], "w");
				break;
			case 'A':
				i++;
				accel_stream = xfdopen(argv[i], "w");
				break;
			default:
				dup2(2, 1);
				usage();
				return EXIT_FAILURE;
			}
		} else {
			dup2(2, 1);
			usage();
			return EXIT_FAILURE;
		}
	}
	if (!(video_stream || depth_stream || accel_stream)) {
		dup2(2, 1);
		usage();
		return EXIT_FAILURE;
	}

	res = freenect_init(&ctx, 0);
	if (res) {
		error(EXIT_FAILURE, 0, "freenect_init: %s", libusb_strerror(res));
	}

	freenect_device_flags devices = 0;
	if (video_stream || depth_stream)
		devices |= FREENECT_DEVICE_CAMERA;
	if (accel_stream)
		devices |= FREENECT_DEVICE_MOTOR;

	freenect_select_subdevices(ctx, devices);
	if ((res = freenect_open_device(ctx, &dev, 0))) {
		error(EXIT_FAILURE, 0, "freenect_open_device: %s", libusb_strerror(res));
	}

	if (video_stream) {
		freenect_set_video_mode(dev, freenect_find_video_mode(FREENECT_RESOLUTION_MEDIUM, FREENECT_VIDEO_RGB));
		freenect_start_video(dev);
		freenect_set_video_callback(dev, handle_video);
	}

	if (depth_stream) {
		freenect_set_depth_mode(dev, freenect_find_depth_mode(FREENECT_RESOLUTION_MEDIUM, FREENECT_DEPTH_11BIT));
		freenect_start_depth(dev);
		freenect_set_depth_callback(dev, handle_depth);
	}

	signal(SIGPIPE, SIG_IGN);
	signal(SIGTERM, stop);
	signal(SIGQUIT, stop);
	signal(SIGINT, stop);
	while (running) {
		if (freenect_process_events(ctx) < 0) {
			stop(0);
			continue;
		}
		if (accel_stream) {
			freenect_raw_tilt_state* state;
			freenect_update_tilt_state(dev);
			state = freenect_get_tilt_state(dev);
			handle_accel(dev, state);
		}
	}
}

void cleanup() {
	if (dev) {
		freenect_stop_depth(dev);
		freenect_stop_video(dev);
		freenect_close_device(dev);
	}
	if (ctx)
		freenect_shutdown(ctx);
	if (video_stream)
		fclose(video_stream);
	if (depth_stream)
		fclose(depth_stream);
	if (accel_stream)
		fclose(accel_stream);
	mpjpg_destroy(&video_encoder);
	mpjpg_destroy(&depth_encoder);
	fflush(stderr);
	if (sd_booted() > 0) sleep(5); /* work around systemd bug dropping log messages */
	_exit(exitcode);
}