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

 * util.cpp : matroska demuxer

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

 * Copyright (C) 2003-2004 VLC authors and VideoLAN

 *

 * Authors: Laurent Aimar <fenrir@via.ecp.fr>

 *          Steve Lhomme <steve.lhomme@free.fr>

 *

 * 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 Lesser General Public License

 * along with this program; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.

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

#include "mkv.hpp"

#include "util.hpp"

#include "demux.hpp"

#include "virtual_segment.hpp"

#include "../../codec/webvtt/helpers.h"

#include "lzokay.hpp"

namespace mkv {

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

 * Local prototypes

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

#ifdef HAVE_ZLIB

int32_t zlib_decompress_extra( demux_t * p_demux, mkv_track_t & tk )

{

    int result;

    z_stream d_stream;

    size_t n = 0;

    uint8_t * p_new_extra = NULL;

    msg_Dbg(p_demux,"Inflating private data");

    d_stream.zalloc = Z_NULL;

    d_stream.zfree = Z_NULL;

    d_stream.opaque = Z_NULL;

    if( inflateInit( &d_stream ) != Z_OK )

    {

        msg_Err( p_demux, "Couldn't initiate inflation ignore track %u",

                 tk.i_number );

        return 1;

    }

    d_stream.next_in = tk.p_extra_data;

    d_stream.avail_in = tk.i_extra_data;

    do

    {

        n++;

        void *alloc = realloc(p_new_extra, n*1024);

        if( alloc == NULL )

        {

            msg_Err( p_demux, "Couldn't allocate buffer to inflate data, ignore track %u",

                      tk.i_number );

            free(p_new_extra);

            inflateEnd( &d_stream );

            return 1;

        }

        p_new_extra = static_cast<uint8_t *>( alloc );

        d_stream.next_out = &p_new_extra[(n - 1) * 1024];

        d_stream.avail_out = 1024;

        result = inflate(&d_stream, Z_NO_FLUSH);

        if( result != Z_OK && result != Z_STREAM_END )

        {

            msg_Err( p_demux, "Zlib decompression failed. Result: %d", result );

            inflateEnd( &d_stream );

            free(p_new_extra);

            return 1;

        }

    }

    while ( d_stream.avail_out == 0 && d_stream.avail_in != 0  &&

            result != Z_STREAM_END );

    free( tk.p_extra_data );

    tk.i_extra_data = d_stream.total_out;

    p_new_extra = static_cast<uint8_t *>( realloc(p_new_extra, tk.i_extra_data) );

    if( !p_new_extra )

    {

        msg_Err( p_demux, "Couldn't allocate buffer to inflate data, ignore track %u",

                 tk.i_number );

        inflateEnd( &d_stream );

        tk.p_extra_data = NULL;

        return 1;

    }

    tk.p_extra_data = p_new_extra;

    inflateEnd( &d_stream );

    return 0;

}

block_t *block_zlib_decompress( vlc_object_t *p_this, block_t *p_in_block ) {

    int result, dstsize, n;

    unsigned char *dst;

    block_t *p_block;

    z_stream d_stream;

    d_stream.zalloc = NULL;

    d_stream.zfree = NULL;

    d_stream.opaque = NULL;

    result = inflateInit(&d_stream);

    if( result != Z_OK )

    {

        msg_Dbg( p_this, "inflateInit() failed. Result: %d", result );

        return NULL;

    }

    d_stream.next_in = (Bytef *)p_in_block->p_buffer;

    d_stream.avail_in = p_in_block->i_buffer;

    n = 0;

    p_block = block_Alloc( 0 );

    dst = NULL;

    do

    {

        n++;

        p_block = block_Realloc( p_block, 0, n * 1000 );

        dst = static_cast<unsigned char *>( p_block->p_buffer );

        d_stream.next_out = (Bytef *)&dst[(n - 1) * 1000];

        d_stream.avail_out = 1000;

        result = inflate(&d_stream, Z_NO_FLUSH);

        if( ( result != Z_OK ) && ( result != Z_STREAM_END ) )

        {

            msg_Err( p_this, "Zlib decompression failed. Result: %d", result );

            inflateEnd( &d_stream );

            block_Release( p_block );

            return p_in_block;

        }

    }

    while( ( d_stream.avail_out == 0 ) && ( d_stream.avail_in != 0 ) &&

           ( result != Z_STREAM_END ) );

    dstsize = d_stream.total_out;

    inflateEnd( &d_stream );

    p_block = block_Realloc( p_block, 0, dstsize );

    p_block->i_buffer = dstsize;

    block_Release( p_in_block );

    return p_block;

}

#endif

bool lzo1x_decompress_extra( demux_t * p_demux, mkv_track_t & tk )

{

    lzokay::EResult result;

    size_t n = 0, total_out = 0;

    uint8_t * p_new_extra = nullptr;

    msg_Dbg(p_demux,"Inflating private data");

    do

    {

        n++;

        void *alloc = realloc(p_new_extra, n * 1024);

        if( alloc == nullptr )

        {

            msg_Err( p_demux, "Couldn't allocate buffer to inflate data, ignore track %u",

                      tk.i_number );

            free(p_new_extra);

            return false;

        }

        p_new_extra = static_cast<uint8_t *>( alloc );

        result = lzokay::decompress( tk.p_extra_data, tk.i_extra_data,

                                     p_new_extra, n * 1024, total_out );

    }

    while ( result == lzokay::EResult::OutputOverrun );

    if( result != lzokay::EResult::Success )

    {

        msg_Err( p_demux, "LZO1X private data decompression failed. Result: %d", (int)result );

        free(p_new_extra);

        return false;

    }

    free( tk.p_extra_data );

    tk.i_extra_data = total_out;

    tk.p_extra_data = p_new_extra;

    return true;

}

block_t *block_lzo1x_decompress( vlc_object_t *p_this, block_t *p_in_block ) {

    lzokay::EResult result;

    size_t dstsize = 0, n;

    block_t *p_block;

    n = 0;

    p_block = block_Alloc( 0 );

    do

    {

        n++;

        p_block = block_Realloc( p_block, 0, n * 1000 );

        result = lzokay::decompress( p_in_block->p_buffer, p_in_block->i_buffer,

                                     p_block->p_buffer, p_block->i_buffer, dstsize );

    }

    while( result == lzokay::EResult::OutputOverrun );

    if( result != lzokay::EResult::Success )

    {

        msg_Err( p_this, "LZO1X decompression failed. Result: %d", (int)result );

        block_Release( p_block );

        return p_in_block;

    }

    p_block = block_Realloc( p_block, 0, dstsize );

    p_block->i_buffer = dstsize;

    block_Release( p_in_block );

    return p_block;

}

/* Utility function for BlockDecode */

block_t *MemToBlock( uint8_t *p_mem, size_t i_mem, size_t offset)

{

    if( unlikely( i_mem > SIZE_MAX - offset ) )

        return NULL;

    block_t *p_block = block_Alloc( i_mem + offset );

    if( likely(p_block != NULL) )

    {

        memcpy( p_block->p_buffer + offset, p_mem, i_mem );

    }

    return p_block;

}

void handle_real_audio(demux_t * p_demux, mkv_track_t * p_tk, block_t * p_blk, vlc_tick_t i_pts)

{

    uint8_t * p_frame = p_blk->p_buffer;

    Cook_PrivateTrackData * p_sys = static_cast<Cook_PrivateTrackData *>(p_tk->p_sys);

    size_t size = p_blk->i_buffer;

    if( p_tk->i_last_dts == VLC_TICK_INVALID )

    {

        for( size_t i = 0; i < p_sys->p_subpackets.size(); i++)

            if( p_sys->p_subpackets[i] )

            {

                block_Release(p_sys->p_subpackets[i]);

                p_sys->p_subpackets[i] = NULL;

            }

        p_sys->i_subpacket = 0;

        if ( !( p_blk->i_flags & BLOCK_FLAG_TYPE_I) )

        {

            msg_Dbg( p_demux, "discard non-key preroll block in track %u at %" PRId64,

                     p_tk->i_number, i_pts );

            return;

        }

    }

    if( p_tk->fmt.i_codec == VLC_CODEC_COOK ||

        p_tk->fmt.i_codec == VLC_CODEC_ATRAC3 )

    {

        const uint16_t i_num = p_sys->i_frame_size / p_sys->i_subpacket_size;

        if ( i_num == 0 )

            return;

        const size_t y = p_sys->i_subpacket / i_num;

        for( uint16_t i = 0; i < i_num; i++ )

        {

            size_t i_index = (size_t) p_sys->i_sub_packet_h * i +

                          ((p_sys->i_sub_packet_h + 1) / 2) * (y&1) + (y>>1);

            if( unlikely(i_index >= p_sys->p_subpackets.size()) )

                return;

            if( size < p_sys->i_subpacket_size )

                return;

            if (likely(p_sys->p_subpackets[i_index] == nullptr))

            {

                // the index was not used

                block_t *p_block = block_Alloc( p_sys->i_subpacket_size );

                if( !p_block )

                    return;

                memcpy( p_block->p_buffer, p_frame, p_sys->i_subpacket_size );

                p_block->i_dts = VLC_TICK_INVALID;

                p_block->i_pts = VLC_TICK_INVALID;

                if( p_sys->i_subpacket == 0 )

                {

                    p_tk->i_last_dts =

                    p_block->i_pts = i_pts;

                }

                p_sys->p_subpackets[i_index] = p_block;

            }

            p_frame += p_sys->i_subpacket_size;

            size -=  p_sys->i_subpacket_size;

            p_sys->i_subpacket++;

        }

    }

    else

    {

        /*TODO*/

    }

    if( p_sys->i_subpacket == p_sys->p_subpackets.size() )

    {

        for( size_t i = 0; i < p_sys->p_subpackets.size(); i++)

        {

            if (likely(p_sys->p_subpackets[i]))

            {

                send_Block( p_demux, p_tk, p_sys->p_subpackets[i], 1, 0 );

                p_sys->p_subpackets[i] = NULL;

            }

        }

        p_sys->i_subpacket = 0;

    }

}

block_t *WEBVTT_Repack_Sample(block_t *p_block, bool b_webm,

                              const uint8_t *p_add, size_t i_add)

{

    struct webvtt_cueelements_s els;

    memset(&els, 0, sizeof(els));

    size_t newsize = 0;

    block_t *newblock = nullptr;

    /* Repack to ISOBMFF samples format */

    if( !b_webm ) /* S_TEXT/WEBVTT */

    {

        /* process addition fields */

        if( i_add )

        {

            const uint8_t *end = p_add + i_add;

            const uint8_t *iden =

                    reinterpret_cast<const uint8_t *>(std::memchr( p_add, '\n', i_add ));

            if( iden && ++iden != end )

            {

                els.sttg.p_data = p_add;

                els.sttg.i_data = &iden[-1] - p_add;

                const uint8_t *comm =

                        reinterpret_cast<const uint8_t *>(std::memchr( iden, '\n', end - iden ));

                els.iden.p_data = iden;

                if( comm )

                    els.iden.i_data = comm - iden;

                else

                    els.iden.i_data = end - iden;

            }

        }

        /* the payload being in the block */

        els.payl.p_data = p_block->p_buffer;

        els.payl.i_data = p_block->i_buffer;

    }

    else /* deprecated D_WEBVTT/ */

    {

        const uint8_t *start = p_block->p_buffer;

        const uint8_t *end = p_block->p_buffer + p_block->i_buffer;

        const uint8_t *sttg =

                reinterpret_cast<const uint8_t *>(std::memchr( start, '\n', p_block->i_buffer ));

        if( !sttg || ++sttg == end )

            goto error;

        const uint8_t *payl =

                reinterpret_cast<const uint8_t *>(std::memchr( sttg, '\n', end - sttg ));

        if( !payl || ++payl == end )

            goto error;

        els.iden.p_data = start;

        els.iden.i_data = &sttg[-1] - start;

        els.sttg.p_data = sttg;

        els.sttg.i_data = &payl[-1] - sttg;

        els.payl.p_data = payl;

        els.payl.i_data = end - payl;

    }

    newsize = WEBVTT_Pack_CueElementsGetNewSize( &els );

    newblock = block_Alloc( newsize );

    if( !newblock )

        goto error;

    WEBVTT_Pack_CueElements( &els, newblock->p_buffer );

    block_CopyProperties( newblock, p_block );

    block_Release( p_block );

    return newblock;

error:

    block_Release( p_block );

    return NULL;

}

int UpdatePCR( demux_t * p_demux )

{

    demux_sys_t *p_sys = (demux_sys_t *)p_demux->p_sys;

    matroska_segment_c *p_segment = p_sys->GetCurrentVSegment()->CurrentSegment();

    vlc_tick_t i_pcr = VLC_TICK_INVALID;

    for( const auto & it : p_segment->tracks )

    {

        const auto &track = it.second;

        if( track->i_last_dts == VLC_TICK_INVALID )

            continue;

        if( track->fmt.i_cat != VIDEO_ES && track->fmt.i_cat != AUDIO_ES )

            continue;

        if( track->i_last_dts < i_pcr || i_pcr == VLC_TICK_INVALID )

        {

            i_pcr = track->i_last_dts;

        }

    }

    if( i_pcr != VLC_TICK_INVALID && i_pcr > p_sys->i_pcr )

    {

        if( es_out_SetPCR( p_demux->out, i_pcr ) )

        {

            msg_Err( p_demux, "ES_OUT_SET_PCR failed, aborting." );

            return VLC_EGENERIC;

        }

        p_sys->i_pcr = i_pcr;

    }

    return VLC_SUCCESS;

}

void send_Block( demux_t * p_demux, mkv_track_t * p_tk, block_t * p_block, unsigned int i_number_frames, uint64_t i_duration )

{

    demux_sys_t *p_sys = (demux_sys_t *)p_demux->p_sys;

    matroska_segment_c *p_segment = p_sys->GetCurrentVSegment()->CurrentSegment();

    if( p_tk->fmt.i_cat == AUDIO_ES && p_tk->i_chans_to_reorder )

    {

        aout_ChannelReorder( p_block->p_buffer, p_block->i_buffer,

                             p_tk->i_chans_to_reorder,

                             p_tk->pi_chan_table, p_tk->fmt.i_codec );

    }

    if( p_block->i_dts != VLC_TICK_INVALID &&

        ( p_tk->fmt.i_cat == VIDEO_ES || p_tk->fmt.i_cat == AUDIO_ES ) )

    {

        p_tk->i_last_dts = p_block->i_dts;

    }

    if( !p_tk->b_no_duration )

    {

        p_block->i_length = VLC_TICK_FROM_NS(i_duration * p_tk->f_timecodescale *

                                             p_segment->i_timescale) / i_number_frames;

    }

    if( p_tk->b_discontinuity )

    {

        p_block->i_flags |= BLOCK_FLAG_DISCONTINUITY;

        p_tk->b_discontinuity = false;

    }

    if ( p_sys->i_pcr == VLC_TICK_INVALID )

        UpdatePCR( p_demux );

    es_out_Send( p_demux->out, p_tk->p_es, p_block);

}

struct real_audio_private

{

