



#include "FileWriterFFmpeg.h"

using namespace std;

namespace mfs
{
    static std::string& GetAVErrorString(int ret, std::string& message)
    {
        char buffer[256]={0};
        message = av_make_error_string(buffer,256, ret);
    }
    FileWriterFFmpeg::FileWriterFFmpeg(const std::string& fileFullName):_fileFullName(fileFullName)
    {
    }
    
    FileWriterFFmpeg::~FileWriterFFmpeg()
    {
        CloseFile();
    }

    int FileWriterFFmpeg::CloseFile()
    {
        if (_oc)
        {
            // 写文件尾 
            av_write_trailer(_oc);
            // 关闭输出流 
            avio_close(_oc->pb);
            // 清理封装输出上下文 
            avformat_free_context(_oc);
            _oc= nullptr;
        }
        if (_c)
        {
            // 关闭编码器
            avcodec_close(_c);
            // 清理编码器上下文
            avcodec_free_context(&_c);
            _c= nullptr;
        }
        if (_frame)
        {
            av_frame_free(&_frame);
            _frame = nullptr;
        }
        if (_packet)
        {
            av_packet_free(&_packet);
            _packet = nullptr;
        }
        if (_actx)
        {
            swr_free(&_actx);
            _actx = nullptr;
        }
        
        
        return 0;
    }
    static int testsamples = 0;
    int FileWriterFFmpeg::WriteData(const uint8_t* data, int dataLength)
    {
        int ret = 0;
        std::string message;
        int samplesCount = dataLength/_perSampleBytes/_c->channels; //一帧音频存放的样本数量
        // 每次编码的采样数不能多余 编码器的frame_size
        // 如果编码器的采样格式不一样，要转换
        int encoderFrameBytes = _c->frame_size * _c->channels * _perSampleBytes;
        while (samplesCount>0)
        {
            int frameSamplesCount = samplesCount > _c->frame_size?_c->frame_size:samplesCount;
            int frameBytesCount = samplesCount > _c->frame_size?encoderFrameBytes:samplesCount * _c->channels * _perSampleBytes;

            const uint8_t* inputBuffer = data;
            int inputBufferLength = frameBytesCount;
            if(_actx)
            {
                // LOG_DEBUG << "CONVERT SAMPLEs" <<LOG_ENDL;
                // _samplesBuffer
                uint8_t *dstBuffer[1] = {_samplesBuffer.data()};
                const uint8_t *srcBuffer[1] = {data};
                ret = swr_convert(_actx, dstBuffer, frameSamplesCount, srcBuffer, frameSamplesCount);
                if (ret < 0)
                {
                    LOG_ERROR << "swr_convert error="<< GetAVErrorString(ret, message) << LOG_ENDL;
                    return -1;
                }
                inputBuffer = _samplesBuffer.data();
                inputBufferLength = ret * av_get_bytes_per_sample(_c->sample_fmt) * _c->channels;
            }

            testsamples+=frameSamplesCount;
            _frame->nb_samples = ret;
            _frame->pts = _pts;
            _frame->format = _c->sample_fmt;
            _frame->channel_layout = _c->channel_layout;
            _frame->sample_rate = _c->sample_rate;
            _pts += ret;
            // LOG_DEBUG<<"SEND sample count:"<<frameSamplesCount<<",testsamples="<<testsamples<<LOG_ENDL;
            // _frame->nb_samples = frameSamplesCount;
            ret = avcodec_fill_audio_frame(_frame, _c->channels, _c->sample_fmt, inputBuffer, inputBufferLength, 1);
            if (ret < 0)
            {
                LOG_ERROR << "avcodec_fill_audio_frame error="<< GetAVErrorString(ret, message) << LOG_ENDL;
                return -1;
            }
            
            
            // AVPacket pkt;
            // av_init_packet(&pkt);

            // 编码
            ret = avcodec_send_frame(_c,_frame);
            if (ret!=0)
            {
                LOG_ERROR << "avcodec_send_frame error="<<GetAVErrorString(ret, message) << LOG_ENDL;
            }
            av_frame_unref(_frame);
            ret = avcodec_receive_packet(_c, _packet);
            if (ret!=0)
            {
                LOG_ERROR << "avcodec_receive_packet error="<<GetAVErrorString(ret, message) << LOG_ENDL;
                continue;
            }

            ret = av_interleaved_write_frame(_oc,_packet);
            if (ret!=0)
            {
                LOG_ERROR << "av_interleaved_write_frame error="<<GetAVErrorString(ret, message) << LOG_ENDL;
                return -1;
            }

	        av_packet_unref(_packet);

            samplesCount -= frameSamplesCount;
            data += frameBytesCount;
        }
        
        return 0;
    } 


