/***************************************************************************
                          ADM_edit.cpp  -  description
                             -------------------
    begin                : Thu Feb 28 2002
    copyright            : (C) 2002/2008 by mean
    email                : fixounet@free.fr
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   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 "ADM_cpp.h"
#include "ADM_default.h"
#include "A_functions.h"
#include "GUI_ui.h"
#include <algorithm>
#include <map>

#include "ADM_edPtsDts.h"
#include "ADM_edit.hxx"
#include "DIA_coreToolkit.h"
#include "fourcc.h"
#include "prefs.h"

#include "ADM_coreDemuxer.h"
#include "ADM_vidMisc.h"

/**
    \fn ADM_Composer

*/
bool ADM_Composer::checkForValidPts(_SEGMENT *seg)
{
    int checkRange = 100;
    _VIDEOS *vid = _segments.getRefVideo(seg->_reference);
    vidHeader *hdr = vid->_aviheader;
    int totalFrames = hdr->getVideoStreamHeader()->dwLength;

    if (checkRange > totalFrames)
        checkRange = totalFrames;
    uint64_t from = seg->_startTimeUs;
    if (!seg->_refStartTimeUs && vid->firstFramePts)
        from += vid->firstFramePts;
    if (!goToTimeVideo(from))
    {
        ADM_warning("Cannot navigate, cannot check file for broken PTS.\n");
        return false;
    }
    getVideoPixelAndColorInfo(&(vid->srcPixFrmt), NULL, NULL, NULL, &(vid->srcColSpace));
    // vid->lastSentFrame=0; // We may not touch this counter here else we confuse the decoder, esp. with MT enabled.
    vid->dontTrustBFramePts = false;
    uint64_t inc = vid->timeIncrementInUs;

    stats.reset();

    bool bFramesPresent = false;
    for (int i = 0; i < totalFrames; i++)
    {
        uint32_t flags;
        if (!hdr->getFlags(i, &flags))
            break;
        if (flags & AVI_B_FRAME)
        {
            bFramesPresent = true;
            break;
        }
    }

    /* If the demuxer knows that B-frames are present but cannot provide
    valid pts, we may proceed with pts reconstruction right away. */
    if (!bFramesPresent || hdr->providePts())
    {
        ADM_info("Checking file for broken PTS...\n");
        ADM_info("Checking %d frames out of %d.\n", checkRange, totalFrames);
        DIA_workingBase *working =
            createWorking(QT_TRANSLATE_NOOP("ADM_Composer", "Checking if timestamps are valid.."));
        for (int i = 0; i < checkRange; i++)
        {
            if (false == working->isAlive())
                break;
            DecodeNextPicture(seg->_reference);
            working->update(i, checkRange);
        }
        goToTimeVideo(from);
        delete working;
        ADM_info("-------- Stats :----------\n");
#define INFO(x) ADM_info(#x ":%d\n", stats.x);
        INFO(nbBFrames)
        INFO(nbPFrames)
        INFO(nbIFrames)
        INFO(nbNoImage)
        INFO(nbPtsgoingBack)
        ADM_info("-------- /Stats ----------\n");

        if (stats.nbBFrames)
            bFramesPresent = true;
    }

    if (!bFramesPresent && hdr->providePts() == false)
    {
        ADM_info("No B-frames and no PTS, setting PTS equal DTS\n");
        return setPtsEqualDts(hdr, inc);
    }
    // check whether DTS are completely missing, ignore the first frame
    bool noDts = true;
    for (int i = 1; i < totalFrames; i++)
    {
        uint64_t pts, dts;
        hdr->getPtsDts(i, &pts, &dts);
        if (dts != ADM_NO_PTS)
        {
            noDts = false;
            break;
        }
    }
    if (!stats.nbBFrames && !stats.nbPtsgoingBack && hdr->providePts() && noDts)
    {
        ADM_info("No B-frames and no DTS, setting DTS equal PTS\n");
        return setPtsEqualDts(hdr, inc);
    }

    if (stats.nbPtsgoingBack > 1 || (bFramesPresent && hdr->providePts() == false))
    {
#ifdef WORK_AROUND_BAD_PTS
        if (!GUI_Question(QT_TRANSLATE_NOOP("ADM_Composer",
                                            "This video contains B-frames, but presentation time stamps (PTS) are "
                                            "either missing or monotonically increasing. "
                                            "Avidemux can try to reconstruct correct PTS by decoding the entire video. "
                                            "This may take a lot of time. Proceed?")))
        {
            vid->dontTrustBFramePts = true;
        }
        else
        {
            goToTimeVideo(from);
            vid->lastSentFrame = 0;

            if (false == ADM_verifyDts(hdr, inc))
                return false;
            if (false == setPtsEqualDts(hdr, inc))
                return false;

            std::map<uint64_t, uint32_t> timing;
            std::vector<uint64_t> ListOfPts, ListOfDts;

            for (uint32_t i = 0; i < totalFrames; i++)
            {
                uint64_t pts, dts;
                hdr->getPtsDts(i, &pts, &dts);
                timing.insert({dts, i});
            }

            EditorCache *cache = vid->_videoCache;
            ADM_assert(cache);
            cache->flush();

            bool interlaced = false; // when true and less than ~ 2/3 of frames produce pics => field-encoded
            uint64_t bfdelay = 0;
            std::map<uint64_t, uint32_t>::iterator it = timing.begin();

            DIA_workingBase *decoding = createWorking(QT_TRANSLATE_NOOP("ADM_Composer", "Decoding video..."));

            while (!vid->decoder->endOfStreamReached())
            {
                if (false == decoding->isAlive())
                {
                    vid->dontTrustBFramePts = true;
                    goToTimeVideo(from);
                    delete decoding;
                    return true;
                }
                decoding->update(vid->lastSentFrame, totalFrames);

                ADMCompressedImage img;
                img.data = compBuffer;
                img.cleanup(vid->lastSentFrame);
                img.dataLength = 0;

                if (!vid->decoder->getDrainingState())
                {
                    if (!hdr->getFrame(vid->lastSentFrame, &img))
                    {
                        ADM_warning("getFrame failed for frame %" PRIu32 "\n", vid->lastSentFrame);
                        img.dataLength = 0;
                        vid->decoder->setDrainingState(true);
                    }
                    else
                    {
                        vid->lastSentFrame++;
                    }
                }

                ADMImage *pic = cache->getFreeImage(seg->_reference);
                if (!pic)
                {
                    ADM_warning("Cannot find free image in cache, aborting PTS reconstruction.\n");
                    rewind();
                    delete decoding;
                    return false;
                }
                if (false == decompressImage(pic, &img, seg->_reference))
                {
                    if (vid->decoder->keepFeeding())
                    {
                        // printf("[checkForValidPts] No pic for frame %u yet.\n",vid->lastSentFrame);
                        continue;
                    }
                    if (vid->decoder->endOfStreamReached())
                    {
                        // printf("[checkForValidPts] End of stream\n");
                        break;
                    }
                    ADM_warning("Decoding error for frame %d\n", vid->lastSentFrame);
                    continue;
                }
                if (pic->Pts == ADM_NO_PTS)
                {
                    // printf("[checkForValidPts] No PTS out of decoder, last frame: %d\n",vid->lastSentFrame);
                    cache->invalidate(pic);
                    continue;
                }
                uint64_t pts = pic->Pts; // used only to match the decoded frame to its source
                if (pic->flags & AVI_FIELD_STRUCTURE)
                    interlaced = true;
                cache->invalidate(pic);

                ListOfPts.push_back(pts);
                ListOfDts.push_back(pts);
            }
            delete decoding;
            decoding = NULL;

            // invalidate pts in the index
            it = timing.begin();
            hdr->setPtsDts(0, 0, it->first);
            it++;
            for (int i = 1; i < totalFrames; i++)
            {
                if (it == timing.end())
                    break;
                hdr->setPtsDts(i, ADM_NO_PTS, it->first);
                it++;
            }
            // now set pts for all frames which produced output
            std::sort(ListOfDts.begin(), ListOfDts.end());
            int nbValid = ListOfDts.size();
            // printf("[checkForValidPts] We have %d pics out of %d frames.\n",nbValid,totalFrames);
            if (interlaced && nbValid < totalFrames * 2 / 3) // arbitrary threshold
            {
                ADM_info("Video seems to be field-encoded.\n");
                vid->fieldEncoded = true;
                inc *= 2;
            }
            {
                int frame = 0;
                for (std::vector<uint64_t>::iterator t = ListOfPts.begin(); t != ListOfPts.end(); t++)
                {
                    if (frame >= nbValid)
                        break;
                    it = timing.find(*t);
                    if (it == timing.end())
                    {
                        ADM_warning("Timestamp %s not found!\n", ADM_us2plain(*t));
                        continue;
                    }
                    uint64_t pts = ListOfDts.at(frame);
                    uint64_t dts = it->first;
                    // calculate B-frame delay
                    if (pts != ADM_NO_PTS && dts != ADM_NO_PTS && dts > pts && bfdelay < dts - pts)
                        bfdelay = dts - pts;
#if 0
                printf("adding entry %d for frame %u with PTS %s ",frame,it->second,ADM_us2plain(pts));
                printf("DTS %s\n",ADM_us2plain(dts));
#endif
                    hdr->setPtsDts(it->second, pts, dts);
                    frame++;
                }
            }

            ADM_info("B-frame delay set to %" PRIu64 " ms\n", bfdelay / 1000);

            uint64_t maxPts = ADM_NO_PTS;
            for (int i = 0; i < totalFrames; i++)
            { // add B-frame delay
                uint64_t pts, dts;
                hdr->getPtsDts(i, &pts, &dts);
                if (pts == ADM_NO_PTS)
                    continue;
                pts += bfdelay;
                if (maxPts == ADM_NO_PTS || maxPts < pts)
                    maxPts = pts;
                hdr->setPtsDts(i, pts, dts);
            }
            // update segment duration
            if (maxPts != ADM_NO_PTS)
            {
                maxPts += inc;
                if (maxPts > seg->_durationUs)
                {
                    printf("[checkForValidPts] Increasing segment duration from %s \n", ADM_us2plain(seg->_durationUs));
                    printf("to %s\n", ADM_us2plain(maxPts));
                    seg->_durationUs = maxPts;
                }
            }

            _segments.updateRefVideo();

            int nbAudTracks = getNumberOfActiveAudioTracks();
            for (int i = 0; i < nbAudTracks; i++)
            { // add the corresponding audio delay
                bool enabled;
                int ms = 0;
                getAudioShift(i, &enabled, &ms);
                ms += (int)(bfdelay / 1000);
                if (i == 0)
                    UI_setTimeShift(true, ms);
                ADM_info("Adding %d ms to audio delay for track %d, delay=%d ms\n", (int)(bfdelay / 1000), i, ms);
                setAudioShift(i, true, ms);
            }

            rewind();
            return true;
        }
#else
        if (!GUI_Question(QT_TRANSLATE_NOOP(
                "ADM_Composer", "Some timing information are incorrect.\nIt happens with some capture software.\n"
                                "If you re encode video we should drop these informations,\n else it will cause "
                                "dropped frame/jerky video.\n"
                                "If you just copy the video without reencoding,\n you should keep them.\n"
                                "Drop timing informations ?")))
            return true;
        // lookup min delta between pts & dts => b frame
        uint64_t delta;
        uint64_t minDelta = 1000 * 1000 * 10;
        for (int i = 0; i < totalFrames; i++)
        {
            uint64_t pts, dts;
            vid->_aviheader->getPtsDts(i, &pts, &dts);
            if (pts != ADM_NO_PTS && dts != ADM_NO_PTS)
            {
                delta = pts - dts;
                if (delta < minDelta)
                    minDelta = delta;
            }
        }
        ADM_info("Found %d to be the min value for PTS/DTS delta\n", (int)minDelta);
        minDelta += (minDelta >> 2);
        int processed = 0;
        for (int i = 1; i < totalFrames; i++) // cannot touch first frame
        {
            uint64_t pts, dts;
            vid->_aviheader->getPtsDts(i, &pts, &dts);
            if (pts != ADM_NO_PTS && dts != ADM_NO_PTS)
            {
                delta = pts - dts;
                if (delta <= minDelta)
                {
                    vid->_aviheader->setPtsDts(i, ADM_NO_PTS, dts);
                    processed++;
                }
            }
        }
        ADM_info("Cancelled %d pts as unreliable \n", processed);
#endif
    }
    goToTimeVideo(from);
    return true;
}
/**
 * \fn checkTiming
 * @param list
 * @param limit
 * @return
 */
