/*
 * Copyright (c) 2023 Fancy Code.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "FfmpegMuxStreamV2.h"
#include "FfmpegOutputStream.h"
#include "ILog.h"
#include "LinuxApi.h"
#include "MediaBase.h"
#include "StatusCode.h"
#ifdef __cplusplus
extern "C" {
#endif
#include <libavcodec/codec_id.h>
#include <libavcodec/packet.h>
#include <libavformat/avformat.h>
#include <libavformat/avio.h>
#include <libavutil/dict.h>
#include <libavutil/error.h>
#ifdef __cplusplus
}
#endif
#include <cstddef>
#include <functional>
#include <memory>
// #include <mutex>
#include <string.h>
#include <string>
constexpr unsigned long long MUXING_NOT_START = 0;
FfmpegMuxStreamV2::FfmpegMuxStreamV2()
    : mOutputFormat(nullptr), mOptions(nullptr), mFilesMuxing(false), mFileMuxingDuration_us(0),
      mStartPts(MUXING_NOT_START), mMuxingFinised(false)
{
}
StatusCode FfmpegMuxStreamV2::OpenOutputFile(const OutputFileInfo &fileInfo)
{
    mOutputFileInfo = std::make_shared<OutputFileInfo>(fileInfo);
    return OpenMuxOutputFile(fileInfo.mFileName);
}
StatusCode FfmpegMuxStreamV2::CloseOutputFile(void)
{
    if (mOutputFormat && mOutputFormat->pb && mFilesMuxing) {
        av_write_trailer(mOutputFormat);
    }
    if (mVideoStream) {
        mVideoStream->UnInit();
    }
    if (mAudioStream) {
        mAudioStream->UnInit();
    }
    if (nullptr == mOutputFormat) {
        return CreateStatusCode(STATUS_CODE_OK);
    }
    if (!(mOutputFormat->oformat->flags & AVFMT_NOFILE)) {
        /* Close the output file. */
        avio_closep(&mOutputFormat->pb);
    }
    avformat_free_context(mOutputFormat);
    fx_system_v2("sync");
    mOutputFileInfo.reset();
    return CreateStatusCode(STATUS_CODE_OK);
}
void FfmpegMuxStreamV2::GetStreamData(const void *data, const size_t &size, const StreamInfo &streamInfo)
{
    if (mMuxingFinised) {
        /**
         * @brief Packaging has been completed according to the recording duration parameters, and the excess data
         * frames will be discarded.
         */
        return;
    }
    if (!MakeSureStreamHeanderOK(data, size)) {
        return;
    }
    if (streamInfo.mType == STREAM_TYPE_VIDEO_H264 && mVideoStream) {
        if (MUXING_NOT_START == mStartPts) {
            mStartPts = streamInfo.mTimeStamp_us;
        }
        /**
         * @brief Use the video's timestamp to count the playback duration of the packaged file.
         */
        CalculatingDuration(streamInfo.mTimeStamp_us);
        mVideoStream->WriteSourceData(data, size, streamInfo.mTimeStamp_us);
    }
    if (streamInfo.mType == STREAM_TYPE_AUDIO_G711A && mAudioStream) {
        mAudioStream->WriteSourceData(data, size, streamInfo.mTimeStamp_us);
    }
}
OutputFileInfo FfmpegMuxStreamV2::GetOutputFileInfo(void)
{
    OutputFileInfo finalFile = {.mDuration_ms = static_cast<unsigned int>(mFileMuxingDuration_us / 1000),
                                .mFinished = nullptr};
    memcpy(finalFile.mFileName, mOutputFileInfo->mFileName, OUTPUT_FILE_NAME_MAX);
    return finalFile;
}
StatusCode inline FfmpegMuxStreamV2::OpenMuxOutputFile(const std::string &fileName)
{
    AVDictionary *opt = nullptr;
    int ret = 0;
    /* allocate the output media context */
    avformat_alloc_output_context2(&mOutputFormat, nullptr, "mp4", fileName.c_str());
    if (!mOutputFormat) {
        LogError("Could not deduce output format from file.\n");
        return CreateStatusCode(STATUS_CODE_NOT_OK);
    }
    /* Add the audio and video streams using the default format codecs
     * and initialize the codecs. */
    if (mOutputFormat->oformat->video_codec != AV_CODEC_ID_NONE) {
        const std::string thumbnailFileName = mOutputFileInfo->mThumbnailFileName;
        mVideoStream = AddStream(mOutputFormat, AV_CODEC_ID_NONE, AV_CODEC_ID_NONE, thumbnailFileName);
        // mVideoStream = AddStream(mOutputFormat, mOutputFormat->oformat->video_codec, AV_CODEC_ID_H264);
        mVideoStream->SetWriteSourceDataCallback(
            std::bind(&FfmpegMuxStreamV2::GetAVPacketDataCallback, this, std::placeholders::_1));
    }
    if (mOutputFormat->oformat->audio_codec != AV_CODEC_ID_NONE) {
        mAudioStream = AddStream(mOutputFormat, mOutputFormat->oformat->audio_codec, AV_CODEC_ID_PCM_ALAW);
        mAudioStream->SetWriteSourceDataCallback(
            std::bind(&FfmpegMuxStreamV2::GetAVPacketDataCallback, this, std::placeholders::_1));
    }
    av_dump_format(mOutputFormat, 0, fileName.c_str(), 1);
    /* open the output file, if needed */
    if (!(mOutputFormat->oformat->flags & AVFMT_NOFILE)) {
        ret = avio_open(&mOutputFormat->pb, fileName.c_str(), AVIO_FLAG_WRITE);
        if (ret < 0) {
            char error_str[AV_ERROR_MAX_STRING_SIZE] = {0};
            LogError("Could not open '%s': %s\n",
                     fileName.c_str(),
                     av_make_error_string(error_str, AV_ERROR_MAX_STRING_SIZE, ret));
        }
    }
    if (mVideoStream) {
        return CreateStatusCode(STATUS_CODE_OK);
    }
    av_dict_set_int(&opt, "use_editlist", 0, 0);
    /* Write the stream header, if any. */
    ret = avformat_write_header(mOutputFormat, &opt);
    if (ret < 0) {
        char error_str[AV_ERROR_MAX_STRING_SIZE] = {0};
        LogError("Error occurred when opening output file: %s\n",
                 av_make_error_string(error_str, AV_ERROR_MAX_STRING_SIZE, ret));
        return CreateStatusCode(STATUS_CODE_NOT_OK);
    }
    av_dict_free(&opt);
    mFilesMuxing = true;
    return CreateStatusCode(STATUS_CODE_OK);
}
void FfmpegMuxStreamV2::GetAVPacketDataCallback(AVPacket *pkt)
{
    // std::lock_guard<std::mutex> locker(mMutex);
    int ret = 0;
    ret = av_interleaved_write_frame(mOutputFormat, pkt);
    /* pkt is now blank (av_interleaved_write_frame() takes ownership of
     * its contents and resets pkt), so that no unreferencing is necessary.
     * This would be different if one used av_write_frame(). */
    if (ret < 0) {
        char error_str[AV_ERROR_MAX_STRING_SIZE] = {0};
        LogInfo("Error while writing output packet: %s\n",
                av_make_error_string(error_str, AV_ERROR_MAX_STRING_SIZE, ret));
    }
}
void FfmpegMuxStreamV2::CalculatingDuration(const unsigned long long &pts_us)
{
    mFileMuxingDuration_us = pts_us - mStartPts;
    if (mFileMuxingDuration_us / 1000 >= mOutputFileInfo->mDuration_ms) {
        LogInfo("Muxing file finished, duration: %lld ms\n", mFileMuxingDuration_us / 1000);
        mMuxingFinised = true;
        if (mOutputFileInfo && mOutputFileInfo->mFinished) {
            *(mOutputFileInfo->mFinished) = static_cast<int>(OUTPUT_FILE_STATUS_FINISHED);
        }
    }
}
bool inline FfmpegMuxStreamV2::MakeSureStreamHeanderOK(const void *data, const size_t &size)
{
    int ret = 0;
    if (!mFilesMuxing) {
        bool fileMuxing = false;
        fileMuxing = mVideoStream->CheckStreamHeader(data, size);
        if (fileMuxing) {
            AVDictionary *opt = nullptr;
            av_dict_set_int(&opt, "use_editlist", 0, 0);
            /* Write the stream header, if any. */
            ret = avformat_write_header(mOutputFormat, nullptr);
            if (ret < 0) {
                char error_str[AV_ERROR_MAX_STRING_SIZE] = {0};
                LogError("Error occurred when opening output file: %s\n",
                         av_make_error_string(error_str, AV_ERROR_MAX_STRING_SIZE, ret));
                return false;
            }
            mFilesMuxing = true;
            av_dict_free(&opt);
        }
        else {
            LogWarning("Stream header not found, skip this frame.\n");
            return false;
        }
    }
    return true;
}
std::shared_ptr<FfmpegOutputStream> FfmpegMuxStreamV2::AddStream(AVFormatContext *outputFormat,
                                                                 enum AVCodecID encodecId, enum AVCodecID decodecId,
                                                                 const std::string &thumbnailFile)
{
    auto stream = std::make_shared<FfmpegOutputStream>(encodecId, decodecId);
    if (thumbnailFile.empty()) {
        stream->Init(outputFormat);
    }
    else {
        stream->Init(outputFormat, thumbnailFile);
    }
    return stream;
}