/* $Id: encoder.c,v 1.10 2009-02-12 10:29:25 potyra Exp $
 *
 * Copyright (C) 2008-2009 FAUmachine Team <info@faumachine.org>.
 * This program is free software. You can redistribute it and/or modify it
 * under the terms of the GNU General Public License, either version 2 of
 * the License, or (at your option) any later version. See COPYING.
 */

#include "config.h"

#include <stdlib.h>
#include <stdio.h>
#include <theora/theora.h>
#include <assert.h>
#include <time.h>
#include <string.h>
#include <errno.h>
#include <limits.h>
#include <inttypes.h>
#include <stdbool.h>

/* for MAX_HEIGHT/MAX_WIDTH */
#include "sig_opt_rgb.h"
#include "rd_types.h"

#define FPS_NUM   12
#define FPS_DENOM  1

/** faum_encoder state structure */
struct faum_encoder {
	/** theora state handle */
	theora_state th;
	/** ogg stream state handle */
	ogg_stream_state vo;

	/** output file */
	FILE *ogg_out;
	/** stream input file */
	FILE *rec_in;

	/** theora info handle */
	theora_info ti;

	/** rgb pixel buffer */
	unsigned char rgb_buf[MAX_WIDTH * MAX_HEIGHT * 3];
	/** y plane */
	unsigned char y_plane[MAX_WIDTH * MAX_WIDTH];
	/** u plane */
	unsigned char u_plane[MAX_WIDTH * MAX_WIDTH];
	/** v plane */
	unsigned char v_plane[MAX_WIDTH * MAX_WIDTH];

	/** actual width of frame */
	unsigned int width;
	/** actual height of frame */
	unsigned int height;

};

/** small contianer class */
struct yuv_pix {
	/** y component */
	unsigned char y;
	/** u component */
	unsigned char u;
	/** v component */
	unsigned char v;
};


static uint16_t
enc_read_dump_magic(FILE *fin)
{
	uint16_t magic;
	size_t sz;

	sz = fread((void*) &magic, sizeof(uint16_t), 1, fin);
	assert(sz == 1);

	return magic;
}

static int
enc_read_dump_event(FILE *fin, uint8_t *event)
{
	size_t sz;

	sz = fread((void*)event, sizeof(uint8_t), 1, fin);
	if (sz == 1) {
		return 1;
	}

	if (feof(fin)) {
		return 0;
	}

	fprintf(stderr, "Error reading dump: %s\n", strerror(errno));
	return -1;
}

static int
enc_read_dump_set(FILE *fin, int *x, int *y)
{
	struct rd_mon_set_data data;
	size_t sz;

	sz = fread((void*) &data, sizeof(data), 1, fin);
	if (sz == 1) {
		*x = data.x;
		*y = data.y;
		return 0;
	}

	if (feof(fin)) {
		/* shouldn't happen */
		return -1;
	}

	fprintf(stderr, "%s: Error reading dump %s\n", __func__, 
		strerror(errno));
	return -1;
}

static int
enc_read_dump_upd(FILE *fin, struct rd_mon_upd_data *data)
{
	size_t sz;

	sz = fread((void*)data, sizeof(struct rd_mon_upd_data), 1, fin);
	if (sz == 1) {
		return 0;
	}

	if (feof(fin)) {
		return -1;
	}

	fprintf(stderr, "Error reading dump %s\n", strerror(errno));
	return -1;
}

/** read a RD_MON_ALL_ID packet and store the pixel data in
 *  data.
 *  @param fin stream to read from
 *  @param data information that will get filled in
 *  @param rgb_buf pointer to rgb pixel buffer that will get filled
 *  @param skip if skip is set, the pixel buffer will not get filled,
 *         and the stream will only be positioned after the data.
 *  @param dw destination width of rgb buffer frame
 *  @return 0 on success, -1 on error
 */
