/*****************************************************************************

 * ts_hotfixes.c : MPEG PMT/PAT less streams fixups

 *****************************************************************************

 * Copyright (C) 2014-2016 - VideoLAN Authors

 *

 * This program is free software; you can redistribute it and/or modify it

 * under the terms of the GNU Lesser General Public License as published by

 * the Free Software Foundation; either version 2.1 of the License, or

 * (at your option) any later version.

 *

 * This program is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

 * GNU Lesser General Public License for more details.

 *

 * You should have received a copy of the GNU General Public License

 * along with this program.  If not, see <http://www.gnu.org/licenses/>.

 *****************************************************************************/

#ifdef HAVE_CONFIG_H

# include "config.h"

#endif

#include <vlc_common.h>

#include <vlc_demux.h>

#include <vlc_es.h>

#ifndef _DVBPSI_DVBPSI_H_

 #include <dvbpsi/dvbpsi.h>

#endif

#include <dvbpsi/descriptor.h>

#include <dvbpsi/pat.h>

#include <dvbpsi/pmt.h>

#include "../../mux/mpeg/streams.h"

#include "../../mux/mpeg/tsutil.h"

#include "../../mux/mpeg/tables.h"

#include "timestamps.h"

#include "pes.h"

#include "ts_streams.h"

#include "ts_psi.h"

#include "ts_pid.h"

#include "ts_streams_private.h"

#include "ts.h"

#include "ts_hotfixes.h"

#include "ts_packet.h"

#include <assert.h>

void ProbePES( demux_t *p_demux, ts_pid_t *pid, const block_t *p_pkt )

{

    demux_sys_t *p_sys = p_demux->p_sys;

    unsigned i_skip = PKTHeaderAndAFSize( p_pkt );

    if ( p_pkt->i_buffer < i_skip )

        return;

    ts_90khz_t pktpcr = GetPCR( p_pkt );

    if( pktpcr != TS_90KHZ_INVALID )

        pid->probed.i_pcr_count++;

    size_t i_data = p_pkt->i_buffer - i_skip;

    const uint8_t *p_pes = &p_pkt->p_buffer[i_skip];

    ts_pes_header_t pesh;

    ts_pes_header_init( &pesh );

    if( ParsePESHeader( NULL, p_pes, i_data, &pesh ) != VLC_SUCCESS )

        return;

    if( pesh.i_dts != TS_90KHZ_INVALID )

        pid->probed.i_dts_count++;

    if( pid->probed.i_fourcc != 0 )

        goto codecprobingend;

    if( i_data < pesh.i_size + 4 )

        return;

    const uint8_t *p_data = &p_pes[pesh.i_size];

    const uint8_t i_stream_id = pid->probed.i_stream_id = p_pes[3];

    /* NON MPEG audio & subpictures STREAM */

    if(i_stream_id == 0xBD)

    {

        if( !memcmp( p_data, "\x7F\xFE\x80\x01", 4 ) )

        {

            pid->probed.i_fourcc = VLC_CODEC_DTS;

            pid->probed.i_cat = AUDIO_ES;

        }

        else if( !memcmp( p_data, "\x0B\x77", 2 ) )

        {

            pid->probed.i_fourcc = VLC_CODEC_EAC3;

            pid->probed.i_cat = AUDIO_ES;

        }

    }

    /* MPEG AUDIO STREAM */

    else if(i_stream_id >= 0xC0 && i_stream_id <= 0xDF)

    {

        pid->probed.i_cat = AUDIO_ES;

        if( p_data[0] == 0xFF && (p_data[1] & 0xE0) == 0xE0 &&

           (p_data[1] & 0x18) != 0x08 && (p_data[1] & 0x06) != 0x00 )

        {

            pid->probed.i_fourcc = VLC_CODEC_MPGA;

        }

        else if( p_data[0] == 0xFF && (p_data[1] & 0xF6) == 0xF0 )

        {

            pid->probed.i_fourcc = VLC_CODEC_MP4A; /* ADTS */

            pid->probed.i_original_fourcc = VLC_FOURCC('A','D','T','S');

        }

    }

    /* VIDEO STREAM */

    else if( i_stream_id >= 0xE0 && i_stream_id <= 0xEF )

    {

        pid->probed.i_cat = VIDEO_ES;

        if( !memcmp( p_data, "\x00\x00\x00\x01", 4 ) )

        {

            pid->probed.i_fourcc = VLC_CODEC_H264;

        }

        else if( !memcmp( p_data, "\x00\x00\x01", 3 ) )

        {

            pid->probed.i_fourcc = VLC_CODEC_MPGV;

        }

    }

codecprobingend:

    /* Track timestamps and flag missing PAT */

    if( !p_sys->patfix.i_timesourcepid && pesh.i_dts != TS_90KHZ_INVALID )

    {

        p_sys->patfix.i_first_dts = FROM_SCALE(pesh.i_dts);

        p_sys->patfix.i_timesourcepid = pid->i_pid;

    }

    else if( p_sys->patfix.i_timesourcepid == pid->i_pid && pesh.i_dts != TS_90KHZ_INVALID &&

             p_sys->patfix.status == PAT_WAITING )

    {

        if( FROM_SCALE(pesh.i_dts) - p_sys->patfix.i_first_dts > MIN_PAT_INTERVAL )

            p_sys->patfix.status = PAT_MISSING;

    }

}

static void BuildPATCallback( void *p_opaque, block_t *p_block )

{

    ts_pid_t *pat_pid = (ts_pid_t *) p_opaque;

    ts_psi_Packet_Push( pat_pid, p_block->p_buffer );

    block_Release( p_block );

}

