/*
 * VPXEncoder.h
 *
 *  Created on: 28/12/2013
 *      Author: lucas
 *      Based on: http://www.webmproject.org/docs/vp8-sdk/example__simple__encoder.html
 *      An part of StepFever project.
 */


#include "VPXEncoder.h"

VPXEncoder::VPXEncoder()  :in_width(0)
,in_height(0)
,out_width(0)
,out_height(0)
,in_pixel_format(AV_PIX_FMT_RGB24)
,out_pixel_format(AV_PIX_FMT_YUV420P)
,sws(NULL)
{}

void VPXEncoder::write_ivf_file_header(FILE *outfile, const vpx_codec_enc_cfg_t *cfg, int frame_cnt) {
    char header[32];

    if(cfg->g_pass != VPX_RC_ONE_PASS && cfg->g_pass != VPX_RC_LAST_PASS)
        return;
    header[0] = 'D';
    header[1] = 'K';
    header[2] = 'I';
    header[3] = 'F';
    mem_put_le16(header+4,  0);                   /* version */
    mem_put_le16(header+6,  32);                  /* headersize */
    mem_put_le32(header+8,  fourcc);              /* headersize */
    mem_put_le16(header+12, cfg->g_w);            /* width */
    mem_put_le16(header+14, cfg->g_h);            /* height */
    mem_put_le32(header+16, cfg->g_timebase.den); /* rate */
    mem_put_le32(header+20, cfg->g_timebase.num); /* scale */
    mem_put_le32(header+24, frame_cnt);           /* length */
    mem_put_le32(header+28, 0);                   /* unused */

    (void) fwrite(header, 1, 32, outfile);
}

void VPXEncoder::write_ivf_frame_header(FILE *outfile, const vpx_codec_cx_pkt_t *pkt)
{
    char             header[12];
    vpx_codec_pts_t  pts;

    if(pkt->kind != VPX_CODEC_CX_FRAME_PKT)
        return;

    pts = pkt->data.frame.pts;
    mem_put_le32(header, pkt->data.frame.sz);
    mem_put_le32(header+4, pts&0xFFFFFFFF);
    mem_put_le32(header+8, pts >> 32);

    (void) fwrite(header, 1, 12, outfile);
}
void VPXEncoder::die_codec(vpx_codec_ctx_t *ctx, const char *s) {
    const char *detail = vpx_codec_error_detail(ctx);
    cout << s << ": " << vpx_codec_error(ctx) << endl;
    if(detail)
        cout << "    " << detail << endl;
    exit(EXIT_FAILURE);
}

VPXEncoder::~VPXEncoder() {
	if(sws)
		close();
}

bool VPXEncoder::open(std::string filename)	{

    if(!vpx_img_alloc(&raw, VPX_IMG_FMT_I420, out_width, out_height, 1))	{
        cout << "Failed to allocate image " << out_width << " " << out_height << endl;
        return false;
    }

	if(out_pixel_format != AV_PIX_FMT_YUV420P) {
		cout << ("At this moment the output format must be AV_PIX_FMT_YUV420P") << endl;
		return false;
	}
	sws = sws_getContext(in_width, in_height, in_pixel_format,
				   out_width, out_height, out_pixel_format,
				   SWS_FAST_BILINEAR, NULL, NULL, NULL);

	if(!sws) {
		cout << "Cannot create SWS context" << endl;
		::exit(EXIT_FAILURE);
	}

	outfile = fopen(filename.c_str(), "w+b");
	if(!outfile) {
	  cout << "Cannot open the VP8 destination file" << endl;
	  return false;
	}

	res = vpx_codec_enc_config_default(interface, &cfg, 0);
	if(res) {
		cout << "Failed to get config: " << vpx_codec_err_to_string(res) << endl;
		return false;
	}

	cfg.rc_target_bitrate = out_width * out_height * cfg.rc_target_bitrate;
	cfg.g_w = out_width;
	cfg.g_h = out_height;

	write_ivf_file_header(outfile, &cfg, 0);
	if(vpx_codec_enc_init(&codec, interface, &cfg, 0))
		die_codec(&codec, "Failed to initialize encoder");

	frame_avail = 1;
	got_data = 0;
}

bool VPXEncoder::encode(char* pixels, bool ReflectY)	{
	vpx_codec_iter_t iter = NULL;
	const vpx_codec_cx_pkt_t *pkt;

	if(!sws) {
		cout << ("Not initialized, so cannot encode") << endl;
		return false;
	}

	// copy the pixels into our "raw input" container.
	int bytes_filled = avpicture_fill(&pic_raw, (uint8_t*)pixels, in_pixel_format, in_width, in_height);
	if(!bytes_filled) {
		cout << ("Cannot fill the raw input buffer") << endl;
		return false;
	}

	// Reflect for OpenGL
	if(ReflectY)	{
		for (int i = 0; i < 4; i++) {
		  pic_raw.data[i] += pic_raw.linesize[i] * (in_height-1);
		  pic_raw.linesize[i] = -pic_raw.linesize[i];
		}
	}
	// convert to I420 for vp8
	int h = sws_scale(sws, pic_raw.data, pic_raw.linesize, 0, in_height, raw.planes, raw.stride);

	if(h != out_height) {
		cout << "scale failed: " << h << endl;
		return false;
	}
	if(vpx_codec_encode(&codec, &raw, frame_cnt, 1, flags, VPX_DL_REALTIME))
		  die_codec(&codec, "Failed to encode frame");

	while( (pkt = vpx_codec_get_cx_data(&codec, &iter)) ) {
		switch(pkt->kind) {
			case VPX_CODEC_CX_FRAME_PKT:
				write_ivf_frame_header(outfile, pkt);
				(void) fwrite(pkt->data.frame.buf, 1, pkt->data.frame.sz, outfile);
				break;
			default:break;
		}
	}
	frame_cnt++;
	return true;
}

bool VPXEncoder::close() {

	if(outfile)	{
		if(!fseek(outfile, 0, SEEK_SET))
			write_ivf_file_header(outfile, &cfg, frame_cnt-1);

		fclose(outfile);
		outfile = NULL;
	}
	return true;
}