#include "video_encoder.h"

extern "C" {
#include <libavutil/opt.h>
}

#include <QDebug>

std::vector<std::string> Encoder<MediaType::VIDEO>::_usableEncoders;

Encoder<MediaType::VIDEO>::Encoder() {}

bool Encoder<MediaType::VIDEO>::Open(const Param& encodeParam, AVFormatContext* fmtCtx)
{
    Close();
    _isOpen = false;
    
    qDebug() << "VideoEncoder::Open: Initializing video encoder with codec" << QString::fromStdString(encodeParam.name)
             << "resolution" << encodeParam.width << "x" << encodeParam.height
             << "fps" << encodeParam.fps << "bitrate" << encodeParam.bitRate;
    
    if (!_Init(encodeParam, fmtCtx)) {
        qDebug() << "VideoEncoder::Open failed: Encoder initialization failed";
        return false;
    }

    // 打开编码器
    int ret = avcodec_open2(_codecCtx, _codec, nullptr);
    if (ret < 0) {
        char errbuf[AV_ERROR_MAX_STRING_SIZE];
        av_strerror(ret, errbuf, AV_ERROR_MAX_STRING_SIZE);
        qDebug() << "VideoEncoder::Open failed: Cannot open codec. Error:" << errbuf << "(" << ret << ")";
        qDebug() << "Codec name:" << _codec->name << "ID:" << _codec->id;
        return false;
    }

    _isOpen = true;
    qDebug() << "VideoEncoder::Open: Video encoder opened successfully";
    return true;
}

bool Encoder<MediaType::VIDEO>::PushFrame(AVFrame* frame, bool isEnd, uint64_t pts)
{
    if (!isEnd) {
        __CheckBool(_Trans(frame));
        frame = _bufferFrame;
        __CheckBool(frame);
    } else {
        frame = nullptr; // 直接刷新编码器缓存
    }
    if (frame != nullptr) {
        frame->pts = pts;
    }
    __CheckBool(avcodec_send_frame(_codecCtx, frame) >= 0);
    return true;
}

void Encoder<MediaType::VIDEO>::AfterEncode()
{
    if (_isHardware) {
        Free(_hwFrame, [this] { av_frame_free(&_hwFrame); });
    }
}

void Encoder<MediaType::VIDEO>::Close()
{
    if (_codecCtx != nullptr) {
        avcodec_free_context(&_codecCtx);
    }

    Free(_codecCtx, [this] { avcodec_free_context(&_codecCtx); });
    Free(_hwDeviceCtx, [this] { av_buffer_unref(&_hwDeviceCtx); });
    _converter = nullptr;
}

const std::vector<std::string>& Encoder<MediaType::VIDEO>::GetUsableEncoders()
{
    if (_usableEncoders.empty()) {
        _FindUsableEncoders();
        
        // 如果仍然没有找到可用编码器，至少添加libx264作为后备
        if (_usableEncoders.empty()) {
            qDebug() << "VideoEncoder: No encoders found during testing, adding libx264 as fallback";
            _usableEncoders.push_back("libx264");
        }
    }
    return _usableEncoders;
}

void Encoder<MediaType::VIDEO>::_FindUsableEncoders()
{
    qDebug() << "VideoEncoder: Finding usable encoders...";
    
    // 尝试打开编码器看看编码器能不能用
    Param param;
    param.bitRate = 1000;
    param.fps = 30;
    param.width = 1920;
    param.height = 1080;
    Encoder encoder;
    AVFormatContext* fmtCtx = nullptr;

    __CheckNo(avformat_alloc_output_context2(&fmtCtx, nullptr, nullptr, "test.mp4") >= 0);
    
    // 测试所有编码器，不做任何假设
    for (const auto& name : _encoderNames) {
        param.name = name;
        qDebug() << "VideoEncoder: Testing encoder:" << name;
        
        if (encoder.Open(param, fmtCtx)) {
            _usableEncoders.push_back(name);
            qDebug() << "VideoEncoder: Encoder" << name << "is usable";
        } else {
            qDebug() << "VideoEncoder: Encoder" << name << "failed to open";
        }
        encoder.Close();
    }
    
    qDebug() << "VideoEncoder: Found" << _usableEncoders.size() << "usable encoders";
    Free(fmtCtx, [&fmtCtx] { avformat_free_context(fmtCtx); });
}

