#include "tracing.h"

#include "rtp/RTPStreamTransponder.h"
#include "waitqueue.h"
#include "vp8/vp8.h"
#include "av1/AV1LayerSelector.h"


RTPStreamTransponder::RTPStreamTransponder(const RTPOutgoingSourceGroup::shared& outgoing, const RTPSender::shared& sender) :
	TimeServiceWrapper<RTPStreamTransponder>(outgoing->GetTimeService()),
	outgoing(outgoing),
	sender(sender)
{
	//Store outgoing streams
	ssrc = outgoing->media.ssrc;

	//Add us as listeners
	outgoing->AddListener(this);

	Debug("-RTPStreamTransponder() | [outgoing:%p,sender:%p,ssrc:%u]\n",outgoing,sender,ssrc);
}

void RTPStreamTransponder::ResetIncoming()
{
	SetIncoming(nullptr, nullptr);
}

void RTPStreamTransponder::SetIncoming(const RTPIncomingMediaStream::shared& incoming, const RTPReceiver::shared& receiver, bool smooth)
{
	AsyncSafe([=](auto now){
		//Check we are not closed
		if (!outgoing)
		{
			//Error
			Error("-RTPStreamTransponder::SetIncoming() | Transponder already closed\n");
			//Exit
			return;
		}

		Debug(">RTPStreamTransponder::SetIncoming() | [incoming:%p,receiver:%p,ssrc:%u,smooth:%d]\n", incoming, receiver, ssrc, smooth);

		if (smooth)
		{
			//If they are the same as next ones already
			if (this->incomingNext == incoming && this->receiverNext == receiver)
				//DO nothing
				return;

			//Remove listener from previues transitioning stream
			if (this->incomingNext)
				this->incomingNext->RemoveListener(this);

			//If they are the same as current ones
			if (this->incoming == incoming && this->receiver == receiver)
			{
				//And don't wait anymore
				incomingNext = nullptr;
				receiverNext = nullptr;
				//DO nothing
				return;
			}

			//Store stream and receiver
			this->incomingNext = incoming;
			this->receiverNext = receiver;

			//Double check
			if (this->incomingNext)
			{
				//Add us as listeners
				this->incomingNext->AddListener(this);

				//Request update on the transitoning one
				if (this->receiverNext) this->receiverNext->SendPLI(this->incomingNext->GetMediaSSRC());
			}

		} else {
			//If we were waiting to transition to a new incoming media stream
			if (this->incomingNext)
			{
				//Stop listening from it
				this->incomingNext->RemoveListener(this);
				//And don't wait anymore
				incomingNext = nullptr;
				receiverNext = nullptr;
			}

			//If they are the same as current ones
			if (this->incoming==incoming && this->receiver==receiver)
				//DO nothing
				return;

			//Remove listener from old stream
			if (this->incoming)
				this->incoming->RemoveListener(this);

			//Reset packets before start listening again
			reset = true;

			//Store stream and receiver
			this->incoming = incoming;
			this->receiver = receiver;

			//Double check
			if (this->incoming)
			{
				//Add us as listeners
				this->incoming->AddListener(this);

				//Request update on the incoming
				if (this->receiver) this->receiver->SendPLI(this->incoming->GetMediaSSRC());
				//Update last requested PLI
				lastSentPLI = now.count();
			}
		}

		Debug("<RTPStreamTransponder::SetIncoming() | [incoming:%p,receiver:%p]\n",incoming,receiver);
	});
}

RTPStreamTransponder::~RTPStreamTransponder()
{
	Debug("~RTPStreamTransponder::~RTPStreamTransponder() [this:%p]\n", this);

	//If not already stopped
	if (outgoing || incoming)
		//Stop listeneing
		Close();
}

void RTPStreamTransponder::Close()
{
	Debug(">RTPStreamTransponder::Close() [this:%p]\n",this);

	//Stop listening
	if (outgoing) outgoing->RemoveListener(this);

	Sync([=](auto now) {
		//Stop listening
		if (incoming) incoming->RemoveListener(this);
		if (incomingNext) incomingNext->RemoveListener(this);
		//Remove sources
		incoming = nullptr;
		receiver = nullptr;
		incomingNext = nullptr;

		receiverNext = nullptr;
		//Remove sources
		outgoing = nullptr;
		sender = nullptr;
	});

	Debug("<RTPStreamTransponder::Close() [this:%p]\n",this);
}


