#include "HxVideoWriter.h"

#include "HxTask.h"

HxVideoWriter::HxVideoWriter() : m_status(false) {}

bool HxVideoWriter::open(AVFormatContext *ifmt_ctx, QString filename)
{
    int ret = -1;

    video_index = -1;
    audio_index = -1;

    video_frame_index = 0;
    audio_frame_index = 0;

    ofmt_ctx = nullptr;

    video_dec_stream = nullptr;
    audio_dec_stream = nullptr;

    /* 查找流中视频与音频位置 */
    if(ifmt_ctx->nb_streams > 2)
    {
        m_error = QObject::tr("数据流异常");
        return false;
    }

    for (uint i = 0; i < ifmt_ctx->nb_streams; i++)
    {
        if (ifmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            video_index = static_cast<int>(i);

            video_dec_stream = ifmt_ctx->streams[i];
        }
        else if (ifmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO)
        {
            audio_index = static_cast<int>(i);

            audio_dec_stream = ifmt_ctx->streams[i];
        }
    }

    if (video_index == -1)
    {
        m_error = QObject::tr("未找到视频流");
        return false;
    }

    avformat_alloc_output_context2(&ofmt_ctx, nullptr, nullptr, filename.toStdString().data());

    if (!ofmt_ctx)
    {
        m_error = QObject::tr("无法打开输出文件");
        return false;
    }

    for (uint i = 0; i < ifmt_ctx->nb_streams; i++)
    {
        /* 根据输入流创建输出流 */
        AVCodec *codec = avcodec_find_decoder(ifmt_ctx->streams[i]->codecpar->codec_id);

        AVStream *out_stream = avformat_new_stream(ofmt_ctx, codec);
        if (!out_stream)
        {
            m_error = QObject::tr("无法创建输出流");
            return false;
        }

        /* 复制AVCodecContext的设置 */
        AVCodecContext *codec_ctx = avcodec_alloc_context3(codec);
        ret = avcodec_parameters_to_context(codec_ctx, ifmt_ctx->streams[i]->codecpar);
        if (ret < 0)
        {
            m_error = QObject::tr("未能将输入流中的codepar复制到编解码器上下文");
            avcodec_free_context(&codec_ctx);
            return false;
        }

        codec_ctx->codec_tag = 0;
        if (ofmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
            codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

        ret = avcodec_parameters_from_context(out_stream->codecpar, codec_ctx);
        if (ret < 0)
        {
            m_error = QObject::tr("无法将编解码器上下文复制到输出流codepar上下文");
            avcodec_free_context(&codec_ctx);
            return false;
        }

        avcodec_free_context(&codec_ctx);
    }

    ofmt_ctx->streams[video_index]->time_base = {1, 25};

    av_dump_format(ofmt_ctx, 0, filename.toStdString().data(), 1);

    if (!(ofmt_ctx->oformat->flags & AVFMT_NOFILE))
    {
        ret = avio_open(&ofmt_ctx->pb, filename.toStdString().data(), AVIO_FLAG_WRITE);
        if (ret < 0)
        {
            m_error = QObject::tr("无法打开输出文件");
            return false;
        }
    }

    ret = avformat_write_header(ofmt_ctx, nullptr);
    if (ret < 0)
    {
        m_error = QObject::tr("无法写入文件头信息");
        return false;
    }

    return (m_status = true);
}

void HxVideoWriter::send(AVPacket *packet)
{
    auto out_stream = ofmt_ctx->streams[packet->stream_index];

    if (packet->stream_index == video_index)
    {
        auto pts = av_rescale_q_rnd(packet->pts, video_dec_stream->time_base, out_stream->time_base, (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
        auto dts = av_rescale_q_rnd(packet->dts, video_dec_stream->time_base, out_stream->time_base, (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));

        if (video_frame_index == 0)
        {
            video_pts = pts < 0 ? 0 : pts;
            video_dts = dts < 0 ? 0 : dts;
        }
        else if (video_frame_index == 1)
        {
            if (video_pts > pts)
                video_pts = pts;
            else
                video_pts = pts - video_pts;

            if (video_dts > dts)
                video_dts = dts;
            else
                video_dts = dts - video_dts;
        }

        packet->pts = video_pts * video_frame_index;
        packet->dts = video_dts * video_frame_index;
        packet->duration = 0;

        video_frame_index++;
    }
    else if (packet->stream_index == audio_index)
    {
        auto pts = av_rescale_q_rnd(packet->pts, audio_dec_stream->time_base, out_stream->time_base, (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
        auto dts = av_rescale_q_rnd(packet->dts, audio_dec_stream->time_base, out_stream->time_base, (AVRounding)(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));

        if (audio_frame_index == 1 && pts > 0)
        {
            audio_pts = pts - audio_pts;
            audio_dts = dts - audio_dts;
        }
        else if (audio_frame_index == 0)
        {
            audio_pts = pts < 0 ? 0 : pts;
            audio_dts = dts < 0 ? 0 : dts;
        }

        packet->pts = audio_pts * audio_frame_index;
        packet->dts = audio_dts * audio_frame_index;
        packet->duration = 0;

        audio_frame_index++;
    }

    auto ret = av_interleaved_write_frame(ofmt_ctx, packet);

    if (ret < 0)
    {
        HxTrace::debug_write_line("VideoWrite", QString("Error muxing packet, ret = %1").arg(ret));

        //这里是否需要做处理  ???
    }
}

void HxVideoWriter::close()
{
    //Write file trailer
    if (m_status)
        av_write_trailer(ofmt_ctx);

    video_pts = 0;
    video_dts = 0;
    audio_pts = 0;
    audio_dts = 0;

    video_frame_index = 0;
    audio_frame_index = 0;

    /* close output */
    if (ofmt_ctx && !(ofmt_ctx->oformat->flags & AVFMT_NOFILE))
        avio_close(ofmt_ctx->pb);

    avformat_free_context(ofmt_ctx);

    m_status = false;
}

QString HxVideoWriter::error() { return m_error; }