#include "video_encoder.h"

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

#include <QDebug>
#include <algorithm>

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) {
        if (!_Trans(frame)) { return false; }
        frame = _bufferFrame;
        if (!frame) { return false; }
    } else {
        frame = nullptr; // 直接刷新编码器缓存
    }
    if (frame != nullptr) {
        frame->pts = pts;
    }
    int sendRet = avcodec_send_frame(_codecCtx, frame);
    if (sendRet < 0) { __DebugPrint("avcodec_send_frame failed: %d", sendRet); return false; }
    return true;
}

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

void Encoder<MediaType::VIDEO>::Close()
{
    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();
    }
    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;

    int ctxRet = avformat_alloc_output_context2(&fmtCtx, nullptr, nullptr, "test.mp4");
    if (ctxRet < 0) { qDebug() << "VideoEncoder: Failed to alloc output context for testing encoders"; return; }
    
    // 测试所有编码器，不做任何假设
    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();
    }

    // 调整优先级排序：优先显卡（硬件）中兼容性较好的，然后再选择软件中兼容性最好的
    if (!_usableEncoders.empty()) {
        auto priority = [](const std::string& n) -> int {
            // 硬件优先（Windows常见顺序：NVIDIA -> Intel -> AMD；Apple放在后面）
            if (n == "h264_nvenc") return 0;        // NVIDIA 硬件
            if (n == "h264_qsv") return 1;          // Intel 硬件
            if (n == "h264_amf") return 2;          // AMD 硬件
            if (n == "h264_videotoolbox") return 3; // Apple 硬件
            // 软件其次
            if (n == "libx264") return 4;           // 软件中兼容性最好
            if (n == "libx265") return 5;           // 兼容性一般
            if (n == "libaom-av1") return 6;        // 兼容性较差/更慢
            return 100;                                // 未知放最后
        };
        std::sort(_usableEncoders.begin(), _usableEncoders.end(), [&](const std::string& a, const std::string& b) {
            int pa = priority(a);
            int pb = priority(b);
            if (pa != pb) return pa < pb;
            return a < b; // 同优先级按名称稳定排序
        });
    }
    
    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;
    }
    _codec = avcodec_find_encoder_by_name(encodeParam.name.c_str());
    if (!_codec) { __DebugPrint("avcodec_find_encoder_by_name failed: %s", encodeParam.name.c_str()); return false; }
    _codecCtx = avcodec_alloc_context3(_codec);
    if (!_codecCtx) { __DebugPrint("avcodec_alloc_context3 failed"); return false; }
    _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);
        }
    }

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

    hwFramesRef = av_hwframe_ctx_alloc(_hwDeviceCtx);
    if (!hwFramesRef) { __DebugPrint("av_hwframe_ctx_alloc failed"); return false; }
    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;
    }
    _codecCtx->hw_frames_ctx = av_buffer_ref(hwFramesRef);
    if (!_codecCtx->hw_frames_ctx) { __DebugPrint("av_buffer_ref failed"); av_buffer_unref(&hwFramesRef); return false; }
    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) {
            // Convert from source pixel format to NV12 for downstream processing
            _converter = std::make_unique<FfmpegConverter>(static_cast<AVPixelFormat>(frame->format), AV_PIX_FMT_NV12);
            _converter->SetSize(frame->width, frame->height);
        }
        _bufferFrame = _converter->Trans(frame);
    }
    if (_isHardware) {
        _bufferFrame = _ToHardware();
        if (!_bufferFrame) { return false; }
    }
    return true;
}

AVFrame* Encoder<MediaType::VIDEO>::_ToHardware()
{
    if (!_hwDeviceCtx) { __DebugPrint("_hwDeviceCtx is null"); return nullptr; }
    if (!(_codecCtx && _codecCtx->hw_frames_ctx)) { __DebugPrint("hw_frames_ctx is null"); return nullptr; }
    if (_hwFrame == nullptr) {
        _hwFrame = av_frame_alloc();
        if (!_hwFrame) { __DebugPrint("av_frame_alloc failed"); return nullptr; }
        _hwFrame->format = ((AVHWFramesContext*)_codecCtx->hw_frames_ctx->data)->format;
        _hwFrame->width = _codecCtx->width;
        _hwFrame->height = _codecCtx->height;
        if (av_hwframe_get_buffer(_codecCtx->hw_frames_ctx, _hwFrame, 0) < 0) { __DebugPrint("av_hwframe_get_buffer failed"); return nullptr; }
    }
    if (av_hwframe_transfer_data(_hwFrame, _bufferFrame, 0) < 0) {
        __DebugPrint("hw transfer failed\n");
        return nullptr;
    }
    return _hwFrame;
}