void RTPStreamTransponder::onRTP(const RTPIncomingMediaStream* stream,const RTPPacket::shared& packet)
{
	//Trace method
	TRACE_EVENT("rtp","RTPStreamTransponder::onRTP", "ssrc", packet->GetSSRC(), "seqnum", packet->GetSeqNum());

	if (!packet)
		//Exit
		return;

	//Run async
	AsyncSafe([=, packet = packet->Clone()](auto now){

		// In theory this check is not necessary as it is already checked in onRTPAsync but 
		// more encapsulated and safer against future maintainence duplicating it here.
		if (stream != incomingNext.get() && stream != incoming.get())
		{
			return;
		}

		onRTPAsync(now, stream, packet);
	});
}

void RTPStreamTransponder::onRTPAsync(std::chrono::milliseconds now, const RTPIncomingMediaStream* stream, const RTPPacket::shared& packet)
{
	//If it is from the next transitioning stream
	if (stream == incomingNext.get())
	{
		//If it is a video packet and not an iframe
		if (packet->GetMediaType()==MediaFrame::Video && !packet->IsKeyFrame())
			//Skip
			return;

		//Remove listener from old stream
		if (this->incoming)
			this->incoming->RemoveListener(this);

		//Reset packets
		reset = true;

		//Transition to new stream and receiver
		this->incoming = incomingNext;
		this->receiver = receiverNext;
		this->incomingNext = nullptr;
		this->receiverNext = nullptr;
	}

	//Check if it is from the correct stream
	if (stream != this->incoming.get())
		//Skip
		return;

	//If muted
	if (muted)
		//Skip
		return;

	//If forwarding only intra frames and video frame is not intra
	if (intraOnlyForwarding && packet->GetMediaType() == MediaFrame::Video && !packet->IsKeyFrame())
	{
		//Drop it
		dropped++;
		//Skip
		return;
	}

	//Check if it is an empty packet
	if (!packet->GetMediaLength())
	{
		UltraDebug("-RTPStreamTransponder::onRTP() | dropping empty packet\n");
		//Drop it
		dropped++;
		//Exit
		return;
	}

	//Check sender
	if (!sender)
		//Nothing
		return;

	//Check if source has changed
	if (source && packet->GetSSRC()!=source)
		//We need to reset
		reset = true;

	//If we need to reset
	if (reset)
	{
		Debug("-StreamTransponder::onRTP() | Reset stream\n");
		//IF last was not completed
		if (!lastCompleted && type==MediaFrame::Video)
		{
			//Create new RTP packet
			RTPPacket::shared rtp = std::make_shared<RTPPacket>(media,codec);
			//Set data
			rtp->SetPayloadType(type);
			rtp->SetSSRC(ssrc);
			rtp->SetExtSeqNum(lastExtSeqNum++);
			rtp->SetMark(true);
			rtp->SetExtTimestamp(lastTimestamp);
			//Send it
			if (sender) sender->Enqueue(rtp);
		}
		//No source
		lastCompleted = true;
		source = 0;
		//Reset first paquet seq num and timestamp
		firstExtSeqNum = NoSeqNum;
		firstTimestamp = NoTimestamp;
		//Store the last send ones
		baseExtSeqNum = lastExtSeqNum+1;
		baseTimestamp = lastTimestamp;
		//None dropped or added
		dropped = 0;
		added = 0;
		//Not selecting
		selector = nullptr;
		//No layer
		spatialLayerId = LayerInfo::MaxLayerId;
		temporalLayerId = LayerInfo::MaxLayerId;

		//Reset frame numbers
		firstFrameNumber = NoFrameNum;
		baseFrameNumber = lastFrameNumber + 1;
		frameNumberExtender.Reset();

		//Reseted
		reset = false;
	}

	//Update source
	source = packet->GetSSRC();
	//Get new seq number
	DWORD extSeqNum = packet->GetExtSeqNum();

	//Check if it the first received packet
	if (firstExtSeqNum==NoSeqNum || firstTimestamp==NoTimestamp)
	{
		//If we have a time offest from last sent packet
		if (lastTime)
		{
			//Calculate time diff
			QWORD offset = getTimeDiff(lastTime)/1000;
			//Get timestamp diff on correct clock rate
			QWORD diff = offset*packet->GetClockRate()/1000;

			//UltraDebug("-ts offset:%llu diff:%llu baseTimestap:%lu firstTimestamp:%llu lastTimestamp:%llu rate:%llu\n",offset,diff,baseTimestamp,firstTimestamp,lastTimestamp,packet->GetClockRate());

			//convert it to rtp time and add to the last sent timestamp
			baseTimestamp = lastTimestamp + diff + 1;
		}

		//Reset drop counter
		dropped = 0;
		//Store seq number
		firstExtSeqNum = extSeqNum;
		//Get first timestamp
		firstTimestamp = packet->GetExtTimestamp();

		UltraDebug("-StreamTransponder::onRTP() | first seq:%u base:%u last:%u ts:%llu baseSeq:%u baseTimestamp:%llu lastTimestamp:%llu\n",firstExtSeqNum,baseExtSeqNum,lastExtSeqNum,firstTimestamp,baseExtSeqNum,baseTimestamp,lastTimestamp);
	}

	//Ensure it is not before first one
	if (extSeqNum<firstExtSeqNum)
		//Exit
		return;

	//Only for viedo
	if (packet->GetMediaType()==MediaFrame::Video)
	{
		//Check if we don't have one or if we have a selector and it is not from the same codec
		if (!selector || (BYTE)selector->GetCodec()!=packet->GetCodec())
		{
			//Create new selector for codec
			selector.reset(VideoLayerSelector::Create((VideoCodec::Type)packet->GetCodec()));
			//Set prev layers
			selector->SelectSpatialLayer(spatialLayerId);
			selector->SelectTemporalLayer(temporalLayerId);
		}
	}

	//Get rtp marking
	bool mark = packet->GetMark();

	//If we have selector for codec
	if (selector)
	{
		//Select layer
		selector->SelectSpatialLayer(spatialLayerId);
		selector->SelectTemporalLayer(temporalLayerId);

		//Select pacekt
		if (!packet->GetMediaLength() || !selector->Select(packet,mark))
		{
			//One more dropperd
			dropped++;
			//If selector is waiting for intra and last PLI was more than 1s ago
			if (selector->IsWaitingForIntra() && getTimeDiffMS(lastSentPLI)>1E3)
			{
				//Log
				//UltraDebug("-RTPStreamTransponder::onRTP() | selector IsWaitingForIntra\n");
				//Request it again
				RequestPLI();
			}
			//Drop
			return;
		}
		//Get current spatial layer id
		lastSpatialLayerId = selector->GetSpatialLayer();
	}

	//Set normalized seq num
	extSeqNum = baseExtSeqNum + (extSeqNum - firstExtSeqNum) - dropped + added;

	//Set normailized timestamp
	uint64_t timestamp = baseTimestamp + (packet->GetExtTimestamp()-firstTimestamp);

	//UPdate media codec and type
	media = packet->GetMediaType();
	codec = packet->GetCodec();
	type  = packet->GetPayloadType();

	//UltraDebug("-ext seq:%lu base:%lu first:%lu current:%lu dropped:%lu added:%d ts:%lu normalized:%llu intra:%d codec=%d\n",extSeqNum,baseExtSeqNum,firstExtSeqNum,packet->GetExtSeqNum(),dropped,added,packet->GetTimestamp(),timestamp,packet->IsKeyFrame(),codec);

	//TODO: this should go into the layer selector??
	if (codec==VideoCodec::VP8 && packet->vp8PayloadDescriptor)
	{
		//Get VP8 description
		auto& vp8PayloadDescriptor = *packet->vp8PayloadDescriptor;

		//Check if we have a pictId
		if (vp8PayloadDescriptor.pictureIdPresent)
		{
			//If we have not received any yet
			if (!pictureId)
			{
				//Use current as starting point
				pictureId = vp8PayloadDescriptor.pictureId;
					
			} 
			//If picture id is different than last received one
			else if (vp8PayloadDescriptor.pictureId != lastSrcPictureId)
			{
				//Increase picture id
				(*pictureId)++;
			}

			//Update last received pict id
			lastSrcPictureId = vp8PayloadDescriptor.pictureId;
		}

		//Check if we have a new base layer
		if (vp8PayloadDescriptor.temporalLevelZeroIndexPresent)
		{
			/*
				* TL0PICIDX:  8 bits temporal level zero index.TL0PICIDX is a
				* running index for the temporal base layer frames, i.e., the
				* frames with TID set to 0.  If TID is larger than 0, TL0PICIDX
				* indicates on which base - layer frame the current image depends.
				* TL0PICIDX MUST be incremented when TID is 0.  The index MAY
				* start at a random value, and it MUST wrap to 0 after reaching
				* the maximum number 255.  Use of TL0PICIDX depends on the
				* presence of TID.Therefore, it is RECOMMENDED that the TID be
				* used whenever TL0PICIDX is.
			*/

			//Check if it is the base temporal layer
			if (vp8PayloadDescriptor.temporalLayerIndex == 0)
			{
				// If we have not received any tl0 index yet
				if (!temporalLevelZeroIndex)
				{
					//Use current as starting point
					temporalLevelZeroIndex = vp8PayloadDescriptor.temporalLevelZeroIndex;
				}
				//If it is different than last received tl0 index
				else if (vp8PayloadDescriptor.temporalLevelZeroIndex != lastSrcTemporalLevelZeroIndex)
				{
					//Increase tl0 index
					(*temporalLevelZeroIndex)++;
				}

				//Update last received tl0 index
				lastSrcTemporalLevelZeroIndex = vp8PayloadDescriptor.temporalLevelZeroIndex;
			}

		}

		// Set if we need rewrite the picture ID or tl0 index
		if (temporalLevelZeroIndex && pictureId && 
			( vp8PayloadDescriptor.pictureId != *pictureId || vp8PayloadDescriptor.temporalLevelZeroIndex != *temporalLevelZeroIndex)
		)
		{
			//Rewrite ids
			packet->rewitePictureIds = true;

			//Rewrite picture id
			vp8PayloadDescriptor.pictureId = *pictureId;
			//Rewrite tl0 index
			vp8PayloadDescriptor.temporalLevelZeroIndex = *temporalLevelZeroIndex;
		}

		//Error("-ext seq:%lu pictureIdPresent:%d rewrite:%d pictId:%d lastPictId:%d origPictId:%d intra:%d mark:%d \n",extSeqNum, vp8PayloadDescriptor.pictureIdPresent, packet->rewitePictureIds, *pictureId, lastSrcPicId, vp8PayloadDescriptor.pictureId : -1, packet->IsKeyFrame(), packet->GetMark());
	}

	//If we have to append h264 sprop parameters set for the first packet of an iframe
	if (h264Parameters && codec==VideoCodec::H264 && packet->IsKeyFrame() && (timestamp!=lastTimestamp || firstExtSeqNum==packet->GetExtSeqNum()))
	{
		//UltraDebug("-addding h264 sprop\n");

		//Clone packet
		auto cloned = h264Parameters->Clone();
		//Set new seq numbers
		cloned->SetExtSeqNum(extSeqNum);
		//Set normailized timestamp
		cloned->SetExtTimestamp(timestamp);
		//Set payload type
		cloned->SetPayloadType(type);
		//Change ssrc
		cloned->SetSSRC(ssrc);
		//Send packet
		if (sender)
			//Create clone on sender thread
			sender->Enqueue(cloned);
		//Add new packet
		added ++;
		extSeqNum ++;

		//UltraDebug("-ext seq:%lu base:%lu first:%lu current:%lu dropped:%lu added:%d ts:%lu normalized:%llu intra:%d codec=%d\n",extSeqNum,baseExtSeqNum,firstExtSeqNum,packet->GetExtSeqNum(),dropped,added,packet->GetTimestamp(),timestamp,packet->IsKeyFrame(),codec);
	}

	//Dependency descriptor active decodte target mask
	std::optional<std::vector<bool>> forwaredDecodeTargets;

	//If it is AV1
	if (codec==VideoCodec::AV1 && selector)
		//Get decode target
		forwaredDecodeTargets = static_cast<AV1LayerSelector*>(selector.get())->GetForwardedDecodeTargets();

	//Continous frame number
	uint64_t continousFrameNumber = NoFrameNum;
	//Get frame number if we have dependency descriptor
	if (packet->HasDependencyDestriptor())
	{
		//Get it
		auto dd = packet->GetDependencyDescriptor();

		//Double check
		if (dd)
		{
			//Extend it
			frameNumberExtender.Extend(dd->frameNumber);
			//Get extended frame number
			uint64_t frameNumber = frameNumberExtender.GetExtSeqNum();

			//If it is first
			if (firstFrameNumber==NoFrameNum)
				//Set it
				firstFrameNumber = frameNumber;
			//If it is the first frame after reset
			if (baseFrameNumber==NoFrameNum)
				//Set it
				baseFrameNumber = frameNumber;
			//Calculate a continous frame number
			continousFrameNumber = baseFrameNumber + frameNumber - firstFrameNumber;

			//UltraDebug("-frameNum first:%llu base:%llu current:%llu(%u) continous=%llu\n", firstFrameNumber, baseFrameNumber, lastFrameNumber, dd->frameNumber, continousFrameNumber);
		}
	}

	//Get last send seq num and timestamp
	lastExtSeqNum = extSeqNum;
	lastTimestamp = timestamp;
	//Update last sent time
	lastTime = getTime();

	//Get last frame number
	lastFrameNumber = continousFrameNumber;

	//Set new seq numbers
	packet->SetExtSeqNum(extSeqNum);
	//Set normailized timestamp
	packet->SetTimestamp(timestamp);
	//Set mark again
	packet->SetMark(mark);
	//Change ssrc
	packet->SetSSRC(ssrc);

	//If it has a dependency descriptor
	if (forwaredDecodeTargets && packet->HasTemplateDependencyStructure())
		//Override mak
		packet->OverrideActiveDecodeTargets(forwaredDecodeTargets);
	//If we have a continous frame number
	if (packet->HasDependencyDestriptor() && continousFrameNumber != NoFrameNum)
		//Update it
		packet->OverrideFrameNumber(static_cast<uint16_t>(continousFrameNumber));

	//Send packet
	if (sender)
		//Enqueue it
		sender->Enqueue(packet);
}