bool Encoder<MediaType::VIDEO>::_Init(const Param& encodeParam, AVFormatContext* fmtCtx)
{
    _isHardware = (encodeParam.name != "libx264" && encodeParam.name != "libx265" && encodeParam.name != "libaom-av1");
    AVHWDeviceType hwType;
    if (encodeParam.name == "libx264" || encodeParam.name == "libx265" || encodeParam.name == "libaom-av1") {
        _pixFmt = AV_PIX_FMT_YUV420P;
    } else if (encodeParam.name == "h264_nvenc") {
        _pixFmt = AV_PIX_FMT_CUDA;
        hwType = AV_HWDEVICE_TYPE_CUDA;
    } else if (encodeParam.name == "h264_qsv") {
        _pixFmt = AV_PIX_FMT_QSV;
        hwType = AV_HWDEVICE_TYPE_QSV;
    } else if (encodeParam.name == "h264_amf") {
        _pixFmt = AV_PIX_FMT_D3D11;
        hwType = AV_HWDEVICE_TYPE_D3D11VA;
    } else if (encodeParam.name == "h264_videotoolbox") {
        _pixFmt = AV_PIX_FMT_VIDEOTOOLBOX;
        hwType = AV_HWDEVICE_TYPE_VIDEOTOOLBOX;
    }
    _isHardware = (_pixFmt != AV_PIX_FMT_YUV420P && _pixFmt != AV_PIX_FMT_NV12);
    if (_isHardware
        && av_hwdevice_ctx_create(&_hwDeviceCtx, hwType, nullptr, nullptr, 0) < 0) { // 硬件解码
        __DebugPrint("av_hwdevice_ctx_create failed\n");
        return false;
    }
    __CheckBool(_codec = avcodec_find_encoder_by_name(encodeParam.name.c_str()));
    __CheckBool(_codecCtx = avcodec_alloc_context3(_codec));
    _codecCtx->bit_rate = encodeParam.bitRate;
    _codecCtx->width = encodeParam.width;
    _codecCtx->height = encodeParam.height;
    _codecCtx->time_base = {1, encodeParam.fps};
    _codecCtx->framerate = {encodeParam.fps, 1};

    // 影响缓冲区大小
    _codecCtx->gop_size = 5;
    _codecCtx->max_b_frames = 0;
    _codecCtx->pix_fmt = _pixFmt;

    /* Some formats want stream headers to be separate. */
    if (fmtCtx->oformat->flags & AVFMT_GLOBALHEADER) {
        _codecCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
    }

    if (!_isHardware) { // 软件编码设置为快，避免占用过高的 CPU
        if (encodeParam.name == "libx264") {
            av_opt_set(_codecCtx->priv_data, "preset", "ultrafast", 0); // 使用最快的预设
            av_opt_set(_codecCtx->priv_data, "tune", "zerolatency", 0); // 零延迟调优
        } else if (encodeParam.name == "libx265") {
            av_opt_set(_codecCtx->priv_data, "preset", "ultrafast", 0);
            av_opt_set(_codecCtx->priv_data, "tune", "zerolatency", 0);
        } else if (encodeParam.name == "libaom-av1") {
            av_opt_set(_codecCtx->priv_data, "cpu-used", "8", 0); // 最快速度
            av_opt_set(_codecCtx->priv_data, "rt", "1", 0); // 实时模式
        }
    } else {
        // 为不同硬件编码器设置特定参数
        if (encodeParam.name == "h264_nvenc") {
            av_opt_set(_codecCtx->priv_data, "preset", "p4", 0); // NVENC预设
            av_opt_set(_codecCtx->priv_data, "rc", "vbr", 0);
        } else if (encodeParam.name == "h264_qsv") {
            av_opt_set(_codecCtx->priv_data, "preset", "medium", 0); // QSV有效预设
            av_opt_set(_codecCtx->priv_data, "look_ahead", "0", 0);
        } else if (encodeParam.name == "h264_amf") {
            av_opt_set(_codecCtx->priv_data, "quality", "speed", 0); // AMF预设
            av_opt_set(_codecCtx->priv_data, "rc", "vbr_peak", 0);
        }
    }

    __CheckBool(!_isHardware || _SetHwFrameCtx());
    return true;
}
bool Encoder<MediaType::VIDEO>::_SetHwFrameCtx()
{
    AVBufferRef* hwFramesRef;
    AVHWFramesContext* framesCtx = nullptr;

    __CheckBool(hwFramesRef = av_hwframe_ctx_alloc(_hwDeviceCtx));
    framesCtx = (AVHWFramesContext*) (hwFramesRef->data);
    framesCtx->format = _pixFmt;
    framesCtx->sw_format = AV_PIX_FMT_NV12;
    framesCtx->width = _codecCtx->width;
    framesCtx->height = _codecCtx->height;
    framesCtx->initial_pool_size = 20;
    if (av_hwframe_ctx_init(hwFramesRef) < 0) {
        __DebugPrint("av_hwframe_ctx_init failed\n");
        av_buffer_unref(&hwFramesRef);
        return false;
    }
    __CheckBool(_codecCtx->hw_frames_ctx = av_buffer_ref(hwFramesRef));
    av_buffer_unref(&hwFramesRef);
    return true;
}

bool Encoder<MediaType::VIDEO>::_Trans(AVFrame* frame)
{
    std::lock_guard<std::mutex> lk(__mtx);
    if (!_isOpen) {
        return false;
    }
    if (frame->format == AV_PIX_FMT_NV12) {
        _bufferFrame = frame;
    } else {
        if (_converter == nullptr) {
            _converter = std::make_unique<FfmpegConverter>(AVPixelFormat(frame->format),
                                                           AV_PIX_FMT_NV12);
            _converter->SetSize(frame->width, frame->height);
        }
        _bufferFrame = _converter->Trans(frame);
    }
    if (_isHardware) {
        _bufferFrame = _ToHardware();
    }
    __CheckBool(_bufferFrame);
    return true;
}

AVFrame* Encoder<MediaType::VIDEO>::_ToHardware()
{
    if (_bufferFrame == nullptr) {
        return nullptr;
    }
    __CheckNullptr(_hwFrame = av_frame_alloc());
    __CheckNullptr(av_hwframe_get_buffer(_codecCtx->hw_frames_ctx, _hwFrame, 0) >= 0);
    __CheckNullptr(_hwFrame->hw_frames_ctx);
    __CheckNullptr(av_hwframe_transfer_data(_hwFrame, _bufferFrame, 0) >= 0);
    return _hwFrame;
}