// ITU-T H.222.0(06/2012)
// Information technology - Generic coding of moving pictures and associated audio information: Systems
// 2.5.3.1 Program stream(p74)

#include "mpeg-pes-internal.h"
#include "mpeg-ps-internal.h"
#include "mpeg-util.h"
#include "mpeg-ps.h"
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <time.h>

#define MAX_PES_HEADER	1024	// pack_header + system_header + psm
#define MAX_PES_PACKET	0xFF80	// 64k pes data, reserved 0x7F for hik

struct ps_muxer_t
{
    struct psm_t psm; 
    struct ps_pack_header_t pack;
    struct ps_system_header_t system;
    
    int h26x_with_aud;
	unsigned int psm_period;
	unsigned int scr_period;

	struct ps_muxer_func_t func;
	void* param;

//	uint8_t packet[MAX_PACKET_SIZE];
};

static struct pes_t* ps_stream_find(struct ps_muxer_t *ps, int streamid)
{
    size_t i;
    for (i = 0; i < ps->psm.stream_count; i++)
    {
        if (streamid == ps->psm.streams[i].sid)
            return &ps->psm.streams[i];
    }
    return NULL;
}

int ps_muxer_input(struct ps_muxer_t* ps, int streamid, int flags, int64_t pts, int64_t dts, const void* data, size_t bytes)
{
	int r, first;
	size_t i, n, sz;
	uint8_t *packet;
    struct pes_t* stream;
    const uint8_t* payload;

	i = 0;
	first = 1;
	payload = (const uint8_t*)data;

    stream = ps_stream_find(ps, streamid);
    if (NULL == stream) return -1; // not found
    stream->data_alignment_indicator = (flags & MPEG_FLAG_IDR_FRAME) ? 1 : 0; // idr frame
    stream->pts = pts;
    stream->dts = dts;

	// Add PSM for IDR frame
	ps->psm_period = ((flags & MPEG_FLAG_IDR_FRAME) && mpeg_stream_type_video(stream->codecid)) ? 0 : ps->psm_period;
    ps->h26x_with_aud = (flags & MPEG_FLAG_H264_H265_WITH_AUD) ? 1 : 0;

	// TODO: 
	// 1. update packet header program_mux_rate
	// 2. update system header rate_bound

	// alloc once (include Multi-PES packet)
	sz = bytes + MAX_PES_HEADER + (bytes/MAX_PES_PACKET+1) * 64; // 64 = 0x000001 + stream_id + PES_packet_length + other
	packet = ps->func.alloc(ps->param, sz);
	if(!packet) return -ENOMEM;

	// write pack_header(p74)
	// 2.7.1 Frequency of coding the system clock reference
	// http://www.bretl.com/mpeghtml/SCR.HTM
	//the maximum allowed interval between SCRs is 700ms 
	//ps->pack.system_clock_reference_base = (dts-3600) % (((int64_t)1)<<33);
	ps->pack.system_clock_reference_base = dts >= 3600 ? (dts - 3600) : 0;
	ps->pack.system_clock_reference_extension = 0;
	ps->pack.program_mux_rate = 6106;
	i += pack_header_write(&ps->pack, packet + i);

#if !defined(MPEG_FIX_VLC_3_X_PS_SYSTEM_HEADER)
    // https://github.com/videolan/vlc/blob/3.0.x/modules/demux/mpeg/ps.h#L488
    // fix ps_pkt_parse_system -> ps_track_fill with default mp1/2 audio codec(without psm)
    
	// write system_header(p76)
	if(0 == (ps->psm_period % 30))
		i += system_header_write(&ps->system, packet + i);
#endif

	// write program_stream_map(p79)
	if (0 == (ps->psm_period % 30))
	{
#if defined(MPEG_CLOCK_EXTENSION_DESCRIPTOR)
		ps->psm.clock = time() * 1000; // todo: gettimeofday
#endif
		i += psm_write(&ps->psm, packet + i);
	}

	// check packet size
	assert(i < MAX_PES_HEADER);

	// write data
	while(bytes > 0)
	{
		uint8_t *p;
		uint8_t *pes = packet + i;
		
		p = pes + pes_write_header(stream, pes, sz - i);
		stream->pts = stream->dts = PTS_NO_VALUE; // clear pts/dts flags
		stream->data_alignment_indicator = 0; // clear flags
		assert(p - pes < 64);

		if(first)
		{
			if (PSI_STREAM_H264 == stream->codecid && !ps->h26x_with_aud)
			{
				// 2.14 Carriage of Rec. ITU-T H.264 | ISO/IEC 14496-10 video
				// Each AVC access unit shall contain an access unit delimiter NAL Unit
				nbo_w32(p, 0x00000001);
				p[4] = 0x09; // AUD
				p[5] = 0xE0; // any slice type (0xe) + rbsp stop one bit
				p += 6;
			}
			else if (PSI_STREAM_H265 == stream->codecid && !ps->h26x_with_aud)
			{
				// 2.17 Carriage of HEVC
				// Each HEVC access unit shall contain an access unit delimiter NAL unit.
				nbo_w32(p, 0x00000001);
				p[4] = 0x46; // 35-AUD_NUT
				p[5] = 0x01;
				p[6] = 0x50; // B&P&I (0x2) + rbsp stop one bit
				p += 7;
			}
			else if (PSI_STREAM_H266 == stream->codecid && !ps->h26x_with_aud)
			{
				// 2.23 Carriage of VVC
				// Each VVC access unit shall contain an access unit delimiter NAL unit
				nbo_w32(p, 0x00000001);
				p[4] = 0x00; // 20-AUD_NUT
				p[5] = 0xA1;
				p[6] = 0x18; // B&P&I (0x1) + rbsp stop one bit
				p += 7;
			}
		}

		// PES_packet_length = PES-Header + Payload-Size
		// A value of 0 indicates that the PES packet length is neither specified nor bounded 
		// and is allowed only in PES packets whose payload consists of bytes from a 
		// video elementary stream contained in transport stream packets
		if((p - pes - 6) + bytes > MAX_PES_PACKET)
		{
			nbo_w16(pes + 4, MAX_PES_PACKET);
			n = MAX_PES_PACKET - (p - pes - 6);
		}
		else
		{
			nbo_w16(pes + 4, (uint16_t)((p - pes - 6) + bytes));
			n = bytes;
		}

		memcpy(p, payload, n);
		payload += n;
		bytes -= n;

		// notify packet already
		i += n + (p - pes);

//		i = 0; // clear value, the next pes packet don't need pack_header
		first = 0; // clear first packet flag
		pts = dts = 0; // only first packet write PTS/DTS
	}

	assert(i < sz);
	r = ps->func.write(ps->param, stream->sid, packet, i);
	ps->func.free(ps->param, packet);

	++ps->psm_period;
	return r;
}