void RTPStreamTransponder::onBye(const RTPIncomingMediaStream* stream)
{
	AsyncSafe([=](auto) {

		//If they are the not same
		if (incoming.get() != stream)
			//DO nothing
			return;

		//Reset packets
		reset = true;
	});
}

void RTPStreamTransponder::onEnded(const RTPIncomingMediaStream* stream)
{
	AsyncSafe([=](auto) {
		// Note: Checks for equality of stream pointer are important to ensure it is still valid on exec!

		//IF it is the current one
		if (this->incoming.get() == stream)
		{
			//Reset packets before start listening again
			reset = true;

			//No stream and receiver
			this->incoming = nullptr;
			this->receiver = nullptr;
		//If it is the next one
		} else if (this->incomingNext.get() == stream) {
			//No transitioning stream and receiver
			this->incomingNext = nullptr;
			this->receiverNext = nullptr;
		}
	});
}
void RTPStreamTransponder::onEnded(const RTPOutgoingSourceGroup* group)
{
	AsyncSafe([=](auto){
		// Note: Checks for equality of group pointer are important to ensure it is still valid on exec!

		//IF it is the current one
		if (outgoing.get() == group)
			//No more outgoing
			outgoing = nullptr;
	});
}


void RTPStreamTransponder::RequestPLI()
{
	//Log("-RTPStreamTransponder::RequestPLI() [receiver:%p,incoming:%p]\n",receiver,incoming);
	AsyncSafe([=](auto now) {
		//Request update on the incoming
		if (receiver && incoming) receiver->SendPLI(incoming->GetMediaSSRC());
		//Update last sent pli
		lastSentPLI = now.count();
	});
}

