video-pc/AV/code/codec/codec_audio_encoder.cpp

#include "codec_audio_encoder.h"
#include "../base/logger.h"
#include <algorithm>
#include <chrono>

namespace av {
namespace codec {

// 静态成员初始化
std::vector<std::string> AudioEncoder::supportedEncoders_;
std::once_flag AudioEncoder::encodersInitFlag_;

// AudioResampler 实现
AudioResampler::AudioResampler()
    : swrCtx_(nullptr)
    , srcFormat_(AV_SAMPLE_FMT_NONE)
    , dstFormat_(AV_SAMPLE_FMT_NONE)
    , srcSampleRate_(0)
    , dstSampleRate_(0)
    , dstFrameSize_(0)
    , initialized_(false) {
    av_channel_layout_default(&srcLayout_, 0);
    av_channel_layout_default(&dstLayout_, 0);
}

AudioResampler::~AudioResampler() {
    if (swrCtx_) {
        swr_free(&swrCtx_);
    }
    av_channel_layout_uninit(&srcLayout_);
    av_channel_layout_uninit(&dstLayout_);
}

bool AudioResampler::init(const AVChannelLayout& srcLayout, AVSampleFormat srcFormat, int srcSampleRate,
                         const AVChannelLayout& dstLayout, AVSampleFormat dstFormat, int dstSampleRate) {
    if (swrCtx_) {
        swr_free(&swrCtx_);
    }
    
    av_channel_layout_uninit(&srcLayout_);
    av_channel_layout_uninit(&dstLayout_);
    
    // 复制声道布局
    if (av_channel_layout_copy(&srcLayout_, &srcLayout) < 0 ||
        av_channel_layout_copy(&dstLayout_, &dstLayout) < 0) {
        AV_LOGGER_ERROR("复制声道布局失败");
        return false;
    }
    
    srcFormat_ = srcFormat;
    dstFormat_ = dstFormat;
    srcSampleRate_ = srcSampleRate;
    dstSampleRate_ = dstSampleRate;
    
    // 创建重采样上下文
    int ret = swr_alloc_set_opts2(&swrCtx_,
                                  &dstLayout_, dstFormat_, dstSampleRate_,
                                  &srcLayout_, srcFormat_, srcSampleRate_,
                                  0, nullptr);
    if (ret < 0) {
        AV_LOGGER_ERRORF("创建重采样上下文失败: {}", ffmpeg_utils::errorToString(ret));
        return false;
    }
    
    // 初始化重采样器
    ret = swr_init(swrCtx_);
    if (ret < 0) {
        AV_LOGGER_ERRORF("初始化重采样器失败: {}", ffmpeg_utils::errorToString(ret));
        swr_free(&swrCtx_);
        return false;
    }
    
    // 计算输出帧大小
    dstFrameSize_ = av_rescale_rnd(1024, dstSampleRate_, srcSampleRate_, AV_ROUND_UP);
    
    // 创建输出帧
    dstFrame_ = makeAVFrame();
    if (!dstFrame_) {
        AV_LOGGER_ERROR("分配输出帧失败");
        return false;
    }
    
    dstFrame_->format = dstFormat_;
    dstFrame_->sample_rate = dstSampleRate_;
    av_channel_layout_copy(&dstFrame_->ch_layout, &dstLayout_);
    dstFrame_->nb_samples = dstFrameSize_;
    
    ret = av_frame_get_buffer(dstFrame_.get(), 0);
    if (ret < 0) {
        AV_LOGGER_ERRORF("分配帧缓冲区失败: {}", ffmpeg_utils::errorToString(ret));
        return false;
    }
    
    initialized_ = true;
    return true;
}

AVFramePtr AudioResampler::resample(const AVFramePtr& srcFrame) {
    if (!initialized_ || !srcFrame) {
        return nullptr;
    }
    
    // 计算输出样本数
    int dstSamples = av_rescale_rnd(srcFrame->nb_samples, dstSampleRate_, srcSampleRate_, AV_ROUND_UP);
    
    // 确保输出帧有足够的空间
    if (dstFrame_->nb_samples < dstSamples) {
        av_frame_unref(dstFrame_.get());
        dstFrame_->nb_samples = dstSamples;
        
        int ret = av_frame_get_buffer(dstFrame_.get(), 0);
        if (ret < 0) {
            AV_LOGGER_ERRORF("重新分配帧缓冲区失败: {}", ffmpeg_utils::errorToString(ret));
            return nullptr;
        }
    }
    
    // 执行重采样
    int convertedSamples = swr_convert(swrCtx_,
                                       dstFrame_->data, dstSamples,
                                       const_cast<const uint8_t**>(srcFrame->data), srcFrame->nb_samples);
    
    if (convertedSamples < 0) {
        AV_LOGGER_ERRORF("音频重采样失败: {}", ffmpeg_utils::errorToString(convertedSamples));
        return nullptr;
    }
    
    dstFrame_->nb_samples = convertedSamples;
    
    // 复制时间戳等信息
    av_frame_copy_props(dstFrame_.get(), srcFrame.get());
    
    // 调整时间戳
    if (srcFrame->pts != AV_NOPTS_VALUE) {
        dstFrame_->pts = av_rescale_q(srcFrame->pts, {1, srcSampleRate_}, {1, dstSampleRate_});
    }
    
    return std::move(dstFrame_);
}

std::vector<AVFramePtr> AudioResampler::flush() {
    std::vector<AVFramePtr> frames;
    
    if (!initialized_) {
        return frames;
    }
    
    while (true) {
        AVFramePtr frame = makeAVFrame();
        if (!frame) {
            break;
        }
        
        frame->format = dstFormat_;
        frame->sample_rate = dstSampleRate_;
        av_channel_layout_copy(&frame->ch_layout, &dstLayout_);
        frame->nb_samples = dstFrameSize_;
        
        int ret = av_frame_get_buffer(frame.get(), 0);
        if (ret < 0) {
            break;
        }
        
        int convertedSamples = swr_convert(swrCtx_,
                                           frame->data, dstFrameSize_,
                                           nullptr, 0);
        
        if (convertedSamples <= 0) {
            break;
        }
        
        frame->nb_samples = convertedSamples;
        frames.push_back(std::move(frame));
    }
    
    return frames;
}

// AudioEncoder 实现
AudioEncoder::AudioEncoder()
    : AbstractEncoder(MediaType::AUDIO) {
    AV_LOGGER_DEBUG("创建音频编码器");
}

AudioEncoder::~AudioEncoder() {
    close();
    AV_LOGGER_DEBUG("音频编码器已销毁");
}

ErrorCode AudioEncoder::open(const CodecParams& params) {
    std::lock_guard<std::mutex> lock(encodeMutex_);
    
    if (state_ != CodecState::IDLE && state_ != CodecState::CLOSED) {
        AV_LOGGER_WARNING("编码器已打开，先关闭再重新打开");
        close();
    }
    
    if (!validateParams(params)) {
        return ErrorCode::INVALID_ARGUMENT;
    }
    
    audioParams_ = static_cast<const AudioEncoderParams&>(params);
    params_ = params;
    
    ErrorCode result = initEncoder();
    if (result != ErrorCode::OK) {
        close();
        return result;
    }
    
    setState(CodecState::OPENED);
    AV_LOGGER_INFOF("音频编码器已打开: {} ({}Hz, {}ch, {}kbps)", 
                 audioParams_.codecName, audioParams_.sampleRate, 
                 audioParams_.channels, audioParams_.bitRate / 1000);
    
    return ErrorCode::OK;
}

void AudioEncoder::close() {
    std::lock_guard<std::mutex> lock(encodeMutex_);
    
    if (state_ == CodecState::CLOSED || state_ == CodecState::IDLE) {
        return;
    }
    
    // 清理资源
    resampler_.reset();
    convertedFrame_.reset();
    
    // 清理编解码上下文
    codecCtx_.reset();
    codec_ = nullptr;
    
    setState(CodecState::CLOSED);
    AV_LOGGER_DEBUG("音频编码器已关闭");
}

ErrorCode AudioEncoder::flush() {
    std::lock_guard<std::mutex> lock(encodeMutex_);
    
    if (state_ != CodecState::OPENED && state_ != CodecState::RUNNING) {
        return ErrorCode::INVALID_STATE;
    }
    
    setState(CodecState::FLUSHING);
    
    // 发送空帧来刷新编码器
    int ret = avcodec_send_frame(codecCtx_.get(), nullptr);
    if (ret < 0 && ret != AVERROR_EOF) {
        AV_LOGGER_ERRORF("刷新编码器失败: {}", ffmpeg_utils::errorToString(ret));
        reportError(static_cast<ErrorCode>(ret));
        return static_cast<ErrorCode>(ret);
    }
    
    setState(CodecState::OPENED);
    return ErrorCode::OK;
}

ErrorCode AudioEncoder::encode(const AVFramePtr& frame, std::vector<AVPacketPtr>& packets) {
    std::lock_guard<std::mutex> lock(encodeMutex_);
    
    if (state_ != CodecState::OPENED && state_ != CodecState::RUNNING) {
        return ErrorCode::INVALID_STATE;
    }
    
    setState(CodecState::RUNNING);
    
    auto startTime = std::chrono::high_resolution_clock::now();
    
    ErrorCode result = encodeFrame(frame, packets);
    
    auto endTime = std::chrono::high_resolution_clock::now();
    double processTime = std::chrono::duration<double, std::milli>(endTime - startTime).count();
    
    updateStats(result == ErrorCode::OK, processTime, 
                frame ? frame->nb_samples * audioParams_.channels * av_get_bytes_per_sample(static_cast<AVSampleFormat>(frame->format)) : 0);
    
    if (frameCallback_ && frame) {
        frameCallback_(frame);
    }
    
    for (const auto& packet : packets) {
        if (packetCallback_) {
            packetCallback_(packet);
        }
    }
    
    return result;
}

ErrorCode AudioEncoder::finishEncode(std::vector<AVPacketPtr>& packets) {
    // 不需要额外的锁，因为调用者已经持有锁
    if (state_ != CodecState::OPENED && state_ != CodecState::RUNNING) {
        return ErrorCode::INVALID_STATE;
    }
    
    // 发送空帧来刷新编码器
    int ret = avcodec_send_frame(codecCtx_.get(), nullptr);
    if (ret < 0 && ret != AVERROR_EOF) {
        AV_LOGGER_ERRORF("刷新编码器失败: {}", ffmpeg_utils::errorToString(ret));
        return static_cast<ErrorCode>(ret);
    }
    
    // 接收剩余的包
    return receivePackets(packets);
}

bool AudioEncoder::validateParams(const CodecParams& params) {
    if (params.type != MediaType::AUDIO) {
        AV_LOGGER_ERROR("参数媒体类型不是音频");
        return false;
    }
    
    const auto& audioParams = static_cast<const AudioEncoderParams&>(params);
    
    if (audioParams.sampleRate <= 0) {
        AV_LOGGER_ERROR("采样率无效");
        return false;
    }
    
    if (audioParams.channels <= 0) {
        AV_LOGGER_ERROR("声道数无效");
        return false;
    }
    
    if (audioParams.bitRate <= 0) {
        AV_LOGGER_ERROR("比特率无效");
        return false;
    }
    
    if (audioParams.codecName.empty()) {
        AV_LOGGER_ERROR("编码器名称为空");
        return false;
    }
    
    return true;
}

ErrorCode AudioEncoder::initEncoder() {
    // 查找编码器
    codec_ = avcodec_find_encoder_by_name(audioParams_.codecName.c_str());
    if (!codec_) {
        AV_LOGGER_ERRORF("未找到编码器: {}", audioParams_.codecName);
        return ErrorCode::CODEC_NOT_FOUND;
    }
    
    if (codec_->type != AVMEDIA_TYPE_AUDIO) {
        AV_LOGGER_ERROR("编码器类型不是音频");
        return ErrorCode::INVALID_ARGUMENT;
    }
    
    // 创建编码上下文
    codecCtx_ = makeAVCodecContext(codec_);
    if (!codecCtx_) {
        AV_LOGGER_ERROR("分配编码上下文失败");
        return ErrorCode::OUT_OF_MEMORY;
    }
    
    // 设置编码器参数
    ErrorCode result = setupEncoderParams();
    if (result != ErrorCode::OK) {
        return result;
    }
    
    // 打开编码器
    int ret = avcodec_open2(codecCtx_.get(), codec_, nullptr);
    if (ret < 0) {
        AV_LOGGER_ERRORF("打开编码器失败: {}", ffmpeg_utils::errorToString(ret));
        return static_cast<ErrorCode>(ret);
    }
    
    return ErrorCode::OK;
}

ErrorCode AudioEncoder::setupEncoderParams() {
    codecCtx_->bit_rate = audioParams_.bitRate;
    codecCtx_->sample_rate = audioParams_.sampleRate;
    
    // 设置声道布局
    ErrorCode result = setupChannelLayout();
    if (result != ErrorCode::OK) {
        return result;
    }
    
    // 设置采样格式
    codecCtx_->sample_fmt = getBestSampleFormat();
    
    // 设置帧大小
    if (codec_->capabilities & AV_CODEC_CAP_VARIABLE_FRAME_SIZE) {
        codecCtx_->frame_size = audioParams_.frameSize;
    }
    
    // 设置配置文件
    if (!audioParams_.profile.empty()) {
        av_opt_set(codecCtx_->priv_data, "profile", audioParams_.profile.c_str(), 0);
    }
    
    return ErrorCode::OK;
}

ErrorCode AudioEncoder::setupChannelLayout() {
    // 设置声道布局
    if (audioParams_.channelLayout.nb_channels > 0) {
        av_channel_layout_copy(&codecCtx_->ch_layout, &audioParams_.channelLayout);
    } else {
        // 根据声道数设置默认布局
        av_channel_layout_default(&codecCtx_->ch_layout, audioParams_.channels);
    }
    
    return ErrorCode::OK;
}

AVFramePtr AudioEncoder::convertFrame(const AVFramePtr& frame) {
    if (!frame) {
        return nullptr;
    }
    
    // 检查是否需要转换
    bool needConvert = false;
    
    if (frame->format != codecCtx_->sample_fmt) {
        needConvert = true;
    }
    
    if (frame->sample_rate != codecCtx_->sample_rate) {
        needConvert = true;
    }
    
    if (av_channel_layout_compare(&frame->ch_layout, &codecCtx_->ch_layout) != 0) {
        needConvert = true;
    }
    
    if (!needConvert) {
        // 创建一个新的帧来返回，避免拷贝构造
        AVFramePtr resultFrame = makeAVFrame();
        if (!resultFrame) {
            return nullptr;
        }
        
        // 复制帧数据
        if (av_frame_ref(resultFrame.get(), frame.get()) < 0) {
            return nullptr;
        }
        
        return resultFrame;
    }
    
    // 创建重采样器
    if (!resampler_) {
        resampler_ = std::make_unique<AudioResampler>();
        if (!resampler_->init(frame->ch_layout, static_cast<AVSampleFormat>(frame->format), frame->sample_rate,
                             codecCtx_->ch_layout, codecCtx_->sample_fmt, codecCtx_->sample_rate)) {
            AV_LOGGER_ERROR("初始化音频重采样器失败");
            return nullptr;
        }
    }
    
    return resampler_->resample(frame);
}

ErrorCode AudioEncoder::encodeFrame(const AVFramePtr& frame, std::vector<AVPacketPtr>& packets) {
    AVFramePtr processedFrame;
    
    if (frame) {
        // 格式转换
        processedFrame = convertFrame(frame);
        if (!processedFrame) {
            AV_LOGGER_ERROR("音频帧格式转换失败");
            return ErrorCode::CONVERSION_FAILED;
        }
    }
    
    // 发送帧到编码器
    int ret = avcodec_send_frame(codecCtx_.get(), processedFrame ? processedFrame.get() : nullptr);
    if (ret < 0) {
        AV_LOGGER_ERRORF("发送帧到编码器失败: {}", ffmpeg_utils::errorToString(ret));
        return static_cast<ErrorCode>(ret);
    }
    
    // 接收编码后的包
    return receivePackets(packets);
}

ErrorCode AudioEncoder::receivePackets(std::vector<AVPacketPtr>& packets) {
    while (true) {
        AVPacketPtr packet = makeAVPacket();
        if (!packet) {
            return ErrorCode::OUT_OF_MEMORY;
        }
        
        int ret = avcodec_receive_packet(codecCtx_.get(), packet.get());
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            break; // 需要更多输入或已结束
        }
        
        if (ret < 0) {
            AV_LOGGER_ERRORF("接收编码包失败: {}", ffmpeg_utils::errorToString(ret));
            return static_cast<ErrorCode>(ret);
        }
        
        packets.push_back(std::move(packet));
    }
    