static int
enc_read_dump_all(
	FILE *fin, 
	struct rd_mon_init_data *data,
	unsigned char *rgb_buf,
	bool skip,
	unsigned int dw
)
{
	size_t sz;
	unsigned int x;
	unsigned int y;
	int ret;
	int j;

	sz = fread((void *)data, sizeof(struct rd_mon_init_data), 1, fin);

	if (sz != 1) {
		if (ferror(fin)) {
			fprintf(stderr, "Error reading dump %s\n", 
				strerror(errno));
		}
		return -1;
	}

	if (skip) {
		ret = fseek(fin, data->max_height * data->max_width * 4,
				SEEK_CUR);
		if (ret == 0) {
			return 0;
		}

		fprintf(stderr, "Error seeking in dump %s\n", 
			strerror(errno));
		return -1;
	}

	j = 0;
	for (y = 0; y < data->max_height; y++) {
		for (x = 0; x < data->max_width; x++) {
			unsigned char r;
			unsigned char g;
			unsigned char b;
			unsigned char a;

			sz = fread((void*) &r, sizeof(r), 1, fin);
			if (sz != 1) {
				if (ferror(fin)) {
					fprintf(stderr, 
						"Error reading dump %s\n",
						strerror(errno));
				}
			return -1;
			}

			sz = fread((void*) &g, sizeof(g), 1, fin);
			if (sz != 1) {
				if (ferror(fin)) {
					fprintf(stderr, 
						"Error reading dump %s\n",
						strerror(errno));
				}
			return -1;
			}

			sz = fread((void*) &b, sizeof(b), 1, fin);
			if (sz != 1) {
				if (ferror(fin)) {
					fprintf(stderr, 
						"Error reading dump %s\n",
						strerror(errno));
				}
			return -1;
			}
			
			sz = fread((void*) &a, sizeof(a), 1, fin);
			if (sz != 1) {
				if (ferror(fin)) {
					fprintf(stderr, 
						"Error reading dump %s\n",
						strerror(errno));
				}
			return -1;
			}

			rgb_buf[(y * dw + x) * 3 + 0] = r;
			rgb_buf[(y * dw + x) * 3 + 1] = g;
			rgb_buf[(y * dw + x) * 3 + 2] = b;
		}
	}

	return 0;
}

/* +-----+
 * |12|12|
 * |34|34|
 * +-----+
 * |12|12|
 * |34|34|
 * +-----+
 */
/** downsample a frame 4:1.
 *  @param buffer buffer containing the raw yuv frame, will get 
 *         overwritten with result.
 *  @param w width of frame
 *  @param h height of frame
 */
static void
enc_downsample_41(char *buffer, unsigned int w, unsigned int h)
{
	unsigned int x;
	unsigned int y;

	assert(w % 2 == 0);
	assert(h % 2 == 0);

	for (y = 0; y < h / 2 - 1; y++) {
		for (x = 0; x < w / 2 - 1; x++) {
			unsigned int pix;
			unsigned int sx = x * 2;
			unsigned int sy = y * 2;

#if 0 /* doesn't seem to be correct */
			pix  = buffer[sy * w + sx]
			     + buffer[sy * w + sx + 1]
			     + buffer[(sy + 1) * w + sx]
			     + buffer[(sy + 1) * w + sx + 1];
			
			pix >>= 2;
#else
			pix = buffer[sy * w + sx + 0];
#endif

			buffer[y * w / 2 + x] = pix;
		}
	}
}

/** convert rgb to yuv
 *  @param r red component
 *  @param g green component
 *  @param b blue component
 *  @param ret structure will get filled with return values 
 */
static void
enc_rgb_to_yuv(
	unsigned char r, 
	unsigned char g, 
	unsigned char b,
	struct yuv_pix *ret
)
{
	ret->y = (0.257 * r) + (0.504 * g) + (0.098 * b) + 16;
	ret->v = (0.439 * r) - (0.368 * g) - (0.071 * b) + 128;
	ret->u = -(0.148 * r) - (0.291 * g) + (0.439 * b) + 128;
}

