#include"node.h"

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

#include<new>

#include<mkvmuxer/mkvwriter.h>

#include<assert.h>

#include"mode.h"

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

#include<smmintrin.h>

#include<queue>

#include<string.h>

#include"minitrace.h"

#include"linearity.h"

#include"yuv.h"

struct CHiMuxWebmNode {
	CHiPubNode pub;
	
	mkvmuxer::MkvWriter w;
	mkvmuxer::Segment seg;
	size_t videoTrack, audioTrack;
	
	std::queue<CHiBSFrame> audioBacklog;
	std::queue<CHiBSFrame> videoBacklog;
};
static int muxwebm_perform(CHiPubNode *pubn) {
	using namespace mkvmuxer;
	
	MTR_BEGIN("CHi", "muxwebm_perform");
	
	CHiMuxWebmNode *alln = (CHiMuxWebmNode*) pubn;
	
	if(pubn->sinks[1].data.linked.to) {
		CHiBSFrames *opus = CHi_Crawl(&pubn->sinks[1])->data.bitstream;
		for(size_t i = 0; i < opus->count; i++) {
			alln->audioBacklog.push(opus->data[i]);
		}
	}
	
	auto vp9bs = CHi_Crawl(&pubn->sinks[0])->data.bitstream;
	if(vp9bs) {
		for(size_t i = 0; i < vp9bs->count; i++) {
			alln->videoBacklog.push(vp9bs->data[i]);
		}
	}
	
	while(pubn->sinks[1].data.linked.to && alln->audioBacklog.size() > 0 && alln->videoBacklog.size() > 0 && alln->audioBacklog.front().timestamp <= alln->videoBacklog.front().timestamp) {
		Frame frame;
		frame.Init((const uint8_t*) alln->audioBacklog.front().ptr, alln->audioBacklog.front().sz);
		frame.set_track_number(alln->audioTrack);
		frame.set_timestamp(alln->audioBacklog.front().timestamp * 1000000L);
		frame.set_is_key(true);
		alln->seg.AddGenericFrame(&frame);
		
		alln->audioBacklog.pop();
	}
	
	if(pubn->sinks[1].data.linked.to == NULL || (alln->audioBacklog.size() > 0 && alln->videoBacklog.size() > 0 && alln->audioBacklog.front().timestamp >= alln->videoBacklog.front().timestamp)) {
		Frame frame;
		if(!frame.Init((const uint8_t*) alln->videoBacklog.front().ptr, alln->videoBacklog.front().sz)) puts("INIT FAIL");
		frame.set_track_number(alln->videoTrack);
		frame.set_timestamp(alln->videoBacklog.front().timestamp * 1000000L);
		frame.set_is_key(!!(alln->videoBacklog.front().flags & CUTIHI_BS_FLAG_KEY));
		if(!alln->seg.AddGenericFrame(&frame)) puts("ADD FAIL");
		
		alln->videoBacklog.pop();
	}
	
	MTR_END("CHi", "muxwebm_perform");
	
	return 1;
}
static void muxwebm_destroy(CHiPubNode *pubn) {
	CHiMuxWebmNode *n = (CHiMuxWebmNode*) pubn;
	
	n->~CHiMuxWebmNode();
	
	free(n);
}
CUTIVIS CHiPubNode *CHi_MuxWebm() {
	CHiMuxWebmNode *n = (CHiMuxWebmNode*) calloc(1, sizeof(*n));
	new (n) CHiMuxWebmNode;
	n->pub.type = CUTIHI_T('CExp','Webm');
	n->pub.Start = CHi_MuxWebm_Start;
	n->pub.Perform = muxwebm_perform;
	n->pub.Destroy = muxwebm_destroy;
	n->pub.Stop = CHi_MuxWebm_Stop;
	n->pub.sinks = (CHiValue*) calloc(sizeof(*n->pub.sinks), n->pub.sinkCount = 3);
	n->pub.sourceCount = 0;
	n->pub.sources = NULL;
	return &n->pub;
}
CUTIVIS int CHi_MuxWebm_Start(CHiPubNode *pubn) {
	using namespace mkvmuxer;
	
	CHiMuxWebmNode *alln = (CHiMuxWebmNode*) pubn;
	
	new (&alln->w) MkvWriter{};
	alln->w.Open(CHi_Crawl(&pubn->sinks[CUTIHI_MUXWEBM_IN_FILENAME])->data.text);
	
	new (&alln->seg) Segment{};
	alln->seg.Init(&alln->w);
	alln->seg.AccurateClusterDuration(true);
	alln->seg.UseFixedSizeClusterTimecode(false);
	alln->seg.set_mode(Segment::kFile);
	alln->seg.OutputCues(true);
	alln->seg.set_duration(pubn->ng->duration * 1000);
	alln->seg.GetSegmentInfo()->set_timecode_scale(1000000);
	alln->seg.GetSegmentInfo()->set_writing_app("Cuticle");
	
	/* Hack into first frame to get resolution */
	CHiPubNode *evp9 = pubn->sinks[0].data.linked.to;
	CHiImage *firstFrame = (CHiImage*) CHi_Crawl(&evp9->sinks[0])->data.sample;
	
	alln->videoTrack = alln->seg.AddVideoTrack(firstFrame->width, firstFrame->height, 0);
	VideoTrack *track = (VideoTrack*) alln->seg.GetTrackByNumber(alln->videoTrack);
	track->set_codec_id(CHi_Crawl(&pubn->sinks[0])->type == CUTIHI_VAL_VP9BS ? Tracks::kVp9CodecId : Tracks::kVp8CodecId);
	track->set_frame_rate(30);
	Colour colourspace;
	colourspace.set_matrix_coefficients(Colour::MatrixCoefficients::kBt709);
	colourspace.set_range(Colour::Range::kBroadcastRange);
	colourspace.set_transfer_characteristics(Colour::TransferCharacteristics::kIturBt709Tc);
	colourspace.set_primaries(Colour::Primaries::kIturBt709P);
	track->SetColour(colourspace);
	alln->seg.CuesTrack(alln->videoTrack);
	
	if(pubn->sinks[1].data.linked.to) {
		struct __attribute__((packed)) {
			uint32_t magic1;
			uint32_t magic2;
			uint8_t version;
			uint8_t channels; // NUMBER OF CHANNELS IS HARDCODED TO ONE
			uint16_t preskip;
			uint32_t sampleRate;
			uint16_t gain;
			uint8_t family;
		} header = {0x7375704f, 0x64616548, 1, 1, 3840, 48000, 0, 0};
		
		alln->audioTrack = alln->seg.AddAudioTrack(48000, 1, 0);
		AudioTrack *atrack = (AudioTrack*) alln->seg.GetTrackByNumber(alln->audioTrack);
		atrack->set_codec_id(Tracks::kOpusCodecId);
		atrack->set_seek_pre_roll(80000000);
		atrack->SetCodecPrivate((const uint8_t*) &header, sizeof(header));
	}
	
	return 1;
}
CUTIVIS int CHi_MuxWebm_Stop(CHiPubNode *pubn) {
	CHiMuxWebmNode *alln =(CHiMuxWebmNode*) pubn;
	alln->seg.Finalize();
	
	alln->w.Close();
	
	return 1;
}