void RTPStreamTransponder::onPLIRequest(const RTPOutgoingSourceGroup* group,DWORD ssrc)
{
	//Log
	UltraDebug("-RTPStreamTransponder::onPLIRequest()\n");
	RequestPLI();
}

void RTPStreamTransponder::onREMB(const RTPOutgoingSourceGroup* group,DWORD ssrc, DWORD bitrate)
{
	UltraDebug("-RTPStreamTransponder::onREMB() [ssrc:%u,bitrate:%u]\n",ssrc,bitrate);
}


void RTPStreamTransponder::SelectLayer(int spatialLayerId,int temporalLayerId)
{
	//Log
	UltraDebug("-RTPStreamTransponder::SelectLayer() | [sid:%d,tid:%d]\n", spatialLayerId, temporalLayerId);

	this->spatialLayerId  = spatialLayerId;
	this->temporalLayerId = temporalLayerId;

	if (lastSpatialLayerId!=spatialLayerId)
		//Request update on the incoming
		RequestPLI();
}

void RTPStreamTransponder::Mute(bool muting)
{
	//Log
	UltraDebug("-RTPStreamTransponder::Mute() | [muting:%d]\n", muting);

	//Check if we are changing state
	if (muting==muted)
		//Do nothing
		return;
	//If unmutting
	if (!muting)
		//Request update
		RequestPLI();
	//Update state
	muted = muting;
}

