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

 * webvtt.c: muxer for raw WEBVTT

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

 * Copyright (C) 2024 VideoLabs, VLC authors and VideoLAN

 *

 * 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.

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

#ifdef HAVE_CONFIG_H

#include "config.h"

#endif

#include <vlc_common.h>

#include <vlc_codec.h>

#include <vlc_memstream.h>

#include <vlc_plugin.h>

#include <vlc_sout.h>

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

#include "../demux/mp4/minibox.h"

typedef struct

{

    bool header_done;

    const sout_input_t *input;

} sout_mux_sys_t;

static void OutputTime(struct vlc_memstream *ms, vlc_tick_t time)

{

    const vlc_tick_t secs = SEC_FROM_VLC_TICK(time);

    vlc_memstream_printf(ms,

                         "%02" PRIi64 ":%02" PRIi64 ":%02" PRIi64 ".%03" PRIi64,

                         secs / 3600,

                         secs % 3600 / 60,

                         secs % 60,

                         (time % CLOCK_FREQ) / 1000);

}

static block_t *FormatCue(const webvtt_cue_t *cue)

{

    struct vlc_memstream ms;

    if (vlc_memstream_open(&ms))

        return NULL;

    if (cue->psz_id != NULL)

        vlc_memstream_printf(&ms, "%s\n", cue->psz_id);

    OutputTime(&ms, cue->i_start);

    vlc_memstream_printf(&ms, " --> ");

    OutputTime(&ms, cue->i_stop);

    if (cue->psz_attrs != NULL)

        vlc_memstream_printf(&ms, " %s\n", cue->psz_attrs);

    else

        vlc_memstream_putc(&ms, '\n');

    vlc_memstream_printf(&ms, "%s\n\n", cue->psz_text);

    if (vlc_memstream_close(&ms))

        return NULL;

    block_t *formatted = block_heap_Alloc(ms.ptr, ms.length);

    formatted->i_length = cue->i_stop - cue->i_start;

    formatted->i_pts = formatted->i_dts = cue->i_stop - cue->i_start;

    return formatted;

}

static void UnpackISOBMFF(const vlc_frame_t *packed, webvtt_cue_t *out)

{

    mp4_box_iterator_t it;

    mp4_box_iterator_Init(&it, packed->p_buffer, packed->i_buffer);

    while (mp4_box_iterator_Next(&it))

    {

        if (it.i_type != ATOM_vttc && it.i_type != ATOM_vttx)

            continue;

        mp4_box_iterator_t vtcc;

        mp4_box_iterator_Init(&vtcc, it.p_payload, it.i_payload);

        while (mp4_box_iterator_Next(&vtcc))

        {

            switch (vtcc.i_type)

            {

                case ATOM_iden:

                    if (out->psz_id == NULL)

                        out->psz_id =

                            strndup((char *)vtcc.p_payload, vtcc.i_payload);

                    break;

                case ATOM_sttg:

                    if (out->psz_attrs)

                    {

                        char *dup = strndup((char *)vtcc.p_payload, vtcc.i_payload);

                        if (dup)

                        {

                            char *psz;

                            if (asprintf(&psz, "%s %s", out->psz_attrs, dup) >= 0)

                            {

                                free(out->psz_attrs);

                                out->psz_attrs = psz;

                            }

                            free(dup);

                        }

                    }

                    else

                        out->psz_attrs =

                            strndup((char *)vtcc.p_payload, vtcc.i_payload);

                    break;

                case ATOM_payl:

                    if (!out->psz_text)

                        out->psz_text =

                            strndup((char *)vtcc.p_payload, vtcc.i_payload);

                    break;

            }

        }

    }

}

static int Control(sout_mux_t *mux, int query, va_list args)

{

    switch (query)

    {

        case MUX_CAN_ADD_STREAM_WHILE_MUXING:

            *(va_arg(args, bool *)) = false;

            return VLC_SUCCESS;

        default:

            return VLC_ENOTSUP;

    }

    (void)mux;

}

static int AddStream(sout_mux_t *mux, sout_input_t *input)