#if 0 /* useful for debugging only */
#include "glue-png.h"

static void
enc_test_dump(const unsigned char *buf, unsigned int w, unsigned int h)
{
	static uint32_t argb[MAX_WIDTH * MAX_HEIGHT];
	unsigned int x;
	unsigned int y;
	static int nr = 0;
	char name[PATH_MAX];
	int ret;

	for (y = 0; y < h; y++) {
		for (x = 0; x < w; x++) {
			unsigned char r;
			unsigned char g;
			unsigned char b;

			r = buf[(y * w + x) * 3 + 0]; 
			g = buf[(y * w + x) * 3 + 1];
			b = buf[(y * w + x) * 3 + 2];

			argb[y * w + x] = r
					| g << 8
					| b << 16;
		}
	}

	nr++;
	ret = snprintf(name, sizeof(name), "test-%03d.png", nr);
	assert(ret < sizeof(name));

	png_write(argb, w, h, 0, 0, w, h, name);
}
#endif /* debug code */

/** convert the frame raw_rgb_buf to yuv.
 *  @param raw_rgb_buf raw rgb buffer
 *  @param w frame width
 *  @param h frame height
 *  @param ret yuv plane (y, u, v must be set already)
 *  @param pf used pixelformat (to know if downsampling must take place)
 */
static void
enc_convert_frame(
	const unsigned char *raw_rgb_buf, 
	unsigned int w, 
	unsigned int h, 
	yuv_buffer *ret,
	theora_pixelformat pf
)
{
	unsigned int x;
	unsigned int y;
	struct yuv_pix yuv_pix;

	assert(w > 1);
	assert(h > 1);

	ret->y_width = w;
	ret->y_height = h;
	ret->y_stride = ret->y_width;

	switch (pf) {
	case OC_PF_420:
		ret->uv_width = w >> 1;
		ret->uv_height = h >> 1;
		break;

	case OC_PF_444:
		ret->uv_width = w;
		ret->uv_height = h;
		break;

	case OC_PF_422:
		assert(0);
		break;

	default:
		assert(0);
		break;
	}

	ret->uv_stride = ret->uv_width;


	assert(ret->y != NULL);
	assert(ret->u != NULL);
	assert(ret->v != NULL);

	for (y = 0; y < h; y++) {
		for (x = 0; x < w; x++) {
			unsigned char r;
			unsigned char g;
			unsigned char b;

			r = raw_rgb_buf[(y * w + x) * 3 + 0];
			g = raw_rgb_buf[(y * w + x) * 3 + 1];
			b = raw_rgb_buf[(y * w + x) * 3 + 2];

			enc_rgb_to_yuv(r, g, b, &yuv_pix);

			ret->y[y * w + x] = yuv_pix.y;
			ret->u[y * w + x] = yuv_pix.u;
			ret->v[y * w + x] = yuv_pix.v;
		}
	}

	switch (pf) {
	case OC_PF_420:
		enc_downsample_41(ret->u, w, h);
		enc_downsample_41(ret->v, w, h);
		break;

	case OC_PF_422:
		/* not yet supported */
		assert(0);
		break;

	case OC_PF_444:
		break;

	default:
		assert(0);
		break;

	}
}

/** create page(s) from the ogg packet op and write these to disks
 *  @param cpssp faum_encoder instance
 *  @param op ogg packet to pipe through libogg and write to disk
 *  @return -1 on error, 0 otherwise (even if no page was flushed)
 */
static int
enc_ogg_flush(struct faum_encoder *cpssp, ogg_packet *op)
{
	int ret;
	ogg_page og;

	ret = ogg_stream_packetin(&cpssp->vo, op);
	assert(ret == 0);

	while (1)  {
		size_t sz;

		ret = ogg_stream_pageout(&cpssp->vo, &og);
		if (ret == 0) {
			return 0;
		}

		sz = fwrite(og.header, 1, og.header_len, cpssp->ogg_out);
		if (sz != og.header_len) {
			return -1;
		}

		sz = fwrite(og.body, 1, og.body_len, cpssp->ogg_out);
		if (sz != og.body_len) {
			return -1;
		}
	}
}