struct ps_muxer_t* ps_muxer_create(const struct ps_muxer_func_t *func, void* param)
{
	struct ps_muxer_t *ps = NULL;

	assert(func);
	ps = (struct ps_muxer_t *)calloc(1, sizeof(struct ps_muxer_t));
	if(!ps)
		return NULL;

	memcpy(&ps->func, func, sizeof(ps->func));
	ps->param = param;

	ps->system.rate_bound = 26234; //10493600~10mbps(50BPS * 8 = 400bps)
//	ps->system.audio_bound = 1; // [0,32] max active audio streams
//	ps->system.video_bound = 1; // [0,16] max active video streams
	ps->system.fixed_flag = 0; // 1-fixed bitrate, 0-variable bitrate
	ps->system.CSPS_flag = 0; // meets the constraints defined in 2.7.9.
	ps->system.packet_rate_restriction_flag = 0; // dependence CSPS_flag
	ps->system.system_audio_lock_flag = 0; // all audio stream sampling rate is constant
	ps->system.system_video_lock_flag = 0; // all video stream frequency is constant

	//ps->psm.ver = 1;
	//ps->psm.stream_count = 2;
	//ps->psm.streams[0].element_stream_id = PES_SID_VIDEO;
	//ps->psm.streams[0].stream_type = PSI_STREAM_H264;
	//ps->psm.streams[1].element_stream_id = PES_SID_AUDIO;
	//ps->psm.streams[1].stream_type = PSI_STREAM_AAC;

	return ps;
}

int ps_muxer_destroy(struct ps_muxer_t* ps)
{
    size_t i;
    for (i = 0; i < ps->psm.stream_count; i++)
    {
        if (ps->psm.streams[i].esinfo)
        {
            free(ps->psm.streams[i].esinfo);
            ps->psm.streams[i].esinfo = NULL;
        }
    }

	free(ps);
	return 0;
}

int ps_muxer_add_stream(struct ps_muxer_t* ps, int codecid, const void* extradata, size_t bytes)
{
	struct psm_t* psm;
	struct pes_t* pes;

	assert(bytes < 512);
	if (!ps || ps->psm.stream_count >= sizeof(ps->psm.streams) / sizeof(ps->psm.streams[0]))
	{
		assert(0);
		return -1;
	}

	psm = &ps->psm;
	pes = &psm->streams[psm->stream_count];

	if (mpeg_stream_type_video(codecid))
	{
		pes->sid = (uint8_t)(PES_SID_VIDEO + ps->system.video_bound);

		assert(ps->system.video_bound + 1 < 16);
		++ps->system.video_bound; // [0,16] max active video streams
		ps->system.streams[ps->system.stream_count].buffer_bound_scale = 1;
		/* FIXME -- VCD uses 46, SVCD uses 230, ffmpeg has 230 with a note that it is small */
		ps->system.streams[ps->system.stream_count].buffer_size_bound = 400 /* 8191-13 bits max value */;
	}
	else if (mpeg_stream_type_audio(codecid))
	{
		pes->sid = (uint8_t)(PES_SID_AUDIO + ps->system.audio_bound);

		assert(ps->system.audio_bound + 1 < 32);
		++ps->system.audio_bound; // [0,32] max active audio streams
		ps->system.streams[ps->system.stream_count].buffer_bound_scale = 0;
		/* This value HAS to be used for VCD (see VCD standard, p. IV-7).
		* Right now it is also used for everything else. */
		ps->system.streams[ps->system.stream_count].buffer_size_bound = 32 /* 4 * 1024 / 128 */;
	}
	else
	{
		assert(0);
		return -1;
	}

	if (bytes > 0)
	{
		pes->esinfo = (uint8_t*)malloc(bytes);
		if (!pes->esinfo)
			return -1;
		memcpy(pes->esinfo, extradata, bytes);
		pes->esinfo_len = (uint16_t)bytes;
	}

    assert(psm->stream_count == ps->system.stream_count);
    ps->system.streams[ps->system.stream_count].stream_id = pes->sid;
	++ps->system.stream_count;

    pes->codecid = (uint8_t)codecid;
	++psm->stream_count;
	++psm->ver;

	ps->psm_period = 0; // immediate update psm
	return pes->sid;
}