//----------------------------------------------------------------------------
//
// TSDuck - The MPEG Transport Stream Toolkit
// Copyright (c) 2005-2025, Thierry Lelegard
// BSD-2-Clause license, see LICENSE.txt file or https://tsduck.io/license
//
//----------------------------------------------------------------------------
//
//  Add stuffing to a TS file to reach a target bitrate.
//
//----------------------------------------------------------------------------

#include "tsMain.h"
#include "tsTSFileInputBuffered.h"
TS_MAIN(MainCode);


//-----------------------------------------------------------------------------
//  Command line options
//-----------------------------------------------------------------------------

static constexpr size_t   MIN_TS_BUFFER_SIZE     = 1024;             // 1 kB
static constexpr size_t   DEFAULT_TS_BUFFER_SIZE = 4 * 1024 * 1024;  // 4 MB
static constexpr size_t   MAX_TS_BUFFER_SIZE     = 16 * 1024 * 1024; // 16 MB
static constexpr uint64_t DEFAULT_MIN_INTERVAL   = 100;              // milliseconds

namespace {
    class Options: public ts::Args
    {
        TS_NOBUILD_NOCOPY(Options);
    public:
        Options(int argc, char *argv[]);

        ts::BitRate target_bitrate = 0;
        ts::PID     reference_pid = ts::PID_NULL;
        size_t      buffer_size = 0;
        uint64_t    leading_packets = 0;
        uint64_t    trailing_packets = 0;
        uint64_t    final_inter_packet = 0;
        uint64_t    initial_inter_packet = 0;
        uint64_t    min_interval_ms = 0;
        bool        dts_based = false;
        bool        dyn_final_inter_packet = false;
        bool        dyn_initial_inter_packet = false;
        ts::UString input_file {};
        ts::UString output_file {};
        ts::TSPacketFormat input_format = ts::TSPacketFormat::AUTODETECT;
        ts::TSPacketFormat output_format = ts::TSPacketFormat::TS;
    };
}

Options::Options(int argc, char *argv[]) :
    Args(u"Add stuffing to a transport stream to reach a target bitrate", u"[options] [input-file]")
{
    option(u"", 0, FILENAME, 0, 1);
    help(u"",
         u"The input file is a TS file, typically with variable bitrate content. "
         u"By default, the standard input is used.");

    option<ts::BitRate>(u"bitrate", 'b', 1, 1, 1);
    help(u"bitrate",
         u"Target constant bitrate of the output file. "
         u"This is mandatory parameter, there is no default.");

    option(u"buffer-size", 0, INTEGER, 0, 1, MIN_TS_BUFFER_SIZE, MAX_TS_BUFFER_SIZE);
    help(u"buffer-size",
         u"Input buffer size, in bytes. Must be large enough to always contain two "
         u"time stamps in the reference PID. Default: " + ts::UString::Decimal(DEFAULT_TS_BUFFER_SIZE) + u" bytes.");

    option(u"dts-based", 'd');
    help(u"dts-based",
         u"Use Decoding Time Stamps (DTS) in the reference PID to evaluate the "
         u"amount of stuffing to insert. The default is to use Program Clock "
         u"References (PCR) instead of DTS.");

    option(u"final-inter-packet", 'f', UNSIGNED);
    help(u"final-inter-packet",
         u"Number of stuffing packets to add between input packets after the last "
         u"time stamp (PCR or DTS). By default, use the same number as in the "
         u"previous segment, between the last two time stamps.");

    option(u"initial-inter-packet", 'i', UNSIGNED);
    help(u"initial-inter-packet",
         u"Number of stuffing packets to add between input packets before the first "
         u"time stamp (PCR or DTS). By default, use the same number as in the "
         u"first segment, between the first two time stamps.");

    option(u"leading-packets", 'l', UNSIGNED);
    help(u"leading-packets",
         u"Number of consecutive stuffing packets to add at the beginning of the "
         u"output file, before the first input packet. The default is zero.");

    option(u"min-interval", 'm', POSITIVE);
    help(u"min-interval",
         u"Minimum interval, in milli-seconds, between two recomputations of the "
         u"amount of stuffing to insert. This duration is based on time-stamps, "
         u"not real time. The default is " + ts::UString::Decimal(DEFAULT_MIN_INTERVAL) + u" ms.");

    option(u"output-file", 'o', FILENAME);
    help(u"output-file", u"filename",
         u"Output file name (standard output by default). The output file is a TS "
         u"file with the same packets as the input file with interspersed stuffing "
         u"packets and a constant bitrate.");

    option(u"reference-pid", 'r', PIDVAL);
    help(u"reference-pid",
         u"PID in which to collect time stamps (PCR or DTS) to use as reference "
         u"for the insertion of stuffing packets. By default, use the first PID "
         u"containing the specified type of time stamps (PCR or DTS).");

    option(u"trailing-packets", 't', UNSIGNED);
    help(u"trailing-packets",
         u"Number of consecutive stuffing packets to add at the end of the "
         u"output file, after the last input packet. The default is zero.");

    ts::DefineTSPacketFormatInputOption(*this, 0, u"input-format");
    ts::DefineTSPacketFormatOutputOption(*this, 0, u"output-format");

    analyze(argc, argv);

    getValue(input_file, u"");
    getValue(output_file, u"output-file");

    input_format = ts::LoadTSPacketFormatInputOption(*this, u"input-format");
    output_format = ts::LoadTSPacketFormatOutputOption(*this, u"output-format");

    getValue(target_bitrate, u"bitrate", 0);
    assert(target_bitrate != 0);

    getIntValue(buffer_size, u"buffer-size", DEFAULT_TS_BUFFER_SIZE);
    dts_based = present(u"dts-based");
    getIntValue(reference_pid, u"reference-pid", ts::PID_NULL);
    getIntValue(final_inter_packet, u"final-inter-packet", 0);
    getIntValue(initial_inter_packet, u"initial-inter-packet", 0);
    getIntValue(leading_packets, u"leading-packets", 0);
    getIntValue(trailing_packets, u"trailing-packets", 0);
    getIntValue(min_interval_ms, u"min-interval", DEFAULT_MIN_INTERVAL);
    dyn_final_inter_packet = !present(u"final-inter-packet");
    dyn_initial_inter_packet = !present(u"initial-inter-packet");

    exitOnError();
}