void RTPStreamTransponder::SetIntraOnlyForwarding(bool intraOnlyForwarding)
{
	//Log
	UltraDebug("-RTPStreamTransponder::SetIntraOnlyForwarding() | [intraOnlyForwarding:%d]\n", intraOnlyForwarding);

	//Set flag
	this->intraOnlyForwarding = intraOnlyForwarding;
}

bool RTPStreamTransponder::AppendH264ParameterSets(const std::string& sprop)
{

	Debug("-RTPStreamTransponder::AppendH264ParameterSets() [sprop:%s]\n",sprop.c_str());

	//Create pakcet
	auto rtp = std::make_shared<RTPPacket>(MediaFrame::Video,VideoCodec::H264);

	//Get current length
	BYTE* data = rtp->AdquireMediaData();
	DWORD len = 0;
	DWORD size = rtp->GetMaxMediaLength();
	//Append stap-a header
	data[len++] = 24;
	//Split by ","
	auto start = 0;

	//Get next separator
	auto end = sprop.find(',');

	//Parse
	while(end!=std::string::npos)
	{
		//Get prop
		auto prop = sprop.substr(start,end-start);

		Debug("-RTPStreamTransponder::AppendH264ParameterSets() [sprop:%s]\n",prop.c_str());
		//Parse and keep space for size
		auto l = av_base64_decode(data+len+2,prop.c_str(),size-len-2);
		//Check result
		if (l<=0)
			return Error("-RTPStreamTransponder::AppendH264ParameterSets() could not decode base64 data [%s]\n",prop.c_str());
		//Set naly length
		set2(data,len,l);
		//Increase length
		len += l+2;
		//Next
		start = end+1;
		//Get nest
		end = sprop.find(',',start);
	}
	//last one
	auto prop = sprop.substr(start,end-start);

	//Parse and keep space for size
	auto l = av_base64_decode(data+len+2,prop.c_str(),size-len-2);
	//Check result
	if (l<=0)
		return Error("-RTPStreamTransponder::AppendH264ParameterSets() could not decode base64 data [%s]\n",prop.c_str());
	//Set naly length
	set2(data,len,l);
	//Increase length
	len += l+2;
	//Set new lenght
	rtp->SetMediaLength(len);

	//Store it
	this->h264Parameters = rtp;

	//Done
	return true;
}