    void FileWriterFFmpeg::InitCodec()
    {
        switch (_c->codec->id)
        {
        case AV_CODEC_ID_AAC:
            _c->bit_rate = 64000;
            _c->sample_rate = 8000;
            _c->sample_fmt = AV_SAMPLE_FMT_FLTP;
            _c->channel_layout = AV_CH_LAYOUT_MONO;
            _c->channels = 1;
            _c->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
	        _c->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;
            break;
        default:
            _c->bit_rate = 64000;
            _c->sample_rate = 8000;
            _c->sample_fmt = AV_SAMPLE_FMT_S16;
            _c->channel_layout = AV_CH_LAYOUT_MONO;
            _c->channels = 1;
            _c->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
	        _c->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;
            break;
        }
    }

    AVCodecID FileWriterFFmpeg::GetCodecIDByFormat()
    {
        auto pos = _fileFullName.rfind('.');
        if (pos != string::npos)
        {
            auto extension = _fileFullName.substr(pos+1);
            if (extension == "aac")
            {
                return AV_CODEC_ID_AAC;
            }
            else
            {
                return AV_CODEC_ID_PCM_S16LE;
            }
        }
    }

    int FileWriterFFmpeg::OpenFile()
    {
        // av_register_all();
        // avcodec_register_all();
        
        std::string message;
        int ret = 0;
        // encoder
        AVCodec *codec = avcodec_find_encoder(GetCodecIDByFormat()); // avcodec_encode_audio2 (fmt_ctx->streams[stream_index]->codec, &enc_pkt, NULL, &got_frame); 
        if (!codec)
        {
            LOG_ERROR << "avcodec_find_encoder error="<<GetAVErrorString(ret, message) << LOG_ENDL;
            return -1;
        }
        _c = avcodec_alloc_context3(codec);
        if (!_c)
        {
            LOG_ERROR << "avcodec_alloc_context3 error="<< GetAVErrorString(ret, message) << LOG_ENDL;
            return -1;
        }

        InitCodec();
        
        ret = avcodec_open2(_c,codec,NULL);
        if (ret < 0)
        {
            LOG_ERROR << "avcodec_open2 error="<<GetAVErrorString(ret, message) << LOG_ENDL;
            return -1;
        }
    // _oc->audio_codec
        // muxer
        avformat_alloc_output_context2(&_oc,NULL,NULL, _fileFullName.c_str());
        if (!_oc)
        {
            LOG_ERROR << "avformat_alloc_output_context2 error="<<GetAVErrorString(ret, message) << LOG_ENDL;
            return -1;
        }

        AVStream *st = avformat_new_stream(_oc, NULL);
        st->codecpar->codec_tag = 0;
        avcodec_parameters_from_context(st->codecpar,_c);
    
        ret = avio_open(&_oc->pb, _fileFullName.c_str(), AVIO_FLAG_WRITE);
        if (ret < 0)
        {
            LOG_ERROR << "avio_open error="<<GetAVErrorString(ret, message) << LOG_ENDL;
            return -1;
        }
        
        ret = avformat_write_header(_oc,NULL);
        if (ret < 0)
        {
            LOG_ERROR << "avformat_write_header error="<<GetAVErrorString(ret, message) << LOG_ENDL;
            return -1;
        }
        
        _perSampleBytes = av_get_bytes_per_sample(AV_SAMPLE_FMT_S16);
        _frame = av_frame_alloc();
	    _packet = av_packet_alloc();
        // _frame->format = _c->sample_fmt;
        // _frame->channels = 1;
        // _frame->channel_layout = AV_CH_LAYOUT_MONO;

    
        if (_c->sample_fmt != AV_SAMPLE_FMT_S16)
        {
            _samplesBuffer.resize(_c->frame_size * _c->channels * av_get_bytes_per_sample(_c->sample_fmt));
            _actx = swr_alloc_set_opts(_actx,
                _c->channel_layout, _c->sample_fmt, _c->sample_rate,//输出格式
                _c->channel_layout, AV_SAMPLE_FMT_S16, _c->sample_rate,//输入格式
            0,0);
            if (!_actx)
            {
                cout << "swr_alloc_set_opts error="<<GetAVErrorString(ret, message) << endl;
                getchar();
                return -1;
            }
            ret = swr_init(_actx);
            if (ret < 0)
            {
                cout << "swr_init error="<<GetAVErrorString(ret, message) << endl;
                getchar();
                return -1;
            }
        }
        
        return 0;
    }

}