//-----------------------------------------------------------------------------
// Define a time stamp in the input TS
//-----------------------------------------------------------------------------

struct TimeStamp
{
    uint64_t tstamp;  // Time stamp in PCR units
    uint64_t packet;  // Packet index in input file

    // Constructor
    TimeStamp(uint64_t t = 0, uint64_t p = 0) : tstamp(t), packet(p) {}
    TS_DEFAULT_COPY_MOVE(TimeStamp);
};


//-----------------------------------------------------------------------------
// This class processes the input file
//-----------------------------------------------------------------------------

class Stuffer
{
    TS_NOBUILD_NOCOPY(Stuffer);
public:
    // Constructor
    Stuffer(Options& opt) :
        _opt(opt),
        _input(opt.buffer_size / ts::PKT_SIZE)
    {
    }

    // Process the content
    void stuff();

private:
    // Private members
    Options&                 _opt;                         // Command-line options.
    ts::TSFileInputBuffered  _input;                       // Input file, including seek buffer for at least one segment.
    ts::TSFile               _output {};                   // Output file.
    std::optional<TimeStamp> _tstamp1 {};                  // First time-stamp in current segment.
    std::optional<TimeStamp> _tstamp2 {};                  // Second time-stamp in current segment.
    uint64_t                 _current_inter_packet = 0;    // Number of null packets to add between all input packets in segment.
    uint64_t                 _current_residue_packets = 0; // Packets to add to inter_packets.
    uint64_t                 _remaining_stuff_count = 0;   // Remaining number of stuffing packets to add before end of segment.
    uint64_t                 _additional_bits = 0;         // Additional bits (less than one packet) to add in next segment.

    // Abort processing (invoked on fatal error, when message already reported)
    [[noreturn]] void fatalError();

    // Get name of time stamps
    ts::UString getTimeStampType() const {return _opt.dts_based ? u"DTS" : u"PCR";}

    // Check if a packet contains a time stamp.
    bool getTimeStamp(const ts::TSPacket& pkt, uint64_t& tstamp) const;

    // Evaluate stuffing need in next segment, between two time stamps.
    // Compute _current_inter_packet, _remaining_stuff_count, _additional_bits.
    void evaluateNextStuffing();

    // Write the specified number of stuffing packets
    void writeStuffing(uint64_t stuffing_packet_count);

    // Read input up to end_packet and perform simple inter-packet stuffing.
    void simpleInterPacketStuffing(uint64_t inter_packet, uint64_t end_packet);
};


//-----------------------------------------------------------------------------
// Abort processing
//-----------------------------------------------------------------------------

void Stuffer::fatalError()
{
    // Maybe something more clever some day...
    std::exit(EXIT_FAILURE);
}


//-----------------------------------------------------------------------------
// Check if a packet contains a time stamp.
//-----------------------------------------------------------------------------

bool Stuffer::getTimeStamp(const ts::TSPacket& pkt, uint64_t& tstamp) const
{
    if (_opt.dts_based) {
        if (pkt.hasDTS()) {
            tstamp = pkt.getDTS() * ts::SYSTEM_CLOCK_SUBFACTOR;
            return true;
        }
    }
    else {
        if (pkt.hasPCR()) {
            tstamp = pkt.getPCR();
            return true;
        }
    }
    return false;
}