static void BuildPMTCallback( void *p_opaque, block_t *p_block )

{

    ts_pid_t *program_pid = (ts_pid_t *) p_opaque;

    assert(program_pid->type == TYPE_PMT);

    while( p_block )

    {

        ts_psi_Packet_Push( program_pid, p_block->p_buffer );

        block_t *p_next = p_block->p_next;

        block_Release( p_block );

        p_block = p_next;

    }

}

void MissingPATPMTFixup( demux_t *p_demux )

{

    demux_sys_t *p_sys = p_demux->p_sys;

    int i_program_number = 1234;

    int i_program_pid = 1337;

    int i_pcr_pid = 0x1FFF;

    int i_num_pes = 0;

    ts_pid_t *p_program_pid = GetPID( p_sys, i_program_pid );

    if( SEEN(p_program_pid) )

    {

        /* Find a free one */

        for( i_program_pid = MIN_ES_PID;

             i_program_pid <= MAX_ES_PID && SEEN(p_program_pid);

             i_program_pid++ )

        {

            p_program_pid = GetPID( p_sys, i_program_pid );

        }

    }

    const ts_pid_t * candidates[4] = { NULL };

    const ts_pid_t *p_pid = NULL;

    ts_pid_next_context_t pidnextctx = ts_pid_NextContextInitValue;

    while( (p_pid = ts_pid_Next( &p_sys->pids, &pidnextctx )) )

    {

        if( !SEEN(p_pid) || p_pid->probed.i_fourcc == 0 )

            continue;

        if( p_pid->probed.i_pcr_count && candidates[0] == NULL && false )

            candidates[0] = p_pid;

        if( p_pid->probed.i_cat == AUDIO_ES &&

            (candidates[1] == NULL ||

             candidates[1]->probed.i_dts_count > p_pid->probed.i_dts_count) )

            candidates[1] = p_pid;

        if( candidates[2] == NULL && p_pid != candidates[1] &&

            p_pid->probed.i_dts_count > 0 )

            candidates[2] = p_pid;

        if( candidates[3] == NULL )

            candidates[3] = p_pid;

        i_num_pes++;

    }

    for(int i=0; i<4; i++)

    {

        if(!candidates[i])

            continue;

        i_pcr_pid = candidates[i]->i_pid;

        p_sys->patfix.b_pcrhasnopcrfield = (candidates[i]->probed.i_pcr_count < 1);

        break;

    }

    if( i_num_pes == 0 )

        return;

    tsmux_stream_t patstream =

    {

        .i_pid = 0,

        .i_continuity_counter = 0x10,

        .b_discontinuity = false

    };

    tsmux_stream_t pmtprogramstream =

    {

        .i_pid = i_program_pid,

        .i_continuity_counter = 0x0,

        .b_discontinuity = false

    };

    dvbpsi_t *handle = dvbpsi_new( NULL, DVBPSI_MSG_DEBUG );

    if( !handle )

        return;

    BuildPAT( handle,

            &p_sys->pids.pat, BuildPATCallback,

            0, 1,

            &patstream,

            1, &pmtprogramstream, &i_program_number );

    /* PAT callback should have been triggered */

    if( p_program_pid->type != TYPE_PMT )

    {

        dvbpsi_delete( handle );

        msg_Err( p_demux, "PAT creation failed" );

        return;

    }

    ts_mux_standard mux_standard = (p_sys->standard == TS_STANDARD_ATSC) ? TS_MUX_STANDARD_ATSC

                                                                         : TS_MUX_STANDARD_DVB;

    struct esstreams_t

    {

        pesmux_stream_t pes;

        tsmux_stream_t ts;

        es_format_t fmt;

    };

    struct esstreams_t *esstreams = calloc( i_num_pes, sizeof(struct esstreams_t) );

    pes_mapped_stream_t *mapped = calloc( i_num_pes, sizeof(pes_mapped_stream_t) );

    if( esstreams && mapped )

    {

        int j=0;

        for( int i=0; i<p_sys->pids.i_all; i++ )

        {

            p_pid = p_sys->pids.pp_all[i];

            if( !SEEN(p_pid) ||

                p_pid->probed.i_fourcc == 0 )

                continue;

            es_format_Init(&esstreams[j].fmt, p_pid->probed.i_cat, p_pid->probed.i_fourcc);

            esstreams[j].fmt.i_original_fourcc = p_pid->probed.i_original_fourcc;

            if( VLC_SUCCESS !=

                FillPMTESParams(mux_standard, &esstreams[j].fmt, &esstreams[j].ts, &esstreams[j].pes ) )

            {

                es_format_Clean( &esstreams[j].fmt );

                continue;

            }

            /* Important for correct remapping: Enforce probed PES stream id */

            esstreams[j].pes.i_stream_id = p_pid->probed.i_stream_id;

            esstreams[j].ts.i_pid = p_pid->i_pid;

            mapped[j].pes = &esstreams[j].pes;

            mapped[j].ts = &esstreams[j].ts;

            mapped[j].fmt = &esstreams[j].fmt;

            j++;

        }

        BuildPMT( handle, VLC_OBJECT(p_demux),

                 mux_standard,

                p_program_pid, BuildPMTCallback,

                0, 1,

                i_pcr_pid,

                NULL,

                1, &pmtprogramstream, &i_program_number,

                j, mapped );

        /* Cleanup */

        for( int i=0; i<j; i++ )

            es_format_Clean( &esstreams[i].fmt );

    }

    free(esstreams);

    free(mapped);

    dvbpsi_delete( handle );

}