/*************************************************************************
    VorbisEncoder.cpp  -  sub encoder base class for Vorbis in an Ogg container
                             -------------------
    begin                : Thu Jan 03 2013
    copyright            : (C) 2013 by Thomas Eschenbacher
    email                : Thomas.Eschenbacher@gmx.de
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/

#include "config.h"

#include <math.h>
#include <string.h>

#include <QTime>
#include <QtGlobal>

#include <KLocalizedString>

#include "libkwave/MessageBox.h"
#include "libkwave/MultiTrackReader.h"
#include "libkwave/Sample.h"
#include "libkwave/SampleArray.h"
#include "libkwave/Utils.h"

#include "VorbisEncoder.h"

/** size of a buffer used for Vorbis encoding */
#define BUFFER_SIZE 1024

/** bitrate to be used when no bitrate has been selected */
#define DEFAULT_BITRATE 64000

/***************************************************************************/
Kwave::VorbisEncoder::VorbisEncoder()
    :m_comments_map(), m_info()
{
    memset(&m_os, 0x00, sizeof(m_os));
    memset(&m_og, 0x00, sizeof(m_og));
    memset(&m_op, 0x00, sizeof(m_op));

    memset(&m_vb, 0x00, sizeof(m_vb));
    memset(&m_vc, 0x00, sizeof(m_vc));
    memset(&m_vd, 0x00, sizeof(m_vd));
    memset(&m_vi, 0x00, sizeof(m_vi));
}

/***************************************************************************/
Kwave::VorbisEncoder::~VorbisEncoder()
{
    close();
}

/***************************************************************************/
void Kwave::VorbisEncoder::encodeProperties(const Kwave::FileInfo &info)
{
    for (VorbisCommentMap::const_iterator it(m_comments_map.constBegin());
         it != m_comments_map.constEnd(); ++it)
    {
	const QString       &key     = it.key();
	Kwave::FileProperty property = it.value();
	if (!info.contains(property)) continue; // skip if not present

	// encode the property as string
	vorbis_comment_add_tag(&m_vc,
	    UTF8(key),
	    UTF8(info.get(property).toString())
	);
    }
}