//-----------------------------------------------------------------------------
// Write the specified number of stuffing packets
//-----------------------------------------------------------------------------

void Stuffer::writeStuffing(uint64_t count)
{
    while (count > 0) {
        if (!_output.writePackets(&ts::NullPacket, nullptr, 1, _opt)) {
            fatalError();
        }
        count--;
    }
}


//-----------------------------------------------------------------------------
// Read input up to end_packet and perform simple inter-packet stuffing.
//-----------------------------------------------------------------------------

void Stuffer::simpleInterPacketStuffing(uint64_t inter_packet, uint64_t end_packet)
{
    assert(_input.readPacketsCount() < end_packet);

    ts::TSPacket pkt;
    while (_input.readPacketsCount() < end_packet && _input.read(&pkt, 1, _opt) == 1) {
        if (!_output.writePackets(&pkt, nullptr, 1, _opt)) {
            fatalError();
        }
        writeStuffing(inter_packet);
    }
}


//-----------------------------------------------------------------------------
// Evaluate stuffing need in next segment, between two time stamps.
//-----------------------------------------------------------------------------

void Stuffer::evaluateNextStuffing()
{
    // Save initial position in the file
    const ts::PacketCounter initial_position = _input.readPacketsCount();
    _opt.debug(u"evaluateNextStuffing: initial_position = %'d", initial_position);

    // Initialize new search. Note that _tstamp1 and _tstamp2 may be unset.
    _tstamp1 = _tstamp2;
    _tstamp2.reset();

    // Read packets until both _tstamp1 and _tstamp2 are set (or end of file)
    ts::TSPacket pkt;
    while (!_tstamp2 && _input.canSeek(initial_position) && _input.read(&pkt, 1, _opt) == 1) {
        uint64_t tstamp = 0;
        if (getTimeStamp(pkt, tstamp)) {
            if (_opt.reference_pid == ts::PID_NULL) {
                // Found the first time stamp, use this PID as reference
                _opt.reference_pid = pkt.getPID();
                _opt.verbose(u"using PID %n as reference", _opt.reference_pid);
            }
            else if (_opt.reference_pid != pkt.getPID()) {
                // Not the reference PID, skip;
                continue;
            }
            const TimeStamp time_stamp(tstamp, _input.readPacketsCount());
            if (!_tstamp1 || tstamp < _tstamp1->tstamp) {
                // 1) Found the first time stamp in the file.
                // 2) Or found a time stamp lower than tstamp1, may be because of a
                //    file rewind or wrapping at 2**42. Use new time stamp as first.
                _tstamp1 = time_stamp;
                _opt.debug(u"evaluateNextStuffing: tstamp1 = %'d at %'d", time_stamp.tstamp, time_stamp.packet);
            }
            else if (((tstamp - _tstamp1->tstamp) * 1000) / ts::SYSTEM_CLOCK_FREQ >= _opt.min_interval_ms) {
                // Found second time stamp (with a sufficiently large interval from first time stamp).
                _tstamp2 = time_stamp;
                _opt.debug(u"evaluateNextStuffing: tstamp2 = %'d at %'d", time_stamp.tstamp, time_stamp.packet);
            }
        }
    }

    // If _tstamp2 not set in first segment or after buffer full, we cannot perform bitrate evaluation
    if (!_tstamp2 && (initial_position == 0 || !_input.canSeek(initial_position))) {
        ts::UString msg(u"no " + getTimeStampType() + u" found");
        if (initial_position > 0) {
            msg += ts::UString::Format(u" after packet %'d", initial_position);
        }
        if (_opt.reference_pid != ts::PID_NULL) {
            msg += ts::UString::Format(u" in PID %n, try another PID or increasing --buffer-size", _opt.reference_pid);
        }
        else {
            msg += u", try increasing --buffer-size";
        }
        _opt.fatal(msg);
    }

    // Restore initial position in the file
    if (!_input.seekBackward(size_t(_input.readPacketsCount() - initial_position), _opt)) {
        fatalError();
    }

    if (_tstamp2.has_value()) {
        // The segment is defined by the two time stamps. Compute new settings.
        assert(_tstamp1.has_value());
        assert(_tstamp1->tstamp < _tstamp2->tstamp);
        assert(_tstamp1->packet < _tstamp2->packet);

        // Segment duration in PCR unit:
        const uint64_t duration = _tstamp2->tstamp - _tstamp1->tstamp;

        // Actual number of input packets in the segment and corresponding bitrate.
        const uint64_t input_packets = _tstamp2->packet - _tstamp1->packet;
        const ts::BitRate input_bitrate = ts::BitRate(input_packets * ts::PKT_SIZE_BITS * ts::SYSTEM_CLOCK_FREQ) / duration;
        _opt.debug(u"segment: %'d ms, %'d packets, %'d b/s", (duration * 1000) / ts::SYSTEM_CLOCK_FREQ, input_packets, input_bitrate);

        // Target output number of bits in segment, plus the previously unstuffed bits (less than one packet).
        const uint64_t target_bits = _additional_bits + ((_opt.target_bitrate * duration) / ts::SYSTEM_CLOCK_FREQ).toInt();

        // Target output number of packets in the segment.
        const uint64_t target_packets = target_bits / ts::PKT_SIZE_BITS;

        // Compute number of stuffing packets to add
        if (input_packets > target_packets) {
            // Input burst greater than target bitrate.
            _opt.warning(u"input bitrate locally higher (%'d b/s) than target bitrate, cannot stuff", input_bitrate);
            _remaining_stuff_count = 0;
            _current_inter_packet = 0;
            _current_residue_packets = 0;
            _additional_bits = 0;
        }
        else {
            _remaining_stuff_count = target_packets - input_packets;
            _current_inter_packet = _remaining_stuff_count / input_packets;
            _current_residue_packets = _remaining_stuff_count % input_packets;
            _additional_bits = target_bits % ts::PKT_SIZE_BITS;
        }
    }
    else {
        // _tstamp2 is not set, we reached the end of file, keep previous settings.
        // Keep same inter-packet stuffing but erase fine tuning values.
        // This means that we do not do anything if the new bitrate is not at least
        // the double of the input one.
        _remaining_stuff_count = 0;
        _current_residue_packets = 0;
        _additional_bits = 0;
    }
}