static bool checkTiming(std::vector<uint64_t> &list, uint64_t limit)
{
    int n = list.size();
    int good = 0, bad = 0;
    for (int i = 0; i < n - 1; i++)
    {
        if ((list[i + 1] - list[i]) < limit)
            bad++;
        else
            good++;
    }
    ADM_info("\tGood : %d\n", good);
    ADM_info("\tBad  : %d\n", bad);
    if (!bad)
        return true;
    return false;
}

/**
    \fn checkForDoubledFps
    \brief Checks if the DTS increases by half the fps

*/
bool ADM_Composer::checkForDoubledFps(vidHeader *hdr, uint64_t timeIncrementUs)
{
    int totalFrames = hdr->getVideoStreamHeader()->dwLength;
    uint64_t dtsCeil = (timeIncrementUs * 18) / 10;
    std::vector<uint64_t> dtsList, ptsList;
    uint64_t lastPts = ADM_NO_PTS;
    uint64_t lastDts = ADM_NO_PTS;
    uint64_t maxPts = 0;
    uint64_t maxDts = 0;
    bool oooPts, oooDts, skipPts, skipDts;
    oooPts = oooDts = skipPts = skipDts = false;
    ADM_info("Checking for doubled FPS.., time increment ceiling = %d\n", (int)dtsCeil);
    for (int i = 0; i < totalFrames; i++)
    {
        uint64_t pts, dts;
        hdr->getPtsDts(i, &pts, &dts);
        if (dts != ADM_NO_PTS && lastDts != ADM_NO_PTS)
        { /* We've got two consecutive frames with valid DTS. We need to add both,
            the previous one now and the current one will be added to the list in
            the next iteration. */
            skipDts = false;
            if (!oooDts)
            {
                if (lastDts > maxDts)
                    maxDts = lastDts;
                if (maxDts > dts)
                    oooDts = true;
            }
        }
        if (!skipDts && lastDts != ADM_NO_PTS)
        {
            dtsList.push_back(lastDts);
            if (i == totalFrames - 1 && dts != ADM_NO_PTS)
                dtsList.push_back(dts); // add the last one
        }
        else
        { /* Stop adding further DTS until we have encountered two consecutive
            frames with valid DTS again. */
            skipDts = true;
        }
        lastDts = dts;

        if (pts != ADM_NO_PTS && lastPts != ADM_NO_PTS)
        { /* We've got two consecutive frames with valid PTS. We need to add both,
            the previous one now and the current one will be added to the list in
            the next iteration. */
            skipPts = false;
            if (!oooPts)
            {
                if (lastPts > maxPts)
                    maxPts = lastPts;
                if (maxPts > pts)
                    oooPts = true;
            }
        }
        if (!skipPts && lastPts != ADM_NO_PTS)
        {
            ptsList.push_back(lastPts);
            if (i == totalFrames - 1 && pts != ADM_NO_PTS)
                ptsList.push_back(pts); // add the last one
        }
        else
        { /* Stop adding further PTS until we have encountered two consecutive
            frames with valid PTS again. */
            skipPts = true;
        }
        lastPts = pts;
    }
    if (dtsList.size() < 2)
    {
        ADM_info("No consecutive frames with valid DTS found, can't safely halve FPS\n");
        return false;
    }
    if (ptsList.size() < 2)
    {
        ADM_info("No consecutive frames with valid PTS found, can't safely halve FPS\n");
        return false;
    }
    if (oooDts)
        std::sort(dtsList.begin(), dtsList.end());
    if (oooPts)
        std::sort(ptsList.begin(), ptsList.end());
    ADM_info("Checking %sDTS...\n", oooDts ? "sorted " : "");
    bool okDts = checkTiming(dtsList, dtsCeil);
    ADM_info("Checking %sPTS...\n", oooPts ? "sorted " : "");
    bool okPts = checkTiming(ptsList, dtsCeil);

    if (okDts && okPts)
    {
        ADM_info("We can safely halve fps\n");
    }
    else
    {
        ADM_info("Cannot halve fps\n");
    }
    return okDts && okPts;
}
// EOF