/** flush all packets available from theora to ogg packets
 *  and write the resulting ogg pages to disk.
 *  @param cpssp faum_encoder instance
 *  @param last_packet is this the last packet?
 *  @return -1 on error, 0 otherwise.
 */
static int
enc_theora_flush(struct faum_encoder *cpssp, int last_packet)
{
	int ret;
	ogg_packet op;

	ret = theora_encode_packetout(&cpssp->th, last_packet, &op);
	if (ret == 0) {
		/* no packet available */
		return 0;
	}

	ret = enc_ogg_flush(cpssp, &op);
	assert(ret == 0);
	return 0;
}

static int
enc_determine_frame_size(struct faum_encoder *cpssp)
{
	uint16_t magic;

	magic = enc_read_dump_magic(cpssp->rec_in);
	if (magic != RD_MAGIC_ID) {
		fprintf(stderr, "Not a valid dump file\n");
		return -1;
	}

	while (1) {
		int ret;
		uint8_t event;
		struct rd_mon_upd_data upd_data;
		struct rd_mon_init_data init_data;
		int x = 0;
		int y = 0;

		ret = enc_read_dump_event(cpssp->rec_in, &event);
		if (ret == 0) {
			break;
		}

		if (ret < 0) {
			return -1;
		}

		switch (event) {
		case RD_MON_SET_ID:
			ret = enc_read_dump_set(cpssp->rec_in, &x, &y);
			if (cpssp->width < x) {
				cpssp->width = x;
			}
			if (cpssp->height < y) {
				cpssp->height = y;
			}
			break;

		case RD_MON_UPD_ID:
			ret = enc_read_dump_upd(cpssp->rec_in, &upd_data);
			break;

		case RD_MON_ALL_ID:
			ret = enc_read_dump_all(cpssp->rec_in, 
						&init_data, 
						cpssp->rgb_buf,
						true,
						cpssp->width);
			assert(init_data.max_height <= MAX_HEIGHT);
			assert(init_data.max_width <= MAX_WIDTH);
			break;

		default:
			fprintf(stderr, "Illegal format code in dump file\n");
			return -1;
		}

		if (ret < 0) {
			return -1;
		}
	}

	/* realign to 16 pixel boundary */
	cpssp->width = (cpssp->width + 0xf) & ~0xf;
	cpssp->height = (cpssp->height + 0xf) & ~0xf;

	assert(cpssp->width < MAX_WIDTH);
	assert(cpssp->height < MAX_HEIGHT);

	fprintf(stderr, "frame size: %d x %d pixels\n", cpssp->width, 
		cpssp->height);

	return 0;
}

