#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;

    // 回到超低延迟配置，但优化画质参数
    av_dict_set(&options, "preset", "ultrafast", 0);
    av_dict_set(&options, "tune", "zerolatency", 0);
    
    // 针对花屏问题的特殊优化：更精细的x264参数控制
    av_dict_set(&options, "x264-params", "keyint=15:min-keyint=15:no-scenecut=1:rc-lookahead=0:aq-mode=3:aq-strength=1.0:deblock=1,1:psy-rd=0.4,0.0:me=umh:subme=4:ref=1", 0);
    av_dict_set(&options, "profile", "high", 0);
    av_dict_set(&options, "level", "4.1", 0);
    // 使用更高的码率而不是更低的CRF来改善画质
    av_dict_set(&options, "crf", "20", 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 * 1.5;  // 适度提高码率，避免过度消耗带宽

        // 回到低延迟GOP设置
        _encoder_ctx->gop_size = 15;

        // 针对花屏优化的量化参数
        _encoder_ctx->qmin = 10;   // 稍微提高最小质量
        _encoder_ctx->qmax = 24;   // 控制最大量化，避免过度压缩
        
        // 设置CRF值 - 针对花屏问题优化
        _encoder_ctx->global_quality = 18 * FF_QP2LAMBDA;  // 降低到18提升画质

        _encoder_ctx->max_b_frames = 0; //NO B Frame
        
        // 恢复最低延迟设置
        _encoder_ctx->delay = 0;
        _encoder_ctx->thread_count = 1;  // 单线程确保最低延迟
        
        _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);
        // 恢复最激进的低延迟处理
        if (!_cond_notify && _ring_buffer->get_pending_frames() == 0) {
            // 使用最短等待时间
            _cond_var.wait_for(lock, std::chrono::milliseconds(1), [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