    uint32_t fourcc;

    uint16_t version;

    uint16_t unknown1;

    uint8_t  unknown2[12];

    uint16_t unknown3;

    uint16_t flavor;

    uint32_t coded_frame_size;

    uint32_t unknown4[3];

    uint16_t sub_packet_h;

    uint16_t frame_size;

    uint16_t sub_packet_size;

    uint16_t unknown5;

};

struct real_audio_private_v4

{

    real_audio_private header;

    uint16_t sample_rate;

    uint16_t unknown;

    uint16_t sample_size;

    uint16_t channels;

};

struct real_audio_private_v5

{

    real_audio_private header;

    uint32_t unknown1;

    uint16_t unknown2;

    uint16_t sample_rate;

    uint16_t unknown3;

    uint16_t sample_size;

    uint16_t channels;

};

bool Cook_PrivateTrackData::Init()

{

    // real_audio_private

    bytes.skip(4);                      // fourcc

    /* FIXME RALF and SIPR */

    uint16_t version = bytes.GetBE16(); // version

    bytes.skip(2);                      // unknown1

    bytes.skip(12);                     // unknown2

    bytes.skip(2);                      // unknown3

    bytes.skip(2);                      // flavor

    coded_frame_size = bytes.GetBE32(); // coded_frame_size

    bytes.skip(3*4);                    // unknown4

    i_sub_packet_h   = bytes.GetBE16(); // sub_packet_h

    i_frame_size     = bytes.GetBE16(); // frame_size

    i_subpacket_size = bytes.GetBE16(); // sub_packet_size

    bytes.skip(2);                      // unknown5

    if ( i_subpacket_size == 0 )

        return false;

    if( version == 4 )

    {

        // real_audio_private_v4

        i_rate          = bytes.GetBE16(); // sample_rate

        bytes.skip(2);                     // unknown

        i_bitspersample = bytes.GetBE16(); // sample_size

        i_channels      = bytes.GetBE16(); // channels

    }

    else if( version == 5 )

    {

        // real_audio_private_v5

        bytes.skip(4);                     // unknown1

        bytes.skip(2);                     // unknown2

        i_rate          = bytes.GetBE16(); // sample_rate

        bytes.skip(2);                     // unknown2

        i_bitspersample = bytes.GetBE16(); // sample_size

        i_channels      = bytes.GetBE16(); // channels

    }

    else

        return false;

    if (bytes.hasErrors())

        return false;

    size_t i_subpackets = (size_t) i_sub_packet_h * (size_t) i_frame_size / (size_t) i_subpacket_size;

    p_subpackets.resize(i_subpackets);

    return true;

}

Cook_PrivateTrackData::~Cook_PrivateTrackData()

{

    for( size_t i = 0; i < p_subpackets.size(); i++ )

        if( p_subpackets[i] )

            block_Release( p_subpackets[i] );

}

static inline void fill_wvpk_block(uint16_t version, uint32_t block_samples, uint32_t flags,

                                   uint32_t crc, uint8_t * src, size_t srclen, uint8_t * dst)

{

    const uint8_t wvpk_header[] = {'w','v','p','k',         /* ckId */

                                    0x0, 0x0, 0x0, 0x0,     /* ckSize */

                                    0x0, 0x0,               /* version */

                                    0x0,                    /* track_no */

                                    0x0,                    /* index_no */

                                    0xFF, 0xFF, 0xFF, 0xFF, /* total_samples */

                                    0x0, 0x0, 0x0, 0x0 };   /* block_index */

    memcpy( dst, wvpk_header, sizeof( wvpk_header ) );

    SetDWLE( dst + 4, srclen + 24 );

    SetWLE( dst + 8, version );

    SetDWLE( dst + 20, block_samples );

    SetDWLE( dst + 24, flags );

    SetDWLE( dst + 28, crc );

    memcpy( dst + 32, src, srclen );

}

block_t * packetize_wavpack( const mkv_track_t & tk, uint8_t * buffer, size_t  size)

{

    uint16_t version = 0x403;

    uint32_t block_samples;

    uint32_t flags;

    uint32_t crc;

    block_t * p_block = NULL;

    if( tk.i_extra_data >= 2 )

        version = GetWLE( tk.p_extra_data );

    if( size < 12 )

        return NULL;

    block_samples = GetDWLE(buffer);

    buffer += 4;

    flags = GetDWLE(buffer);

    size -= 4;

    /* Check if WV_INITIAL_BLOCK and WV_FINAL_BLOCK are present */

    if( ( flags & 0x1800 ) == 0x1800 )

    {

        crc = GetDWLE(buffer+4);

        buffer += 8;

        size -= 8;

        p_block = block_Alloc( size + 32 );

        if( !p_block )

            return NULL;

        fill_wvpk_block(version, block_samples, flags, crc, buffer, size, p_block->p_buffer);

    }

    else

    {

        /* Multiblock */

        size_t total_size = 0;

        p_block = block_Alloc( 0 );

        if( !p_block )

            return NULL;

        while(size >= 12)

        {

            flags = GetDWLE(buffer);

            buffer += 4;

            crc = GetDWLE(buffer);

            buffer += 4;

            uint32_t bsz = GetDWLE(buffer);

            buffer+= 4;

            size -= 12;

            bsz = (bsz < size)?bsz:size;

            total_size += bsz + 32;

            assert(total_size >= p_block->i_buffer);

            p_block = block_Realloc( p_block, 0, total_size );

            if( !p_block )

                return NULL;

            fill_wvpk_block(version, block_samples, flags, crc, buffer, bsz,

                            p_block->p_buffer + total_size - bsz - 32 );

            buffer += bsz;

            size -= bsz;

        }

    }

    return p_block;

}

void MkvTree_va( demux_t& demuxer, int i_level, const char* fmt, va_list args)

{

    static const char indent[] = "|   ";

    static const char prefix[] = "+ ";

    static int  const   indent_len = sizeof( indent ) - 1;

    static int  const   prefix_len = sizeof( prefix ) - 1;

    char   fixed_buffer[256] = {};

    size_t const  static_len = sizeof( fixed_buffer );

    char *            buffer = fixed_buffer;

    size_t         total_len = indent_len * i_level + prefix_len + strlen( fmt ) + 1;

    if( total_len >= static_len ) {

        buffer = new (std::nothrow) char[total_len] ();

        if (buffer == NULL) {

            msg_Err (&demuxer, "Unable to allocate memory for format string");

            return;

        }

    }

    char * dst = buffer;

    for (int i = 0; i < i_level; ++i, dst += indent_len)

        memcpy( dst, indent, indent_len );

    strcat( dst, prefix );

    strcat( dst, fmt );

    msg_GenericVa( &demuxer, VLC_MSG_DBG, buffer, args );

    if (buffer != fixed_buffer)

        delete [] buffer;

}

void MkvTree( demux_t & demuxer, int i_level, const char *psz_format, ... )

{

    va_list args; va_start( args, psz_format );

    MkvTree_va( demuxer, i_level, psz_format, args );

    va_end( args );

}

} // namespace