//----------------------------------------------------------------------------
//
// 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
//
//----------------------------------------------------------------------------
//
//  Fix continuity counters in a TS file
//
//----------------------------------------------------------------------------

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


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

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

        bool         test = false;          // Test mode
        bool         circular = false;      // Add empty packets to enforce circular continuity
        bool         no_replicate = false;  // Option --no-replicate-duplicated
        ts::UString  filename {};           // File name
        std::fstream file {};               // File buffer

        // Check if there was an I/O error on the file.
        // Print an error message if this is the case.
        bool fileError(const ts::UChar* message);
    };
}

// Constructor.
Options::Options(int argc, char *argv[]) :
    Args(u"Fix continuity counters in a transport stream", u"[options] filename")
{
    option(u"", 0, FILENAME, 1, 1);
    help(u"", u"MPEG capture file to be modified.");

    option(u"circular", 'c');
    help(u"circular",
         u"Enforce continuity when the file is played repeatedly. "
         u"Add empty packets, if necessary, on each PID so that the "
         u"continuity is preserved between end and beginning of file.");

    option(u"noaction");
    help(u"noaction", u"Legacy equivalent of --no-action.");

    option(u"no-action", 'n');
    help(u"no-action", u"Display what should be performed but do not modify the file.");

    option(u"no-replicate-duplicated");
    help(u"no-replicate-duplicated",
         u"Two successive packets in the same PID are considered as duplicated if they have "
         u"the same continuity counter and same content (except PCR, if any). "
         u"By default, duplicated input packets are replicated as duplicated on output "
         u"(the corresponding output packets have the same continuity counters). "
         u"When this option is specified, the input packets are not considered as duplicated and "
         u"the output packets receive individually incremented countinuity counters.");

    analyze(argc, argv);

    filename = value(u"");
    circular = present(u"circular");
    test = present(u"no-action") || present(u"noaction");
    no_replicate = present(u"no-replicate-duplicated");

    exitOnError();
}

// Check error on file
bool Options::fileError(const ts::UChar* message)
{
    if (file) {
        return false;
    }
    else {
        error(u"%s: %s", filename, message);
        return true;
    }
}


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

int MainCode(int argc, char *argv[])
{
    Options opt(argc, argv);
    ts::ContinuityAnalyzer fixer(ts::AllPIDs(), &opt);

    // Configure the CC analyzer.
    fixer.setDisplay(true);
    fixer.setFix(!opt.test);
    fixer.setReplicateDuplicated(!opt.no_replicate);
    fixer.setMessageSeverity(opt.test ? ts::Severity::Info : ts::Severity::Verbose);

    // Open file in read/write mode (CC are overwritten)
    std::ios::openmode mode = std::ios::in | std::ios::binary;
    if (!opt.test) {
        mode |= std::ios::out;
    }

    opt.file.open(opt.filename.toUTF8().c_str(), mode);

    if (!opt.file) {
        opt.error(u"cannot open file %s", opt.filename);
        return EXIT_FAILURE;
    }

    // Process all packets in the file
    ts::TSPacket pkt;

    for (;;) {

        // Save position of current packet
        const std::ios::pos_type pos = opt.file.tellg();
        if (opt.fileError(u"error getting file position")) {
            break;
        }

        // Read a TS packet
        if (!pkt.read(opt.file, true, opt)) {
            break; // end of file
        }

        // Process packet
        if (!fixer.feedPacket(pkt) && !opt.test) {
            // Packet was modified, need to rewrite it.
            // Rewind to beginning of current packet
            opt.file.seekp(pos);
            if (opt.fileError(u"error setting file position")) {
                break;
            }
            // Rewrite the packet
            pkt.write(opt.file, opt);
            if (opt.fileError(u"error rewriting packet")) {
                break;
            }
            // Make sure the get position is ok
            opt.file.seekg(opt.file.tellp());
            if (opt.fileError(u"error setting file position")) {
                break;
            }
        }
    }

    opt.verbose(u"%'d packets read, %'d discontinuities, %'d packets updated", fixer.totalPackets(), fixer.errorCount(), fixer.fixCount());

    // Append empty packet to ensure circular continuity
    if (opt.circular && opt.valid()) {

        // Create an empty packet (no payload, 184-byte adaptation field)
        pkt = ts::NullPacket;
        pkt.b[3] = 0x20;    // adaptation field, no payload
        pkt.b[4] = 183;     // adaptation field length
        pkt.b[5] = 0x00;    // nothing in adaptation field

        // Ensure write position is at end of file
        if (!opt.test) {
            // First, need to clear the eof bit
            opt.file.clear();
            // Set write position at eof
            opt.file.seekp(0, std::ios::end);
            // Returned value ignored on purpose, just report error when needed.
            // coverity[CHECKED_RETURN]
            opt.fileError(u"error setting file position");
        }

        // Loop through all PIDs, adding packets where some are missing
        for (ts::PID pid = 0; opt.valid() && pid < ts::PID_MAX; pid++) {
            const uint8_t first_cc = fixer.firstCC(pid);
            uint8_t last_cc = fixer.lastCC(pid);
            if (first_cc != ts::INVALID_CC && first_cc != ((last_cc + 1) & ts::CC_MASK)) {
                // We must add some packets on this PID
                opt.verbose(u"PID: 0x%04X, adding %2d empty packets", pid, ts::ContinuityAnalyzer::MissingPackets(last_cc, first_cc));
                if (!opt.test) {
                    for (;;) {
                        last_cc = (last_cc + 1) & ts::CC_MASK;
                        if (first_cc == last_cc) {
                            break; // complete
                        }
                        // Update PID and CC in the packet
                        pkt.setPID(ts::PID(pid));
                        pkt.setCC(last_cc);
                        // Write the new packet
                        pkt.write(opt.file, opt);
                        if (opt.fileError(u"error writing extra packet")) {
                            break;
                        }
                    }
                }
            }
        }
    }

    opt.file.close();

    return opt.valid() ? EXIT_SUCCESS : EXIT_FAILURE;
}
