/*
 * 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 "FfmpegOutputStream.h"
#include "FfmpegDecoder.h"
#include "FfmpegEncoder.h"
#include "ILog.h"
#ifdef __cplusplus
extern "C" {
#endif
#include <libavcodec/codec_id.h>
#include <libavcodec/packet.h>
#include <libavformat/avformat.h>
#include <libavutil/avutil.h>
#include <libavutil/frame.h>
#include <libavutil/mathematics.h>
#include <libavutil/mem.h>
#include <libavutil/pixfmt.h>
#ifdef __cplusplus
}
#endif
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <string.h>
FfmpegOutputStream::FfmpegOutputStream(const AVCodecID &encodecId, const AVCodecID &dncodecId)
    : mEncodecId(encodecId), mDeccodecId(dncodecId), mTmpPkt(nullptr), mStream(nullptr), mStreamHeaderWritten(false)
{
}
bool FfmpegOutputStream::Init(AVFormatContext *outputFormat)
{
    mDecodeCallback = std::bind(&FfmpegOutputStream::GetDecodeDataCallback, this, std::placeholders::_1);
    mTmpPkt = av_packet_alloc();
    if (!mTmpPkt) {
        LogError("Could not allocate AVPacket\n");
        return false;
    }
    mStream = avformat_new_stream(outputFormat, nullptr);
    if (!mStream) {
        LogError("Could not allocate stream\n");
        return false;
    }
    if (mDeccodecId != AV_CODEC_ID_NONE) {
        mDecoder = std::make_shared<FfmpegDecoder>(mDeccodecId);
        mDecoder->Init();
    }
    mStream->id = outputFormat->nb_streams - 1;
    if (mEncodecId != AV_CODEC_ID_NONE) {
        mEncoder = std::make_shared<FfmpegEncoder>(mEncodecId);
        mEncoder->Init(outputFormat->flags);
        mStream->time_base = mEncoder->GetTimeBase();
        mEncoder->OpenEncoder(nullptr, mStream);
    }
    else {
        /**
         * @brief There is no need to set time_base here, time_base will be automatically corrected inside ffmpeg.
         *
         */
        mStream->time_base = (AVRational){1, 15};
        mStream->codecpar->codec_id = AV_CODEC_ID_H264;
        mStream->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
        mStream->codecpar->width = 1920;
        mStream->codecpar->height = 2160;
        // mStream->codecpar->bit_rate = 2073;
        mStream->codecpar->format = AV_PIX_FMT_YUV420P;
        mStream->codecpar->codec_tag = 0;
        mStream->codecpar->extradata = nullptr;
        mStream->codecpar->extradata_size = 0;
    }
    return true;
}
void FfmpegOutputStream::UnInit(void)
{
    if (mEncoder) {
        mEncoder->UnInit();
        mEncoder.reset();
    }
    if (mDecoder) {
        mDecoder->UnInit();
        mDecoder.reset();
    }
    av_packet_free(&mTmpPkt);
}
void FfmpegOutputStream::WriteSourceData(const void *data, const size_t &size, const unsigned long long &pts)
{
    if (mDecoder) {
        mDecoder->DecodeData(data, size, pts, mDecodeCallback);
        return;
    }
    AVPacket *tmpPkt = av_packet_alloc();
    static unsigned long long u64Interval = 0;
    // AVRational in_timebase = (AVRational){1, 15};
    if (mEncodeCallback) {
        tmpPkt->data = (uint8_t *)data;
        tmpPkt->size = size;
        tmpPkt->stream_index = mStream->index;
        constexpr int64_t ZERO_MEANS_UNKNOWN = 0;
        tmpPkt->duration = ZERO_MEANS_UNKNOWN;
        // tmpPkt->pts = u64Interval * 1000; // ת���� us
        tmpPkt->pts = av_rescale_q(pts, (AVRational){1, 1000000}, mStream->time_base);
        // LogInfo("vvvvvvvvvvvvvvvvvvvvvvvvv num:%d, den:%d\n", mStream->time_base.num, mStream->time_base.den);
        // LogInfo("vvvvvvvvvvvvvvvvvvvvvvvvv pts:%llu, duration:%d\n", tmpPkt->pts, tmpPkt->duration);
        // tmpPkt->pts = pts;
        u64Interval++;
        tmpPkt->dts = tmpPkt->pts;
        /* copy packet */
        // av_packet_rescale_ts(mTmpPkt, in_timebase, mStream->time_base);
        tmpPkt->pos = -1;
        mEncodeCallback(tmpPkt);
    }
    av_packet_unref(tmpPkt);
    av_packet_free(&tmpPkt);
}
void FfmpegOutputStream::SetWriteSourceDataCallback(std::function<void(AVPacket *pkt)> callback)
{
    mEncodeCallback = callback;
}
bool FfmpegOutputStream::CheckStreamHeader(const void *data, const size_t &size)
{
    if (mStreamHeaderWritten || mEncodecId != AV_CODEC_ID_NONE) {
        return true;
    }
    char *pData = (char *)data;
    for (size_t i = 0; i < size; i++) {
        if ((0x00 == pData[i]) && (0x00 == pData[i + 1]) && (0x00 == pData[i + 2]) && (0x01 == pData[i + 3]) &&
            (0x5 == (pData[i + 4] & 0x1F))) {
            uint8_t *extradata = (uint8_t *)av_mallocz(i + 1);
            if (!extradata) {
                LogError("Could not allocate extradata\n");
                return false;
            }
            LogInfo("Found extradata\n");
            memcpy(extradata, pData, i);
            mStream->codecpar->extradata = extradata;
            mStream->codecpar->extradata_size = i;
            mStreamHeaderWritten = true;
            return mStreamHeaderWritten;
        }
    }
    return false;
}
void FfmpegOutputStream::GetDecodeDataCallback(AVFrame *frame)
{
    if (mEncoder) {
        mEncoder->EncodeData(frame, mStream, mEncodeCallback);
        return;
    }
}