#include "encoder_aac.h"
#include "headers_ffmpeg.h"

#include "ring_buffer.h"

#include "error_define.h"
#include "log_helper.h"

namespace am {
encoder_aac::encoder_aac()
{
    ffmpeg_register_all();

    _ring_buffer = new ring_buffer<AVFrame>(1024 * 1024 * 10);

    _inited = false;
    _running = false;

    _encoder = NULL;
    _encoder_ctx = NULL;
    _frame = NULL;
    _buff = NULL;
    _buff_size = 0;

    _bsf_ctx = NULL;

    _cond_notify = false;

#ifdef SAVE_AAC
    _aac_io_ctx = NULL;
    _aac_stream = NULL;
    _aac_fmt_ctx = NULL;
#endif
}

encoder_aac::~encoder_aac()
{
    stop();

    cleanup();

    delete _ring_buffer;
}

int encoder_aac::init(int nb_channels, int sample_rate, AVSampleFormat fmt, int bit_rate)
{
    int err = AE_NO;
    int ret = 0;

    if (_inited == true)
        return err;

    do {
        _encoder = (AVCodec*)avcodec_find_encoder(AV_CODEC_ID_AAC);
        if (!_encoder) {
            err = AE_FFMPEG_FIND_ENCODER_FAILED;
            break;
        }

        _encoder_ctx = avcodec_alloc_context3(_encoder);
        if (!_encoder_ctx) {
            err = AE_FFMPEG_ALLOC_CONTEXT_FAILED;
            break;
        }

        ffmpeg_set_channels(_encoder_ctx, nb_channels);
        ffmpeg_set_channel_layout(_encoder_ctx, ffmpeg_get_default_channel_layout(nb_channels));
        _encoder_ctx->sample_rate = sample_rate;
        _encoder_ctx->sample_fmt = fmt;
        _encoder_ctx->bit_rate = bit_rate;

        // 设置AAC profile为LC，确保RTMP兼容性
        _encoder_ctx->profile = FF_PROFILE_AAC_LOW;
        _encoder_ctx->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;
        _encoder_ctx->time_base.den = sample_rate;
        _encoder_ctx->time_base.num = 1;
        
        // 设置编码器选项，确保生成高质量AAC数据
        AVDictionary *options = NULL;
        av_dict_set(&options, "aac_coder", "twoloop", 0);
        av_dict_set(&options, "aac_pns", "1", 0);
        av_dict_set(&options, "aac_is", "1", 0);
        av_dict_set(&options, "aac_ms", "1", 0);
        av_dict_set(&options, "aac_ltp", "0", 0);
        av_dict_set(&options, "aac_pred", "0", 0);
        
        // 强制启用ADTS头部格式，确保RTMP兼容性
        // 注意："f"选项对编码器无效，需要在muxer层面处理ADTS格式
        // av_dict_set(&options, "f", "adts", 0);
        
        // 对于RTMP推流，确保生成ADTS头部
        // 不设置GLOBAL_HEADER，让每个包都包含ADTS头
        // _encoder_ctx->flags &= ~AV_CODEC_FLAG_GLOBAL_HEADER;
        // 由容器/复用层决定是否需要全局头，编码器保持默认行为（输出原始AAC帧）
        /* no manual change to GLOBAL_HEADER here */
        
        ret = avcodec_open2(_encoder_ctx, _encoder, &options);
        av_dict_free(&options);
        if (ret < 0) {
            err = AE_FFMPEG_OPEN_CODEC_FAILED;
            break;
        }

        // 调试信息：输出编码器配置
        al_debug("AAC encoder initialized: frame_size=%d, extradata_size=%d, profile=%d", 
                 _encoder_ctx->frame_size, _encoder_ctx->extradata_size, _encoder_ctx->profile);
        
        // 初始化BSF以适配必要的比特流转换
        // 目前不使用AAC相关BSF，编码器输出原始AAC帧，由复用器/容器处理（如写入ASC/ADTS等）
        _bsf_ctx = NULL;
        al_debug("No AAC BSF used; raw AAC frames will be passed to muxer for container-specific framing");
        
        /*
        const AVBitStreamFilter *bsf = av_bsf_get_by_name("aac_adtstoasc");
        if (!bsf) {
            al_error("Failed to find aac_adtstoasc BSF");
            err = AE_FFMPEG_FIND_BSF_FAILED;
            break;
        }
        
        ret = av_bsf_alloc(bsf, &_bsf_ctx);
        if (ret < 0) {
            al_error("Failed to allocate BSF context: %d", ret);
            err = AE_FFMPEG_ALLOC_BSF_FAILED;
            break;
        }
        
        // 复制编码器参数到BSF
        ret = avcodec_parameters_from_context(_bsf_ctx->par_in, _encoder_ctx);
        if (ret < 0) {
            al_error("Failed to copy parameters to BSF: %d", ret);
            err = AE_FFMPEG_BSF_PARAMS_FAILED;
            break;
        }
        
        ret = av_bsf_init(_bsf_ctx);
        if (ret < 0) {
            al_error("Failed to initialize BSF: %d", ret);
            err = AE_FFMPEG_BSF_INIT_FAILED;
            break;
        }
        
        al_debug("BSF initialized successfully for ADTS conversion");
        */
        
        // 确保frame_size有效
        if (_encoder_ctx->frame_size <= 0) {
            al_warn("AAC encoder frame_size is invalid: %d, setting to 1024", _encoder_ctx->frame_size);
            _encoder_ctx->frame_size = 1024;
        }

        _buff_size = av_samples_get_buffer_size(NULL, nb_channels, _encoder_ctx->frame_size, fmt, 1);
        _buff = (uint8_t *) av_malloc(_buff_size);

        _frame = av_frame_alloc();
        if (!_frame) {
            err = AE_FFMPEG_ALLOC_FRAME_FAILED;
            break;
        }

        ffmpeg_set_frame_channels(_frame, nb_channels);
        _frame->nb_samples = _encoder_ctx->frame_size;
        ffmpeg_set_frame_channel_layout(_frame, ffmpeg_get_default_channel_layout(nb_channels));
        _frame->format = fmt;
        _frame->sample_rate = sample_rate;

        ret = avcodec_fill_audio_frame(_frame, nb_channels, fmt, _buff, _buff_size, 0);

#ifdef SAVE_AAC
        ret = avio_open(&_aac_io_ctx, "save.aac", AVIO_FLAG_READ_WRITE);
        if (ret < 0) {
            err = AE_FFMPEG_OPEN_IO_FAILED;
            break;
        }

        _aac_fmt_ctx = avformat_alloc_context();
        if (!_aac_fmt_ctx) {
            err = AE_FFMPEG_ALLOC_CONTEXT_FAILED;
            break;
        }

        _aac_fmt_ctx->pb = _aac_io_ctx;
        _aac_fmt_ctx->oformat = av_guess_format(NULL, "save.aac", NULL);
        _aac_fmt_ctx->url = av_strdup("save.aac");

        _aac_stream = avformat_new_stream(_aac_fmt_ctx, NULL);
        if (!_aac_stream) {
            err = AE_FFMPEG_CREATE_STREAM_FAILED;
            break;
        }

        ret = avcodec_parameters_from_context(_aac_stream->codecpar, _encoder_ctx);
        if (ret < 0) {
            err = AE_FFMPEG_COPY_PARAMS_FAILED;
            break;
        }

        if (_aac_fmt_ctx->oformat->flags | AV_CODEC_FLAG_GLOBAL_HEADER) {
            _aac_stream->codec->extradata_size
                = _encoder_ctx->extradata_size; // +AV_INPUT_BUFFER_PADDING_SIZE;
            _aac_stream->codec->extradata = (uint8_t *) av_memdup(_encoder_ctx->extradata,
                                                                  _encoder_ctx->extradata_size);
        }
#endif

        _inited = true;

    } while (0);

    if (err != AE_NO) {
        al_debug("%s,error:%d", err2str(err), ret);
        cleanup();
    }

    return err;
}

int encoder_aac::get_extradata_size()
{
    return _encoder_ctx->extradata_size;
}

const uint8_t *encoder_aac::get_extradata()
{
    return (const uint8_t *) _encoder_ctx->extradata;
}

int encoder_aac::get_nb_samples()
{
    return _encoder_ctx->frame_size;
}

int encoder_aac::start()
{
    int error = AE_NO;

    if (_running == true) {
        return error;
    }

    if (_inited == false) {
        return AE_NEED_INIT;
    }

    _running = true;
    _thread = std::thread(std::bind(&encoder_aac::encode_loop, this));

    return error;
}

void encoder_aac::stop()
{
    _running = false;

    _cond_notify = true;
    _cond_var.notify_all();

    if (_thread.joinable())
        _thread.join();
}

int encoder_aac::put(const uint8_t *data, int data_len, AVFrame *frame)
{
    std::unique_lock<std::mutex> lock(_mutex);

    AVFrame frame_cp;
    memcpy(&frame_cp, frame, sizeof(AVFrame));

    _ring_buffer->put(data, data_len, frame_cp);

    _cond_notify = true;
    _cond_var.notify_all();

    return 0;
}

const AVRational &encoder_aac::get_time_base()
{
    return _encoder_ctx->time_base;
}

AVCodecID encoder_aac::get_codec_id()
{
    if (_inited == false)
        return AV_CODEC_ID_NONE;

    return _encoder->id;
}

int encoder_aac::encode(AVFrame *frame, AVPacket *packet)
{
    // 添加输入帧验证
    if (frame) {
        al_debug("Input frame: samples=%d, channels=%d, format=%d, pts=%lld", 
                 frame->nb_samples, ffmpeg_get_frame_channels(frame), frame->format, frame->pts);
    }
    
    int ret = avcodec_send_frame(_encoder_ctx, frame);
    al_debug("avcodec_send_frame returned: %d", ret);
    if (ret < 0) {
        al_error("avcodec_send_frame failed: %d", ret);
        return AE_FFMPEG_ENCODE_FRAME_FAILED;
    }

    bool packetReceived = false;
    while (ret == 0) {
        av_init_packet(packet);

        ret = avcodec_receive_packet(_encoder_ctx, packet);
        al_debug("avcodec_receive_packet returned: %d", ret);
        if (ret == AVERROR(EAGAIN)) {
            al_debug("Encoder needs more input: %d (this is normal)", ret);
            break; // 需要更多输入，这是正常情况
        }
        
        if (ret == AVERROR_EOF) {
            al_debug("Encoder EOF: %d", ret);
            break; // 编码器已结束
        }

        if (ret < 0) {
            al_error("avcodec_receive_packet failed: %d", ret);
            return AE_FFMPEG_READ_PACKET_FAILED;
        }

        if (ret == 0) {
            packetReceived = true;
            
            // 如果输入帧有有效的时间戳，将其传递给输出包
            if (frame && frame->pts != AV_NOPTS_VALUE) {
                packet->pts = frame->pts;
                packet->dts = frame->pts; // 对于音频，通常dts等于pts
                al_debug("Timestamp inherited from frame: pts=%lld, dts=%lld", packet->pts, packet->dts);
            } else {
                al_warn("Input frame has invalid timestamp (AV_NOPTS_VALUE), packet will have invalid timestamp");
            }
            
            // 验证编码器输出的包有效性
            if (packet->data == nullptr || packet->size <= 0) {
                al_warn("AAC packet is empty, skipping");
                av_packet_unref(packet);
                continue;
            }
            
            // 不再手动添加ADTS头部，完全由复用器/容器处理需要的打包格式
            al_debug("AAC packet ready: size=%d, pts=%lld, dts=%lld", 
                     packet->size, packet->pts, packet->dts);
            
            if (_on_data)
                _on_data(packet);
        }

#ifdef SAVE_AAC
        av_packet_rescale_ts(packet, _encoder_ctx->time_base, _aac_stream->time_base);
        packet->stream_index = _aac_stream->index;
        av_write_frame(_aac_fmt_ctx, packet);
#endif

        av_packet_unref(packet);
    }

    // 如果没有收到数据包且不是EAGAIN错误，可能有问题
    if (!packetReceived && ret != AVERROR(EAGAIN) && ret != AVERROR_EOF) {
        al_warn("No packet received and unexpected return code: %d", ret);
    }

    return AE_NO;
}

void encoder_aac::encode_loop()
{
    int len = 0;
    int error = AE_NO;

    AVPacket *packet = av_packet_alloc();
    AVFrame pcm_frame;

#ifdef SAVE_AAC
    avformat_write_header(_aac_fmt_ctx, NULL);
#endif

    while (_running) {
        std::unique_lock<std::mutex> lock(_mutex);
        // 改为基于条件的等待，避免固定50ms超时带来的延迟抖动
        _cond_var.wait(lock, [this] { return _cond_notify || !_running; });
        if (!_running)
            break;

        while ((len = _ring_buffer->get(_buff, _buff_size, pcm_frame))) {
            // 验证获取的音频数据
            if (len <= 0) {
                al_warn("Invalid audio data length: %d", len);
                continue;
            }
            
            // 确保音频数据正确填充到frame中
            if (len != _buff_size) {
                al_warn("Audio data size mismatch: expected=%d, actual=%d", _buff_size, len);
            }
            
            // 添加详细的调试信息
            static int frameCounter = 0;
            if (++frameCounter % 100 == 0) { // 每100帧打印一次
                al_debug("Ring buffer data: len=%d, buff_size=%d, frame_size=%d, channels=%d, format=%d", 
                         len, _buff_size, _encoder_ctx->frame_size, 
                         ffmpeg_get_frame_channels(_frame), _frame->format);
            }
            
            _frame->pts = pcm_frame.pts;
#if FFMPEG_VERSION_MAJOR >= 7
            // In FFmpeg 7+, pkt_pts is removed, use pts instead
            // _frame->pts is already set above
#else
            _frame->pkt_pts = pcm_frame.pkt_pts;
#endif
            _frame->pkt_dts = pcm_frame.pkt_dts;
            
            // 验证编码器帧的配置
            if (_frame->nb_samples != _encoder_ctx->frame_size) {
                al_warn("Frame samples mismatch: frame->nb_samples=%d, encoder->frame_size=%d", 
                        _frame->nb_samples, _encoder_ctx->frame_size);
            }

            int error = encode(_frame, packet);
            if (error != AE_NO) {
                if (_on_error)
                    _on_error(error);
                al_fatal("encode aac packet failed:%d", error);
                break;
            }
        }

        _cond_notify = false;
    }

    //flush pcm data in encoder
    al_debug("Flushing remaining AAC data from encoder");
    int flush_error = encode(NULL, packet);
    if (flush_error != AE_NO) {
        al_warn("Error during AAC encoder flush: %d", flush_error);
    }

    av_packet_free(&packet);
    
    // 输出编码统计信息
    al_info("AAC encoder stopped. Encoding session completed.");
    
#ifdef SAVE_AAC
    av_write_trailer(_aac_fmt_ctx);
#endif
}

void encoder_aac::cleanup()
{
    if (_encoder) {
        avcodec_close(_encoder_ctx);
    }

    _encoder = NULL;

    if (_encoder_ctx) {
        avcodec_free_context(&_encoder_ctx);
    }

    if (_frame)
        av_free(_frame);
    _frame = NULL;

    if (_buff)
        av_free(_buff);

    _buff = NULL;

    if (_bsf_ctx) {
        av_bsf_free(&_bsf_ctx);
        _bsf_ctx = NULL;
    }

    _encoder_ctx = NULL;

#ifdef SAVE_AAC
    if (_aac_fmt_ctx) {
        avio_closep(&_aac_fmt_ctx->pb);
        avformat_free_context(_aac_fmt_ctx);
    }
#endif
}

int encoder_aac::add_adts_header(AVPacket *packet)
{
    if (!packet || !packet->data || packet->size <= 0) {
        al_error("Invalid packet for ADTS header addition");
        return AE_INVALID_CONTEXT;
    }

    // ADTS头部长度为7字节
    const int adts_header_size = 7;
    const int new_size = packet->size + adts_header_size;

    // 生成ADTS头部
    uint8_t adts_header[adts_header_size];

    // 获取编码器参数
    int sample_rate = _encoder_ctx->sample_rate;
    int channels = ffmpeg_get_codec_context_channels(_encoder_ctx);
    int profile = _encoder_ctx->profile; // AAC-LC = 1

    // 采样率索引映射
    int freq_idx = 4; // 默认44100Hz
    switch (sample_rate) {
        case 96000: freq_idx = 0; break;
        case 88200: freq_idx = 1; break;
        case 64000: freq_idx = 2; break;
        case 48000: freq_idx = 3; break;
        case 44100: freq_idx = 4; break;
        case 32000: freq_idx = 5; break;
        case 24000: freq_idx = 6; break;
        case 22050: freq_idx = 7; break;
        case 16000: freq_idx = 8; break;
        case 12000: freq_idx = 9; break;
        case 11025: freq_idx = 10; break;
        case 8000:  freq_idx = 11; break;
        default:    freq_idx = 4; break;
    }

    // 声道配置
    int channel_config = channels;
    if (channels > 6) channel_config = 0; // 自定义配置

    // 构建ADTS头部
    adts_header[0] = 0xFF; // syncword (12 bits) - 高8位
    adts_header[1] = 0xF0 | (0 << 3) | (0 << 2) | 1; // syncword低4位 + MPEG版本 + Layer + protection_absent
    adts_header[2] = ((profile - 1) << 6) | (freq_idx << 2) | (0 << 1) | ((channel_config >> 2) & 1);
    adts_header[3] = ((channel_config & 3) << 6) | (0 << 5) | (0 << 4) | (0 << 3) | ((new_size >> 11) & 3);
    adts_header[4] = (new_size >> 3) & 0xFF;
    adts_header[5] = ((new_size & 7) << 5) | 0x1F;
    adts_header[6] = 0xFC; // number_of_raw_data_blocks_in_frame (2 bits) + 6位填充

    // 备份原始包的时间戳与元数据
    int64_t old_pts = packet->pts;
    int64_t old_dts = packet->dts;
    int64_t old_duration = packet->duration;
    int old_flags = packet->flags;
    int old_stream_index = packet->stream_index;
    int64_t old_pos = packet->pos;

    // 创建新包并拷贝数据（头部+原数据），以避免清空原包导致PTS/DTS丢失
    AVPacket *new_pkt = av_packet_alloc();
    if (!new_pkt) {
        al_error("Failed to allocate AVPacket for ADTS header");
        return AE_FFMPEG_ALLOC_FRAME_FAILED;
    }

    if (av_new_packet(new_pkt, new_size) < 0) {
        av_packet_free(&new_pkt);
        al_error("Failed to allocate packet buffer for ADTS header");
        return AE_FFMPEG_ALLOC_FRAME_FAILED;
    }

    // 拷贝ADTS头部和原始数据
    memcpy(new_pkt->data, adts_header, adts_header_size);
    memcpy(new_pkt->data + adts_header_size, packet->data, packet->size);

    // 恢复时间戳与元数据
    new_pkt->pts = old_pts;
    new_pkt->dts = old_dts;
    new_pkt->duration = old_duration;
    new_pkt->flags = old_flags;
    new_pkt->stream_index = old_stream_index;
    new_pkt->pos = old_pos;

    // 用新包替换旧包
    av_packet_unref(packet);
    av_packet_move_ref(packet, new_pkt);
    av_packet_free(&new_pkt);

    al_debug("ADTS header added: total_size=%d, profile=%d, freq_idx=%d, channels=%d (pts=%lld, dts=%lld)",
             new_size, profile, freq_idx, channels, packet->pts, packet->dts);

    return AE_NO;
}

} // namespace am