{

    sout_mux_sys_t *sys = mux->p_sys;

    if ((input->fmt.i_codec != VLC_CODEC_WEBVTT &&

         input->fmt.i_codec != VLC_CODEC_TEXT) ||

        sys->input)

        return VLC_ENOTSUP;

    sys->input = input;

    return VLC_SUCCESS;

}

static void DelStream(sout_mux_t *mux, sout_input_t *input)

{

    sout_mux_sys_t *sys = mux->p_sys;

    if (input == sys->input)

        sys->input = NULL;

}

static int Mux(sout_mux_t *mux)

{

    sout_mux_sys_t *sys = mux->p_sys;

    if (mux->i_nb_inputs == 0)

        return VLC_SUCCESS;

    sout_input_t *input = mux->pp_inputs[0];

    if (!sys->header_done)

    {

        block_t *data = NULL;

        if (input->fmt.i_extra > 8 && !memcmp(input->fmt.p_extra, "WEBVTT", 6))

        {

            data = block_Alloc(input->fmt.i_extra + 2);

            if (unlikely(data == NULL))

                return VLC_ENOMEM;

            memcpy(data->p_buffer, input->fmt.p_extra, input->fmt.i_extra);

            data->p_buffer[data->i_buffer - 2] = '\n';

            data->p_buffer[data->i_buffer - 1] = '\n';

        }

        else

        {

            data = block_Alloc(8);

            if (unlikely(data == NULL))

                return VLC_ENOMEM;

            memcpy(data->p_buffer, "WEBVTT\n\n", 8);

        }

        if (data)

        {

            data->i_flags |= BLOCK_FLAG_HEADER;

            sout_AccessOutWrite(mux->p_access, data);

        }

        sys->header_done = true;

    }

    vlc_fifo_Lock(input->p_fifo);

    vlc_frame_t *chain = vlc_fifo_DequeueAllUnlocked(input->p_fifo);

    int status = VLC_SUCCESS;

    while (chain != NULL)

    {

        // Ephemer subtitles stop being displayed at the next SPU display time.

        // We need to delay if subsequent SPU aren't available yet.

        const bool is_ephemer = (chain->i_length == VLC_TICK_INVALID);

        if (is_ephemer && chain->p_next == NULL)

        {

            vlc_fifo_QueueUnlocked(input->p_fifo, chain);

            chain = NULL;

            break;

        }

        webvtt_cue_t cue = {};

        if (sys->input->fmt.i_codec != VLC_CODEC_WEBVTT)

            cue.psz_text = strndup((char *)chain->p_buffer, chain->i_buffer);

        else

            UnpackISOBMFF(chain, &cue);

        if (unlikely(cue.psz_text == NULL))

        {

            webvtt_cue_Clean(&cue);

            status = VLC_ENOMEM;

            break;

        }

        cue.i_start = chain->i_pts - VLC_TICK_0;

        if(is_ephemer)

            cue.i_stop = chain->p_next->i_pts - VLC_TICK_0;

        else

            cue.i_stop = cue.i_start + chain->i_length;

        block_t *to_write = FormatCue(&cue);

        webvtt_cue_Clean(&cue);

        ssize_t written = -1;

        if (likely(to_write != NULL))

            written = sout_AccessOutWrite(mux->p_access, to_write);

        if (written == -1)

        {

            status = VLC_EGENERIC;

            break;

        }

        vlc_frame_t *next = chain->p_next;

        vlc_frame_Release(chain);

        chain = next;

    }

    vlc_frame_ChainRelease(chain);

    vlc_fifo_Unlock(input->p_fifo);

    return status;

}

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

 * Open:

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

int webvtt_OpenMuxer(vlc_object_t *this)

{

    sout_mux_t *mux = (sout_mux_t *)this;

    sout_mux_sys_t *sys;

    mux->p_sys = sys = malloc(sizeof(*sys));

    if (unlikely(sys == NULL))

        return VLC_ENOMEM;

    sys->header_done = false;

    sys->input = NULL;

    mux->pf_control = Control;

    mux->pf_addstream = AddStream;

    mux->pf_delstream = DelStream;

    mux->pf_mux = Mux;

    return VLC_SUCCESS;

}

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

 * Close:

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

void webvtt_CloseMuxer(vlc_object_t *this)

{

    sout_mux_t *mux = (sout_mux_t *)this;

    free(mux->p_sys);

}