/***************************************************************************/
bool Kwave::VorbisEncoder::open(QWidget *widget, const Kwave::FileInfo &info,
                                Kwave::MultiTrackReader &src)
{
    int ret = -1;

    Q_UNUSED(src)

    // get info: tracks, sample rate, bitrate(s)
    m_info = info;
    const unsigned int tracks = info.tracks();
    long int sample_rate = static_cast<long int>(info.rate());

    if (tracks > 2) {
	Kwave::MessageBox::sorry(widget,
	    i18n("This codec supports only mono or stereo files, "
	         "%1 channels are not supported.", tracks));
	return false;
    }

    // ABR bitrates
    int bitrate_nominal = info.contains(Kwave::INF_BITRATE_NOMINAL) ?
        QVariant(info.get(Kwave::INF_BITRATE_NOMINAL)).toInt() : -1;
    int bitrate_lower = info.contains(Kwave::INF_BITRATE_LOWER) ?
        QVariant(info.get(Kwave::INF_BITRATE_LOWER)).toInt() : -1;
    int bitrate_upper = info.contains(Kwave::INF_BITRATE_UPPER) ?
        QVariant(info.get(Kwave::INF_BITRATE_UPPER)).toInt() : -1;

    // VBR quality
    int vbr_quality = info.contains(Kwave::INF_VBR_QUALITY) ?
        QVariant(info.get(Kwave::INF_VBR_QUALITY)).toInt() : -1;

    qDebug("OggEncoder: ABR=%d...%d...%d Bits/s, VBR=%d%%",
           bitrate_lower,bitrate_nominal,bitrate_upper,vbr_quality);

    if ((vbr_quality < 0) && (bitrate_nominal <= 0)) {
	// no quality and no bitrate given -> complain !
	if (Kwave::MessageBox::warningContinueCancel(widget,
	    i18n("You have not selected any bitrate for the encoding. "
	         "Do you want to continue and encode with %1 kBit/s "
	         "or cancel and choose a different bitrate?",
	         DEFAULT_BITRATE / 1000)) != KMessageBox::Continue)
	    return false; // <- canceled

	bitrate_nominal = DEFAULT_BITRATE;
	bitrate_lower = -1;
	bitrate_upper = -1;
    }

    // some checks first
    Q_ASSERT(tracks < 255);
    if (tracks > 255) return false;

    /********** Encode setup ************/
    vorbis_info_init(&m_vi);

    if ((bitrate_lower > 0) || (bitrate_upper > 0)) {
	// Encoding using ABR mode.
	bitrate_nominal = (bitrate_upper + bitrate_lower) / 2;
	ret = vorbis_encode_init(&m_vi, tracks, sample_rate,
	                         bitrate_upper,
	                         bitrate_nominal,
	                         bitrate_lower);
	qDebug("VorbisEncoder: ABR with %d...%d...%d Bits/s",
	       bitrate_lower, bitrate_nominal, bitrate_upper);
    } else if ((vbr_quality < 0) && (bitrate_nominal > 0)) {
	// Encoding using constant bitrate in ABR mode
	ret = vorbis_encode_setup_managed(&m_vi, tracks, sample_rate,
	      -1, bitrate_nominal, -1);

	// If OV_ECTL_RATEMANAGE2_SET is not defined, then your
	// libvorbis is just too old.
	if (!ret) ret =
	      vorbis_encode_ctl(&m_vi, OV_ECTL_RATEMANAGE2_SET, Q_NULLPTR) ||
	      vorbis_encode_setup_init(&m_vi);

	qDebug("VorbisEncoder: CBR with %d Bits/s", bitrate_nominal);
    } else if (vbr_quality >= 0) {
	// Encoding using VBR mode.
	ret = vorbis_encode_init_vbr(&m_vi, tracks, sample_rate,
	    static_cast<float>(vbr_quality) / 100.0f);
	qDebug("OggEncoder: VBR with %d%%", vbr_quality);
    } else {
	// unknown setup !?
	qWarning("unknown Ogg/Vorbis setup: VBR quality=%d%%, "
	         "ABR lower=%d, ABR highest=%d, ABR nominal=%d",
	         vbr_quality, bitrate_lower, bitrate_upper,
	         bitrate_nominal);
	return false;
    }

    /* do not continue if setup failed; this can happen if we ask for a
       mode that libVorbis does not support (eg, too low a bitrate, etc,
       will return 'OV_EIMPL') */
    if (ret) {
	Kwave::MessageBox::sorry(widget, i18n("One or more encoding "
	    "parameters are not supported. Please change the "
	    "settings and try again."));
	return false;
    }

    // add all supported properties as file comments
    vorbis_comment_init(&m_vc);
    encodeProperties(info);

    // set up the analysis state and auxiliary encoding storage
    vorbis_analysis_init(&m_vd, &m_vi);
    vorbis_block_init(&m_vd, &m_vb);


    // set up our packet->stream encoder
    // pick a random serial number; that way we can more likely build
    // chained streams just by concatenation
    qsrand(QTime::currentTime().msec() ^ qrand());
    ogg_stream_init(&m_os, qrand());

    return true;
}