static int
enc_encode(struct faum_encoder *cpssp)
{
	uint16_t magic;
	int ret;
	uint8_t event;
	struct rd_mon_upd_data upd_data;
	struct rd_mon_init_data init_data;
	int x = 0;
	int y = 0;
	yuv_buffer buf;
	int idx;
	unsigned long long next_time = 0;
	unsigned long long time_hz = 0;

	buf.y = cpssp->y_plane;
	buf.u = cpssp->u_plane;
	buf.v = cpssp->v_plane;

	magic = enc_read_dump_magic(cpssp->rec_in);
	if (magic != RD_MAGIC_ID) {
		fprintf(stderr, "Not a valid dump file\n");
		return -1;
	}

	while (1) {
		ret = enc_read_dump_event(cpssp->rec_in, &event);
		if (ret == 0) {
			break;
		}

		if (ret < 0) {
			return -1;
		}

		switch (event) {
		case RD_MON_SET_ID:
			ret = enc_read_dump_set(cpssp->rec_in, &x, &y);
			break;

		case RD_MON_UPD_ID:
			ret = enc_read_dump_upd(cpssp->rec_in, &upd_data);

			if (next_time < upd_data.t) {
				enc_convert_frame(cpssp->rgb_buf,
						cpssp->width,
						cpssp->height,
						&buf,
						cpssp->ti.pixelformat);
			}

			/* keep encoding previously stored frames until the 
			 * just read time is reached */
			while (next_time < upd_data.t) {
				ret = theora_encode_YUVin(&cpssp->th, &buf);
				if (ret != 0) {
					return -1;
				}
				ret = enc_theora_flush(cpssp, 0);
				if (ret < 0) {
					return -1;
				}

				next_time += (time_hz / FPS_NUM) * FPS_DENOM;
			}
			/* then put the submitted pixel into our pixel buffer 
			 */
			idx = (upd_data.y * cpssp->width + upd_data.x) * 3;
			cpssp->rgb_buf[idx + 0] = upd_data.r;
			cpssp->rgb_buf[idx + 1] = upd_data.g;
			cpssp->rgb_buf[idx + 2] = upd_data.b;

			break;

		case RD_MON_ALL_ID:
			ret = enc_read_dump_all(cpssp->rec_in, 
						&init_data, 
						cpssp->rgb_buf,
						false,
						cpssp->width);
			if (ret < 0) {
				return -1;
			}

			enc_convert_frame(cpssp->rgb_buf, 
						cpssp->width, 
						cpssp->height, 
						&buf,
						cpssp->ti.pixelformat);

			ret = theora_encode_YUVin(&cpssp->th, &buf);
			if (ret != 0) {
				fprintf(stderr, "YUVin fail %d\n", ret);
				return -1;
			}
			ret = enc_theora_flush(cpssp, 0);
			time_hz = init_data.time_hz;
			next_time = init_data.t 
			          + (time_hz / FPS_NUM) * FPS_DENOM;
			break;

		default:
			fprintf(stderr, "Illegal format code in dump file\n");
			return -1;
		}

		if (ret < 0) {
			return -1;
		}
	}

	/* flush last buffer contents */
	enc_convert_frame(cpssp->rgb_buf, cpssp->width, cpssp->height, &buf,
			cpssp->ti.pixelformat);
	ret = theora_encode_YUVin(&cpssp->th, &buf);
	if (ret != 0) {
		fprintf(stderr, "YUVin fail %d\n", ret);
		return -1;
	}
	ret = enc_theora_flush(cpssp, 1);

	return ret;
}


static int
enc_theora_create(struct faum_encoder *cpssp)
{
	int ret;

	memset(&cpssp->ti, 0, sizeof(cpssp->ti));

	cpssp->ti.width = cpssp->width;
	cpssp->ti.height = cpssp->height;
	cpssp->ti.frame_width = cpssp->width;
	cpssp->ti.frame_height = cpssp->height;
	cpssp->ti.offset_x = 0;
	cpssp->ti.offset_y = 0;
	/* 12 / 1 FPS */
	cpssp->ti.fps_numerator = FPS_NUM;
	cpssp->ti.fps_denominator = FPS_DENOM;

	cpssp->ti.aspect_numerator = 1;
	cpssp->ti.aspect_denominator = 1;

	cpssp->ti.colorspace = OC_CS_UNSPECIFIED;
	cpssp->ti.quality = 40;
	cpssp->ti.quick_p = 0;
	
	cpssp->ti.keyframe_auto_p = 0;
	cpssp->ti.keyframe_frequency = 64;
	cpssp->ti.keyframe_frequency_force = 128;

	/* First try 4:4:4, then 4:2:1 */
	cpssp->ti.pixelformat = OC_PF_444;

	ret = theora_encode_init(&cpssp->th, &cpssp->ti);
	if (ret != 0) {
		fprintf(stderr, "Pixelformat 4:4:4 is not supported, "
			"trying 4:2:0 instead\n");
		cpssp->ti.pixelformat = OC_PF_420;
		ret = theora_encode_init(&cpssp->th, &cpssp->ti);
	}
	if (ret != 0) {
		fprintf(stderr, "Init failed: %d\n", ret);
		return -1;
	}
	return 0;
}

