#include"node.h"

#include<vpx/vpx_encoder.h>
#include<vpx/vp8cx.h>

#include<assert.h>

#include"mode.h"

#include"img.h"
#include<math.h>

#include<smmintrin.h>

#include<string.h>

#include"minitrace.h"

#include"linearity.h"

#include"yuv.h"

typedef struct CHiEncodeVP9Node {
	CHiPubNode pub;
	
	vpx_codec_ctx_t codec;
	vpx_codec_enc_cfg_t cfg;
	
	enum {
		WAITING, IN_PROGRESS
	} state;
	uint8_t *outY, *outU, *outV;
	uint16_t strideY, strideU, strideV;
	
	vpx_codec_iface_t *iface;
} CHiEncodeVP9Node;

static int encodevpx_perform(CHiPubNode *pub) {
	CHiEncodeVP9Node *node = (void*) pub;
	
	MTR_BEGIN("CHi", "encodevp9_perform");
	
	pub->sources[0].type = CUTIHI_VAL_VP9BS;
	pub->sources[0].data.bitstream = NULL;
	
	if(node->state == WAITING) return 1;
	
	CHiImage *rgbIn = (CHiImage*) CHi_Crawl(&pub->sinks[0])->data.sample;
	
	bgra64toycbcr(rgbIn->data8, rgbIn->stride, rgbIn->width, rgbIn->height, node->outY, node->outU, node->outV, node->strideY, node->strideU, node->strideV);
	
	vpx_image_t vpxraw;
	vpxraw.fmt = VPX_IMG_FMT_I420;
	vpxraw.cs = VPX_CS_BT_709;
	vpxraw.range = VPX_CR_STUDIO_RANGE;
	vpxraw.bit_depth = 8;
	vpxraw.w = vpxraw.d_w = node->cfg.g_w;
	vpxraw.h = vpxraw.d_h = node->cfg.g_h;
	vpxraw.r_w = vpxraw.r_h = 0;
	vpxraw.x_chroma_shift = vpxraw.y_chroma_shift = 1;
	vpxraw.img_data_owner = 0;
	vpxraw.self_allocd = 0;
	vpxraw.bps = 12;
	vpxraw.stride[VPX_PLANE_Y] = node->strideY;
	vpxraw.planes[VPX_PLANE_Y] = node->outY;
	vpxraw.stride[VPX_PLANE_U] = node->strideU;
	vpxraw.planes[VPX_PLANE_U] = node->outU;
	vpxraw.stride[VPX_PLANE_V] = node->strideV;
	vpxraw.planes[VPX_PLANE_V] = node->outV;
	
	vpx_codec_encode(&node->codec, &vpxraw, CHi_Time_Get(pub->ng) * 1000.f, 1, 0, VPX_DL_REALTIME);
	
	CHiBSFrames *ret = malloc(sizeof(CHiBSFrames));
	ret->count = 0;
	
	vpx_codec_iter_t iter = NULL;
	const vpx_codec_cx_pkt_t *pkt;
	while((pkt = vpx_codec_get_cx_data(&node->codec, &iter)) != NULL) {
		if(pkt->kind == VPX_CODEC_CX_FRAME_PKT) {
			ret = (CHiBSFrames*) realloc(ret, sizeof(CHiBSFrames) + sizeof(CHiBSFrame) * (ret->count + 1));
			ret->data[ret->count].timestamp = pkt->data.frame.pts;
			ret->data[ret->count].sz = pkt->data.frame.sz;
			ret->data[ret->count].flags = pkt->data.frame.flags & VPX_FRAME_IS_KEY;
			ret->data[ret->count].ptr = malloc(ret->data[ret->count].sz);
			memcpy(ret->data[ret->count].ptr, pkt->data.frame.buf, ret->data[ret->count].sz);
			ret->count++;
		}
	}
	
	pub->sources[0].data.bitstream = ret;
	
	MTR_END("CHi", "encodevp9_perform");
	
	return 1;
}

static void encodevpx_destroy(CHiPubNode *pub) {
	CHiEncodeVP9Node *node = (void*) pub;
	
	free(node);
}

