#include "encoder_video_x264.h"
#include "headers_ffmpeg.h"

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

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

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

encoder_video_x264::~encoder_video_x264()
{
    stop();

    cleanup();
}

int encoder_video_x264::init(
    int pic_width, int pic_height, int frame_rate, int bit_rate, int qb, int key_pic_sec)
{
    if (_inited == true)
        return AE_NO;

    int err = AE_NO;
    int ret = 0;

    AVDictionary *options = 0;

    // 使用ultrafast预设以获得最低延迟
    av_dict_set(&options, "preset", "ultrafast", 0);
    av_dict_set(&options, "tune", "zerolatency", 0);
    
    // 设置x264参数以匹配低延迟配置
    av_dict_set(&options, "x264-params", "keyint=30:min-keyint=30:no-scenecut=1", 0);
    av_dict_set(&options, "profile", "high", 0);
    av_dict_set(&options, "level", "4.1", 0);
    // 与命令行一致，明确设置CRF为18
    av_dict_set(&options, "crf", "18", 0);

    do {
        _encoder = avcodec_find_encoder(AV_CODEC_ID_H264);
        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;
        }

        _encoder_ctx->codec_id = AV_CODEC_ID_H264;
        _encoder_ctx->codec_type = AVMEDIA_TYPE_VIDEO;
        _encoder_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
        _encoder_ctx->width = pic_width;
        _encoder_ctx->height = pic_height;
        _encoder_ctx->time_base.num = 1;
        _encoder_ctx->time_base.den = frame_rate;
        _encoder_ctx->framerate = {frame_rate, 1};
        _encoder_ctx->bit_rate = bit_rate;

        // 设置GOP大小为30帧，与用户ffmpeg命令一致
        _encoder_ctx->gop_size = 30;

        // 使用CRF 18以获得高质量和低延迟的平衡
        // 注意：FFmpeg的libx264编码器通过crf参数控制质量
        // 这里我们设置合理的qmin和qmax值
        _encoder_ctx->qmin = 10;
        _encoder_ctx->qmax = 25;
        
        // 设置CRF值（通过全局质量参数）
        _encoder_ctx->global_quality = 18 * FF_QP2LAMBDA;

        _encoder_ctx->max_b_frames = 0; //NO B Frame
        _encoder_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

        ret = avcodec_open2(_encoder_ctx, _encoder, &options);
        if (ret != 0) {
            err = AE_FFMPEG_OPEN_CODEC_FAILED;
            break;
        }

        _frame = av_frame_alloc();

        _buff_size = av_image_get_buffer_size(_encoder_ctx->pix_fmt,
                                              _encoder_ctx->width,
                                              _encoder_ctx->height,
                                              1);

        _buff = (uint8_t *) av_malloc(_buff_size);
        if (!_buff) {
            break;
        }

        av_image_fill_arrays(_frame->data,
                             _frame->linesize,
                             _buff,
                             _encoder_ctx->pix_fmt,
                             _encoder_ctx->width,
                             _encoder_ctx->height,
                             1);

        _frame->format = _encoder_ctx->pix_fmt;
        _frame->width = _encoder_ctx->width;
        _frame->height = _encoder_ctx->height;

        _y_size = _encoder_ctx->width * _encoder_ctx->height;

        _time_base = _encoder_ctx->time_base;

        _inited = true;
    } while (0);

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

    if (options)
        av_dict_free(&options);

    return err;
}

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

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

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

    return _encoder->id;
}

void encoder_video_x264::cleanup()
{
    if (_frame)
        av_free(_frame);
    _frame = NULL;

    if (_buff)
        av_free(_buff);

    _buff = NULL;

    if (_encoder)
        avcodec_close(_encoder_ctx);

    _encoder = NULL;

    if (_encoder_ctx)
        avcodec_free_context(&_encoder_ctx);

    _encoder_ctx = NULL;
}

int encoder_video_x264::encode(AVFrame *frame, AVPacket *packet)
{
    int ret = avcodec_send_frame(_encoder_ctx, frame);
    if (ret < 0) {
        return AE_FFMPEG_ENCODE_FRAME_FAILED;
    }

    while (ret >= 0) {
        av_init_packet(packet);

        ret = avcodec_receive_packet(_encoder_ctx, packet);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            break;
        }

        if (ret < 0) {
            return AE_FFMPEG_READ_PACKET_FAILED;
        }

        if (ret == 0 && _on_data)
            _on_data(packet);

        av_packet_unref(packet);
    }

    return AE_NO;
}

void encoder_video_x264::encode_loop()
{
    AVPacket *packet = av_packet_alloc();
    AVFrame yuv_frame;

    int error = AE_NO;

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

        while (_ring_buffer->get(_buff, _buff_size, yuv_frame)) {
            // Normalize incoming frame timestamps to encoder time_base (1/fps)
            // Source pts is in AV_TIME_BASE (microseconds) from capturer
            if (yuv_frame.pts != AV_NOPTS_VALUE) {
                _frame->pts = av_rescale_q(yuv_frame.pts, AVRational{1, AV_TIME_BASE}, _encoder_ctx->time_base);
            } else {
                _frame->pts = AV_NOPTS_VALUE;
            }
            // Keep pkt_dts consistent with pts for no-B-frames encoders
            _frame->pkt_dts = _frame->pts;

            if ((error = encode(_frame, packet)) != AE_NO) {
                if (_on_error)
                    _on_error(error);

                al_fatal("encode 264 packet failed:%d", error);

                break;
            }
        }

        _cond_notify = false;
    }

    av_packet_free(&packet);
}
} // namespace am