static int
enc_ogg_stream_create(struct faum_encoder *cpssp)
{
	time_t ts = time(NULL);
	int r;
	int ret;

	srand(ts);
	r = rand();

	memset(&cpssp->vo, 0, sizeof(cpssp->vo));
	ret = ogg_stream_init(&cpssp->vo, r);
	return ret;
}

static int
enc_create(struct faum_encoder *cpssp)
{
	int ret;
	ogg_packet op;
	theora_comment tc;

	ret = enc_theora_create(cpssp);
	assert(ret == 0);
	ret = enc_ogg_stream_create(cpssp);
	assert(ret == 0);

	ret = theora_encode_header(&cpssp->th, &op);
	if (ret != 0) {
		fprintf(stderr, "Header fail: %d\n", ret);
		return -1;
	}

	ret = enc_ogg_flush(cpssp, &op);
	assert(ret == 0);

	theora_comment_init(&tc);
	ret = theora_encode_comment(&tc, &op);
	assert(ret == 0);
	theora_comment_clear(&tc);
	ret = ogg_stream_packetin(&cpssp->vo, &op);
	assert(ret == 0);

	ret = theora_encode_tables(&cpssp->th, &op);
	assert(ret == 0);
	ret = ogg_stream_packetin(&cpssp->vo, &op);
	assert(ret == 0);

	return ret;
}

static int
enc_destroy(struct faum_encoder *cpssp)
{
	int ret;
	int r = EXIT_SUCCESS;

	theora_clear(&cpssp->th);
	ret = ogg_stream_clear(&cpssp->vo);
	assert(ret == 0);

	ret = fclose(cpssp->ogg_out);
	if (ret != 0) {
		fprintf(stderr, "Error closing stream: %s\n",
			strerror(errno));
		r = EXIT_FAILURE;
	}

	ret = fclose(cpssp->rec_in);
	if (ret != 0) {
		fprintf(stderr, "Error closing stream: %s\n",
			strerror(errno));
		r = EXIT_FAILURE;
	}
	return r;
}

static void
usage(const char *prog_name)
{
	fprintf(stderr, "usage: %s <dumpfile>\n", prog_name);
}

int
main(int argc, char **argv)
{
	int ret;
	/* faum_encoder is too big for the stack */
	static struct faum_encoder cpssp;
	char out_name[PATH_MAX];

	if (argc != 2) {
		usage(argv[0]);
		return EXIT_FAILURE;
	}

	ret = snprintf(out_name, sizeof(out_name), "%s.ogv", argv[1]);
	assert(ret < sizeof(out_name));

	cpssp.rec_in = fopen(argv[1], "r");
	if (cpssp.rec_in == NULL) {
		fprintf(stderr, "Could not open %s for reading: %s\n",
			argv[1], strerror(errno));
		return EXIT_FAILURE;
	}

	ret = enc_determine_frame_size(&cpssp);
	if (ret < 0) {
		return EXIT_FAILURE;
	}

	cpssp.ogg_out = fopen(out_name, "w");
	if (cpssp.ogg_out == NULL) {
		fprintf(stderr, "Could not open %s for writing: %s\n",
			out_name, strerror(errno));
		fclose(cpssp.rec_in);
		return EXIT_FAILURE;
	}

	ret = enc_create(&cpssp);
	if (ret < 0) {
		return EXIT_FAILURE;
	}

	rewind(cpssp.rec_in);

	ret = enc_encode(&cpssp);
	if (ret < 0) {
		enc_destroy(&cpssp);
		return EXIT_FAILURE;
	}

	ret = enc_destroy(&cpssp);
	return ret;
}