    return ErrorCode::OK;
}

AVSampleFormat AudioEncoder::getBestSampleFormat() const {
    if (!codec_->sample_fmts) {
        return audioParams_.sampleFormat;
    }
    
    // 首先尝试使用指定的格式
    for (const AVSampleFormat* fmt = codec_->sample_fmts; *fmt != AV_SAMPLE_FMT_NONE; fmt++) {
        if (*fmt == audioParams_.sampleFormat) {
            return *fmt;
        }
    }
    
    // 返回第一个支持的格式
    return codec_->sample_fmts[0];
}

int AudioEncoder::getBestSampleRate() const {
    if (!codec_->supported_samplerates) {
        return audioParams_.sampleRate;
    }
    
    // 首先尝试使用指定的采样率
    for (const int* rate = codec_->supported_samplerates; *rate != 0; rate++) {
        if (*rate == audioParams_.sampleRate) {
            return *rate;
        }
    }
    
    // 返回第一个支持的采样率
    return codec_->supported_samplerates[0];
}

std::vector<std::string> AudioEncoder::getSupportedEncoders() {
    std::call_once(encodersInitFlag_, findSupportedEncoders);
    return supportedEncoders_;
}

std::string AudioEncoder::getRecommendedEncoder() {
    auto encoders = getSupportedEncoders();
    
    // 优先选择高质量编码器
    const char* preferredEncoders[] = {"libfdk_aac", "aac", "libopus", "opus", "libmp3lame", "mp3"};
    
    for (const char* encoder : preferredEncoders) {
        if (std::find(encoders.begin(), encoders.end(), encoder) != encoders.end()) {
            return encoder;
        }
    }
    
    return encoders.empty() ? "" : encoders[0];
}

bool AudioEncoder::isSampleFormatSupported(const std::string& codecName, AVSampleFormat format) {
    const AVCodec* codec = avcodec_find_encoder_by_name(codecName.c_str());
    if (!codec || !codec->sample_fmts) {
        return false;
    }
    
    for (const AVSampleFormat* fmt = codec->sample_fmts; *fmt != AV_SAMPLE_FMT_NONE; fmt++) {
        if (*fmt == format) {
            return true;
        }
    }
    
    return false;
}

bool AudioEncoder::isSampleRateSupported(const std::string& codecName, int sampleRate) {
    const AVCodec* codec = avcodec_find_encoder_by_name(codecName.c_str());
    if (!codec) {
        return false;
    }
    
    if (!codec->supported_samplerates) {
        return true; // 支持所有采样率
    }
    
    for (const int* rate = codec->supported_samplerates; *rate != 0; rate++) {
        if (*rate == sampleRate) {
            return true;
        }
    }
    
    return false;
}

bool AudioEncoder::isChannelLayoutSupported(const std::string& codecName, const AVChannelLayout& layout) {
    const AVCodec* codec = avcodec_find_encoder_by_name(codecName.c_str());
    if (!codec) {
        return false;
    }
    
    if (!codec->ch_layouts) {
        return true; // 支持所有声道布局
    }
    
    for (const AVChannelLayout* ch_layout = codec->ch_layouts; ch_layout->nb_channels != 0; ch_layout++) {
        if (av_channel_layout_compare(ch_layout, &layout) == 0) {
            return true;
        }
    }
    
    return false;
}

void AudioEncoder::findSupportedEncoders() {
    AV_LOGGER_INFO("查找可用的音频编码器...");
    
    for (const char* encoder : AUDIO_ENCODERS) {
        if (CodecFactory::isCodecSupported(encoder, CodecType::ENCODER, MediaType::AUDIO)) {
            supportedEncoders_.emplace_back(encoder);
            AV_LOGGER_INFOF("找到音频编码器: {}", encoder);
        }
    }
    
    AV_LOGGER_INFOF("总共找到 {} 个可用的音频编码器", supportedEncoders_.size());
}

// AudioEncoderFactory 实现
std::unique_ptr<AudioEncoder> AudioEncoderFactory::create(const std::string& codecName) {
    auto encoder = std::make_unique<AudioEncoder>();
    
    if (!codecName.empty()) {
        if (!CodecFactory::isCodecSupported(codecName, CodecType::ENCODER, MediaType::AUDIO)) {
            AV_LOGGER_ERRORF("不支持的编码器: {}", codecName);
            return nullptr;
        }
    }
    
    return encoder;
}

std::unique_ptr<AudioEncoder> AudioEncoderFactory::createBest() {
    std::string codecName = AudioEncoder::getRecommendedEncoder();
    
    if (codecName.empty()) {
        AV_LOGGER_ERROR("未找到可用的音频编码器");
        return nullptr;
    }
    
    return create(codecName);
}

std::unique_ptr<AudioEncoder> AudioEncoderFactory::createLossless() {
    auto encoders = AudioEncoder::getSupportedEncoders();
    
    // 优先选择无损编码器
    const char* losslessEncoders[] = {"flac", "pcm_s24le", "pcm_s16le"};
    
    for (const char* encoder : losslessEncoders) {
        if (std::find(encoders.begin(), encoders.end(), encoder) != encoders.end()) {
            return create(encoder);
        }
    }
    
    AV_LOGGER_WARNING("未找到无损音频编码器，使用默认编码器");
    return createBest();
}

} // namespace codec
} // namespace av