/***************************************************************************/
bool Kwave::VorbisEncoder::writeHeader(QIODevice &dst)
{
    // Vorbis streams begin with three headers; the initial header (with
    // most of the codec setup parameters) which is mandated by the Ogg
    // bitstream spec.  The second header holds any comment fields.  The
    // third header holds the bitstream codebook.  We merely need to
    // make the headers, then pass them to libvorbis one at a time;
    // libvorbis handles the additional Ogg bitstream constraints
    ogg_packet header;
    ogg_packet header_comm;
    ogg_packet header_code;

    vorbis_analysis_headerout(&m_vd, &m_vc, &header, &header_comm,
                              &header_code);
    // automatically placed in its own page
    ogg_stream_packetin(&m_os, &header);
    ogg_stream_packetin(&m_os, &header_comm);
    ogg_stream_packetin(&m_os, &header_code);

    // This ensures the actual audio data will start on a
    // new page, as per spec
    while (ogg_stream_flush(&m_os, &m_og)) {
	dst.write(reinterpret_cast<char *>(m_og.header),
	          m_og.header_len);
	dst.write(reinterpret_cast<char *>(m_og.body),
	          m_og.body_len);
    }

    return true;
}

/***************************************************************************/
bool Kwave::VorbisEncoder::encode(Kwave::MultiTrackReader &src,
                                  QIODevice &dst)
{
    bool                 eos    = false;
    const unsigned int   tracks = m_info.tracks();
    const sample_index_t length = m_info.length();

    sample_index_t rest = length;
    while (!eos && !src.isCanceled()) {
	if (src.eof()) {
	    // end of file.  this can be done implicitly in the mainline,
	    // but it's easier to see here in non-clever fashion.
	    // Tell the library we're at end of stream so that it can handle
	    // the last frame and mark end of stream in the output properly
	    vorbis_analysis_wrote(&m_vd, 0);
	} else {
	    // data to encode

	    // expose the buffer to submit data
	    float **buffer = vorbis_analysis_buffer(&m_vd, BUFFER_SIZE);
	    unsigned int pos = 0;
	    unsigned int len = (rest > BUFFER_SIZE) ? BUFFER_SIZE :
		                                      Kwave::toUint(rest);
	    Kwave::SampleArray samples(BUFFER_SIZE);
	    for (unsigned int track = 0; track < tracks; ++track) {
		float *p = buffer[track];
		unsigned int l = src[track]->read(samples, 0, len);
		const sample_t *s = samples.constData();

		const unsigned int block = 8;
		pos = 0;
		while (pos + block < l) {
		    for (unsigned int i = 0; i < block; ++i, ++pos)
			p[pos] = sample2float(s[pos]);
		}
		while (pos < l) {
		    p[pos] = sample2float(s[pos]);
		    pos++;
		}
		while (pos < len)
		    p[pos++] = 0;
	    }

	    // tell the library how much we actually submitted
	    vorbis_analysis_wrote(&m_vd, pos);
	}

	// vorbis does some data preanalysis, then divvies up blocks for
	// more involved (potentially parallel) processing.  Get a single
	// block for encoding now
	while (vorbis_analysis_blockout(&m_vd, &m_vb) == 1) {
	    // analysis, assume we want to use bitrate management
            vorbis_analysis(&m_vb, Q_NULLPTR);
	    vorbis_bitrate_addblock(&m_vb);

	    while (vorbis_bitrate_flushpacket(&m_vd, &m_op)) {
		// weld the packet into the bitstream
		ogg_stream_packetin(&m_os, &m_op);

		// write out pages (if any)
		while (!eos) {
		    int result = ogg_stream_pageout(&m_os, &m_og);
		    if (!result) break;
		    dst.write(reinterpret_cast<char*>(m_og.header),
		              m_og.header_len);
		    dst.write(reinterpret_cast<char *>(m_og.body),
		              m_og.body_len);

		    // this could be set above, but for illustrative
		    // purposes, I do it here (to show that vorbis
		    // does know where the stream ends)
		    if (ogg_page_eos(&m_og)) eos = true;
		}
	    }
	}
    }

    return true;
}

/***************************************************************************/
void Kwave::VorbisEncoder::close()
{
    ogg_stream_clear(&m_os);

    vorbis_block_clear(&m_vb);
    vorbis_dsp_clear(&m_vd);
    vorbis_comment_clear(&m_vc);
    vorbis_info_clear(&m_vi);   // <- must be called last
}

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