static int encodevpx_start(CHiPubNode *pubn) {
	CHiEncodeVP9Node *node = (void*) pubn;
	
	node->state = IN_PROGRESS;
	
	CHiImage *firstFrame = (CHiImage*) CHi_Crawl(&pubn->sinks[0])->data.sample;
	
	vpx_codec_enc_config_default(node->iface, &node->cfg, 0);
	node->cfg.g_w = firstFrame->width;
	node->cfg.g_h = firstFrame->height;
	node->cfg.g_timebase.num = 1;
	node->cfg.g_timebase.den = 30;
	node->cfg.g_lag_in_frames = 0;
	node->cfg.g_threads = 8;
	node->cfg.kf_mode = VPX_KF_AUTO;
	node->cfg.kf_max_dist = 300;
	node->cfg.rc_end_usage = VPX_VBR;
	node->cfg.rc_target_bitrate = 512;
	node->cfg.rc_min_quantizer = 4;
	node->cfg.rc_max_quantizer = 48;
	
	vpx_codec_enc_init(&node->codec, node->iface, &node->cfg, 0);
	vpx_codec_control(&node->codec, VP8E_SET_CPUUSED, 8);
	vpx_codec_control(&node->codec, VP9E_SET_ROW_MT, 1);
	vpx_codec_control(&node->codec, VP9E_SET_TILE_COLUMNS, 2);
	vpx_codec_control(&node->codec, VP9E_SET_TILE_ROWS, 1);
	vpx_codec_control(&node->codec, VP9E_SET_TUNE_CONTENT, VP9E_CONTENT_SCREEN);
	
	node->strideY = (node->cfg.g_w + 64) & ~63;
	node->strideU = (node->cfg.g_w / 2 + 64) & ~63;
	node->strideV = (node->cfg.g_w / 2 + 64) & ~63;
	
	node->outY = (uint8_t*) _mm_malloc(node->strideY * node->cfg.g_h, 16);
	node->outU = (uint8_t*) _mm_malloc(node->strideU * node->cfg.g_h / 2, 16);
	node->outV = (uint8_t*) _mm_malloc(node->strideV * node->cfg.g_h / 2, 16);
	
	return 1;
}

static int encodevpx_stop(CHiPubNode *pubn) {
	CHiEncodeVP9Node *node = (void*) pubn;
	
	node->state = WAITING;
	
	_mm_free(node->outY);
	_mm_free(node->outU);
	_mm_free(node->outV);
	
	vpx_codec_destroy(&node->codec);
	
	return 1;
}

CUTIVIS CHiPubNode *CHi_EncodeVP8() {
	CHiEncodeVP9Node *n = calloc(1, sizeof(*n));
	n->pub.type = CUTIHI_T('CEnc','GVP8');
	n->pub.Start = encodevpx_start;
	n->pub.Perform = encodevpx_perform;
	n->pub.Stop = encodevpx_stop;
	n->pub.Destroy = encodevpx_destroy;
	n->pub.sinks = (CHiValue*) calloc(sizeof(*n->pub.sinks), n->pub.sinkCount = 1);
	n->pub.sources = (CHiValue*) calloc(sizeof(*n->pub.sources), n->pub.sourceCount = 1);
	n->state = WAITING;
	n->iface = vpx_codec_vp8_cx();
	return &n->pub;
}

CUTIVIS CHiPubNode *CHi_EncodeVP9() {
	CHiEncodeVP9Node *n = calloc(1, sizeof(*n));
	n->pub.type = CUTIHI_T('CEnc','GVP9');
	n->pub.Start = encodevpx_start;
	n->pub.Perform = encodevpx_perform;
	n->pub.Stop = encodevpx_stop;
	n->pub.Destroy = encodevpx_destroy;
	n->pub.sinks = (CHiValue*) calloc(sizeof(*n->pub.sinks), n->pub.sinkCount = 1);
	n->pub.sources = (CHiValue*) calloc(sizeof(*n->pub.sources), n->pub.sourceCount = 1);
	n->state = WAITING;
	n->iface = vpx_codec_vp9_cx();
	return &n->pub;
}