//-----------------------------------------------------------------------------
// Process the content
//-----------------------------------------------------------------------------

void Stuffer::stuff()
{
    // Open input file
    if (!_input.openRead(_opt.input_file, 1, 0, _opt, _opt.input_format)) {
        fatalError();
    }

    _opt.debug(u"input file buffer size: %'d packets", _input.getBufferSize());

    // Remaining number of bits to stuff representing less than one packet
    _additional_bits = 0;

    // Locate first two time stamps,
    _tstamp1.reset();
    _tstamp2.reset();
    evaluateNextStuffing();
    assert(_tstamp1.has_value());
    assert(_tstamp2.has_value());

    // Create output file
    if (!_output.open(_opt.output_file, ts::TSFile::WRITE | ts::TSFile::SHARED, _opt, _opt.output_format)) {
        fatalError();
    }

    // Write leading stuffing packets
    writeStuffing(_opt.leading_packets);

    // Perform initial stuffing, up to the first time stamp
    simpleInterPacketStuffing(_opt.dyn_initial_inter_packet ? _current_inter_packet : _opt.initial_inter_packet, _tstamp1->packet);

    // Perform stuffing, segment after segment
    while (_tstamp2.has_value()) {
        assert(_input.readPacketsCount() < _tstamp2->packet);

        // Perform stuffing on current segment, loop on input packets, one by one.
        ts::TSPacket pkt;
        while (_input.readPacketsCount() < _tstamp2->packet && _input.read(&pkt, 1, _opt) == 1) {
            // Write the input packet.
            if (!_output.writePackets(&pkt, nullptr, 1, _opt)) {
                fatalError();
            }
            // Write stuffing packets after each input packet.
            const uint64_t residue = _current_residue_packets > 0 ? 1 : 0;
            const uint64_t count = std::min(_current_inter_packet + residue, _remaining_stuff_count);
            writeStuffing(count);
            _remaining_stuff_count -= count;
            _current_residue_packets -= residue;
        }

        // Optionally write a burst of null packets at the end of segment.
        writeStuffing(_remaining_stuff_count);
        _remaining_stuff_count = 0;

        // Evaluate stuffing need for next segment
        evaluateNextStuffing();
    }

    // Perform final stuffing, up to end of file
    simpleInterPacketStuffing(_opt.dyn_final_inter_packet ? _current_inter_packet : _opt.final_inter_packet, std::numeric_limits<uint64_t>::max());

    // Write trailing stuffing packets
    writeStuffing(_opt.trailing_packets);

    _opt.verbose(u"stuffing completed, read %'d packets, written %'d packets", _input.readPacketsCount(), _output.writePacketsCount());

    // Close files
    _output.close(_opt);
    _input.close(_opt);
}


//----------------------------------------------------------------------------
//  Program entry point
//----------------------------------------------------------------------------

int MainCode(int argc, char *argv[])
{
    Options opt(argc, argv);
    Stuffer stuffer(opt);
    stuffer.stuff();
    return EXIT_SUCCESS;
}
