GrabBag/Module/FFMediaStream/Src/CVrFFMediaPuller.cpp

#include "CVrFFMediaPuller.h"
#include "VrError.h"
#include "VrLog.h"

#include <cstring>

#if !defined(FFMEDIA_PLATFORM_DUMMY)
#include <linux/videodev2.h>
#endif

bool IVrFFMediaPuller::CreateObject(IVrFFMediaPuller** ppPuller)
{
    if (!ppPuller) return false;
    try { *ppPuller = new CVrFFMediaPuller(); return true; }
    catch (...) { return false; }
}

CVrFFMediaPuller::CVrFFMediaPuller() = default;

CVrFFMediaPuller::~CVrFFMediaPuller()
{
    try { UnInit(); } catch (...) {}
}

void CVrFFMediaPuller::SetFrameCallback(FFPulledFrameCallback callback)
{
    std::lock_guard<std::mutex> lk(m_mutex);
    m_callback = std::move(callback);
}

bool CVrFFMediaPuller::IsStreaming() const
{
    return m_bStarted.load();
}

#if defined(FFMEDIA_PLATFORM_DUMMY)

// =================== Windows / 非 RK3588 占位实现 ===================

int CVrFFMediaPuller::Init(const VrFFPullConfig& config)
{
    std::lock_guard<std::mutex> lk(m_mutex);
    m_config = config;
    m_bInited = true;
    LOG_INFO("FFMediaPuller: Init (dummy)\n");
    return SUCCESS;
}

int CVrFFMediaPuller::Start()
{
    if (!m_bInited.load()) return ERR_CODE(DEV_NO_OPEN);
    m_bStarted = true;
    LOG_INFO("FFMediaPuller: Start (dummy)\n");
    return SUCCESS;
}

int CVrFFMediaPuller::Stop()
{
    m_bStarted = false;
    return SUCCESS;
}

int CVrFFMediaPuller::UnInit()
{
    Stop();
    m_bInited = false;
    return SUCCESS;
}

#else

// =================== RK3588 真正实现 MPP + FFmpeg-rockchip ===================

// ---- 辅助函数 ----

MppCodingType CVrFFMediaPuller::AvCodecIdToMppCoding(AVCodecID id)
{
    switch (id)
    {
    case AV_CODEC_ID_H264:  return MPP_VIDEO_CodingAVC;
    case AV_CODEC_ID_H265:  return MPP_VIDEO_CodingHEVC;
    case AV_CODEC_ID_MJPEG: return MPP_VIDEO_CodingMJPEG;
    default:                return MPP_VIDEO_CodingAVC;
    }
}

int CVrFFMediaPuller::ToAvPixelFormat(EFFPullPixelFormat fmt)
{
    switch (fmt)
    {
    case EFFPullPixelFormat::NV12:   return AV_PIX_FMT_NV12;
    case EFFPullPixelFormat::RGBA32: return AV_PIX_FMT_RGBA;
    case EFFPullPixelFormat::BGR24:  return AV_PIX_FMT_BGR24;
    }
    return AV_PIX_FMT_RGBA;
}

unsigned int CVrFFMediaPuller::PixelBpp(EFFPullPixelFormat fmt)
{
    switch (fmt)
    {
    case EFFPullPixelFormat::NV12:   return 1;
    case EFFPullPixelFormat::RGBA32: return 4;
    case EFFPullPixelFormat::BGR24:  return 3;
    }
    return 4;
}

// ---- MPP 解码器初始化 ----

int CVrFFMediaPuller::initMppDecoder()
{
    if (!m_fmtCtx) return ERR_CODE(DEV_OPEN_ERR);

    AVCodecParameters* codecpar = m_fmtCtx->streams[m_videoStreamIndex]->codecpar;
    MppCodingType coding = AvCodecIdToMppCoding(codecpar->codec_id);

    // 1) 创建 MPP 上下文
    MPP_RET ret = mpp_create(&m_mppCtx, &m_mppApi);
    if (ret != MPP_SUCCESS || !m_mppCtx || !m_mppApi)
    {
        LOG_ERROR("FFMediaPuller: mpp_create fail %d\n", ret);
        return ERR_CODE(DEV_OPEN_ERR);
    }

    // 2) 启用 parser split 模式（自动分割 Annex B 流中的 SPS/PPS/Slice）
    RK_U32 needSplit = 1;
    ret = m_mppApi->control(m_mppCtx, MPP_DEC_SET_PARSER_SPLIT_MODE, &needSplit);
    if (ret != MPP_SUCCESS)
    {
        LOG_ERROR("FFMediaPuller: set parser split mode fail %d\n", ret);
        // 非致命，继续
    }

    ret = mpp_init(m_mppCtx, MPP_CTX_DEC, coding);
    if (ret != MPP_SUCCESS)
    {
        LOG_ERROR("FFMediaPuller: mpp_init fail %d\n", ret);
        return ERR_CODE(DEV_OPEN_ERR);
    }

    // 3) 发送 extradata（SPS/PPS）到解码器
    if (codecpar->extradata && codecpar->extradata_size > 0)
    {
        MppPacket cfgPkt = nullptr;
        mpp_packet_init(&cfgPkt, codecpar->extradata,
                        static_cast<size_t>(codecpar->extradata_size));
        mpp_packet_set_extra_data(cfgPkt);

        ret = m_mppApi->decode_put_packet(m_mppCtx, cfgPkt);
        if (ret == MPP_SUCCESS)
        {
            MppFrame infoFrame = nullptr;
            ret = m_mppApi->decode_get_frame(m_mppCtx, &infoFrame);
            if (ret == MPP_SUCCESS && infoFrame)
            {
                if (mpp_frame_get_info_change(infoFrame))
                {
                    m_decWidth     = static_cast<int>(mpp_frame_get_width(infoFrame));
                    m_decHeight    = static_cast<int>(mpp_frame_get_height(infoFrame));
                    m_decHorStride = static_cast<int>(mpp_frame_get_hor_stride(infoFrame));
                    m_decVerStride = static_cast<int>(mpp_frame_get_ver_stride(infoFrame));
                    LOG_INFO("FFMediaPuller: decoder info_change %dx%d str=%dx%d\n",
                             m_decWidth, m_decHeight, m_decHorStride, m_decVerStride);
                }
                mpp_frame_deinit(&infoFrame);
            }
            m_mppApi->control(m_mppCtx, MPP_DEC_SET_INFO_CHANGE_READY, nullptr);
        }
        mpp_packet_deinit(&cfgPkt);
    }

    return SUCCESS;
}

// ---- 格式转换初始化 ----

int CVrFFMediaPuller::initSwscale()
{
    // 计算输出尺寸
    m_outWidth  = (m_config.outputWidth  > 0) ? static_cast<int>(m_config.outputWidth)  : m_decWidth;
    m_outHeight = (m_config.outputHeight > 0) ? static_cast<int>(m_config.outputHeight) : m_decHeight;

    m_needConvert = (m_config.outputFormat != EFFPullPixelFormat::NV12);

    if (!m_needConvert)
        return SUCCESS;

    int outFmt = ToAvPixelFormat(m_config.outputFormat);
    m_outStride = m_outWidth * static_cast<int>(PixelBpp(m_config.outputFormat));

    m_swsCtx = sws_getContext(m_decWidth, m_decHeight, AV_PIX_FMT_NV12,
                              m_outWidth, m_outHeight, static_cast<AVPixelFormat>(outFmt),
                              SWS_BILINEAR, nullptr, nullptr, nullptr);
    if (!m_swsCtx)
    {
        LOG_ERROR("FFMediaPuller: sws_getContext fail\n");
        return ERR_CODE(DEV_OPEN_ERR);
    }

    // 分配转换输出缓冲区
    m_outBufSize = m_outHeight * m_outStride;
    m_outBuffer = static_cast<uint8_t*>(av_malloc(m_outBufSize + AV_INPUT_BUFFER_PADDING_SIZE));
    if (!m_outBuffer)
    {
        LOG_ERROR("FFMediaPuller: out buffer alloc fail\n");
        return ERR_CODE(DEV_OPEN_ERR);
    }

    LOG_INFO("FFMediaPuller: swscale enabled (NV12 → pix %d, %dx%d → %dx%d)\n",
             outFmt, m_decWidth, m_decHeight, m_outWidth, m_outHeight);
    return SUCCESS;
}

// ---- 解码线程 ----

void CVrFFMediaPuller::decodeLoop()
{
    LOG_INFO("FFMediaPuller: decode thread started\n");

    AVPacket* avpkt = av_packet_alloc();
    if (!avpkt)
    {
        LOG_ERROR("FFMediaPuller: av_packet_alloc fail\n");
        return;
    }

    while (m_bRunning.load())
    {
        int ret = av_read_frame(m_fmtCtx, avpkt);
        if (ret < 0)
        {
            if (ret == AVERROR_EOF)
            {
                LOG_INFO("FFMediaPuller: stream EOF\n");
            }
            else if (ret != AVERROR_EXIT && m_bRunning.load())
            {
                LOG_DEBUG("FFMediaPuller: av_read_frame err %d, retrying...\n", ret);
                av_packet_unref(avpkt);
                // 短暂休眠后重试
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
                continue;
            }
            av_packet_unref(avpkt);
            break;
        }

        // 仅处理视频流
        if (avpkt->stream_index != m_videoStreamIndex)
        {
            av_packet_unref(avpkt);
            continue;
        }

        // 1) 构造 MppPacket
        MppPacket mppPkt = nullptr;
        mpp_packet_init(&mppPkt, avpkt->data, static_cast<size_t>(avpkt->size));
        mpp_packet_set_pts(mppPkt, avpkt->pts);

        // 2) 同步解码
        MppFrame frame = nullptr;
        MPP_RET mret = m_mppApi->decode(m_mppCtx, mppPkt, &frame);
        mpp_packet_deinit(&mppPkt);

        if (mret != MPP_SUCCESS || !frame)
        {
            av_packet_unref(avpkt);
            continue;
        }

        // 检查 info_change（中间可能发生分辨率变化）
        if (mpp_frame_get_info_change(frame))
        {
            m_decWidth     = static_cast<int>(mpp_frame_get_width(frame));
            m_decHeight    = static_cast<int>(mpp_frame_get_height(frame));
            m_decHorStride = static_cast<int>(mpp_frame_get_hor_stride(frame));
            m_decVerStride = static_cast<int>(mpp_frame_get_ver_stride(frame));
            m_mppApi->control(m_mppCtx, MPP_DEC_SET_INFO_CHANGE_READY, nullptr);
            LOG_INFO("FFMediaPuller: mid-stream info_change %dx%d\n", m_decWidth, m_decHeight);
            mpp_frame_deinit(&frame);
            av_packet_unref(avpkt);
            continue;
        }

        // 获取解码数据
        MppBuffer buf = mpp_frame_get_buffer(frame);
        if (!buf)
        {
            mpp_frame_deinit(&frame);
            av_packet_unref(avpkt);
            continue;
        }

        void*   decData    = mpp_buffer_get_ptr(buf);
        int     decW       = static_cast<int>(mpp_frame_get_width(frame));
        int     decH       = static_cast<int>(mpp_frame_get_height(frame));
        int     decHStride = static_cast<int>(mpp_frame_get_hor_stride(frame));
        int     decVStride = static_cast<int>(mpp_frame_get_ver_stride(frame));
        int64_t decPts     = mpp_frame_get_pts(frame);

        if (!decData || decW <= 0 || decH <= 0)
        {
            mpp_frame_deinit(&frame);
            av_packet_unref(avpkt);
            continue;
        }

        // 3) 构建回调帧
        VrFFPulledFrame pulledFrame;
        pulledFrame.ptsUs  = decPts;
        pulledFrame.width  = static_cast<unsigned int>(decW);
        pulledFrame.height = static_cast<unsigned int>(decH);

        if (m_needConvert && m_swsCtx && m_outBuffer)
        {
            // NV12 → RGBA/BGR；UV 平面偏移必须用 ver_stride（通常 > height，如 1088>1080）
            uint8_t* srcSlice[4] = {
                static_cast<uint8_t*>(decData),
                static_cast<uint8_t*>(decData) + static_cast<size_t>(decHStride) * decVStride,
                nullptr, nullptr
            };
            int srcStride[4] = { decHStride, decHStride, 0, 0 };

            uint8_t* dstSlice[4] = { m_outBuffer, nullptr, nullptr, nullptr };
            int dstStride[4] = { m_outStride, 0, 0, 0 };

            sws_scale(m_swsCtx, srcSlice, srcStride, 0, decH,
                      dstSlice, dstStride);

            pulledFrame.data   = m_outBuffer;
            pulledFrame.size   = m_outBufSize;
            pulledFrame.width  = static_cast<unsigned int>(m_outWidth);
            pulledFrame.height = static_cast<unsigned int>(m_outHeight);
            pulledFrame.stride = static_cast<unsigned int>(m_outStride);
            pulledFrame.format = m_config.outputFormat;
        }
        else
        {
            // 直接输出 NV12（零拷贝，buffer 在回调期间有效；偏移按 ver_stride）
            pulledFrame.data   = decData;
            pulledFrame.size   = static_cast<size_t>(decHStride) * decVStride * 3 / 2;
            pulledFrame.stride = static_cast<unsigned int>(decHStride);
            pulledFrame.format = EFFPullPixelFormat::NV12;
        }

        // 4) 通知回调
        {
            FFPulledFrameCallback cb;
            {
                std::lock_guard<std::mutex> lk(m_mutex);
                cb = m_callback;
            }
            if (cb)
            {
                try { cb(pulledFrame); } catch (...) {}
            }
        }

        mpp_frame_deinit(&frame);
        av_packet_unref(avpkt);
    }

    av_packet_free(&avpkt);
    LOG_INFO("FFMediaPuller: decode thread stopped\n");
}

// ---- 主流程 ----

int CVrFFMediaPuller::Init(const VrFFPullConfig& config)
{
    std::lock_guard<std::mutex> lk(m_mutex);
    if (m_bInited.load()) return SUCCESS;

    m_config = config;

    // 1) 初始化 FFmpeg 网络 + 降低日志噪声
    avformat_network_init();
    av_log_set_level(AV_LOG_ERROR);

    // 2) 打开 RTSP 流
    AVDictionary* opts = nullptr;
    if (config.useTcp)
    {
        av_dict_set(&opts, "rtsp_transport", "tcp", 0);
    }
    else
    {
        av_dict_set(&opts, "rtsp_transport", "udp", 0);
    }
    // socket 读超时（微秒）。FFmpeg 6.x RTSP demuxer 用 "timeout"；
    // 同时设置旧名 "stimeout" 以兼容（未知项被忽略，无害）。
    av_dict_set(&opts, "timeout",  "5000000", 0);
    av_dict_set(&opts, "stimeout", "5000000", 0);
    // 降低拉流端到端延迟
    av_dict_set(&opts, "max_delay", "500000", 0);

    int ret = avformat_open_input(&m_fmtCtx, config.rtspUrl.c_str(), nullptr, &opts);
    av_dict_free(&opts);
    if (ret < 0)
    {
        LOG_ERROR("FFMediaPuller: avformat_open_input fail %d (%s)\n",
                  ret, config.rtspUrl.c_str());
        return ERR_CODE(DEV_OPEN_ERR);
    }

    // 3) 获取流信息
    ret = avformat_find_stream_info(m_fmtCtx, nullptr);
    if (ret < 0)
    {
        LOG_ERROR("FFMediaPuller: avformat_find_stream_info fail %d\n", ret);
        avformat_close_input(&m_fmtCtx);
        return ERR_CODE(DEV_OPEN_ERR);
    }

    // 4) 查找视频流
    m_videoStreamIndex = -1;
    for (unsigned int i = 0; i < m_fmtCtx->nb_streams; ++i)
    {
        if (m_fmtCtx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            m_videoStreamIndex = static_cast<int>(i);
            break;
        }
    }
    if (m_videoStreamIndex < 0)
    {
        LOG_ERROR("FFMediaPuller: no video stream found\n");
        avformat_close_input(&m_fmtCtx);
        return ERR_CODE(DEV_OPEN_ERR);
    }

    AVCodecParameters* codecpar = m_fmtCtx->streams[m_videoStreamIndex]->codecpar;
    LOG_INFO("FFMediaPuller: video stream #%d, codec=%d, %dx%d\n",
             m_videoStreamIndex, codecpar->codec_id,
             codecpar->width, codecpar->height);

    // 5) 初始化解码器默认尺寸（可能被 info_change 更新）
    m_decWidth     = codecpar->width;
    m_decHeight    = codecpar->height;
    m_decHorStride = codecpar->width;   // 先猜测，info_change 后会更新
    m_decVerStride = codecpar->height;

    // 6) 初始化 MPP 解码器
    ret = initMppDecoder();
    if (ret != SUCCESS)
    {
        LOG_ERROR("FFMediaPuller: initMppDecoder fail\n");
        avformat_close_input(&m_fmtCtx);
        return ret;
    }

    // 7) 初始化格式转换
    ret = initSwscale();
    if (ret != SUCCESS)
    {
        LOG_ERROR("FFMediaPuller: initSwscale fail\n");
        avformat_close_input(&m_fmtCtx);
        return ret;
    }

    m_bInited = true;
    LOG_INFO("FFMediaPuller: init done (%dx%d dec → %dx%d out)\n",
             m_decWidth, m_decHeight, m_outWidth, m_outHeight);
    return SUCCESS;
}

int CVrFFMediaPuller::Start()
{
    if (!m_bInited.load()) return ERR_CODE(DEV_NO_OPEN);
    if (m_bStarted.load()) return SUCCESS;

    // 设置中断回调（用于 Stop 时安全退出阻塞的 av_read_frame）
    m_bStopRunning = false;   // 复位，支持重启
    if (m_fmtCtx)
    {
        m_fmtCtx->interrupt_callback.callback = [](void* ctx) -> int {
            return *(static_cast<std::atomic<bool>*>(ctx)) ? 1 : 0;
        };
        m_fmtCtx->interrupt_callback.opaque = &m_bStopRunning;
    }

    m_bRunning = true;
    m_decodeThread = std::thread(&CVrFFMediaPuller::decodeLoop, this);

    m_bStarted = true;
    return SUCCESS;
}

int CVrFFMediaPuller::Stop()
{
    if (!m_bStarted.load()) return SUCCESS;

    m_bStopRunning = true;  // 中断 FFmpeg 阻塞调用
    m_bRunning = false;     // 退出解码循环

    if (m_decodeThread.joinable())
    {
        m_decodeThread.join();
    }

    m_bStarted = false;
    return SUCCESS;
}

int CVrFFMediaPuller::UnInit()
{
    Stop();
    std::lock_guard<std::mutex> lk(m_mutex);

    // 释放 FFmpeg 输入上下文
    if (m_fmtCtx)
    {
        avformat_close_input(&m_fmtCtx);
        m_fmtCtx = nullptr;
    }

    // 销毁 MPP 解码器上下文
    if (m_mppCtx)
    {
        mpp_destroy(m_mppCtx);
        m_mppCtx = nullptr;
        m_mppApi = nullptr;
    }

    // 释放 swscale 上下文
    if (m_swsCtx)
    {
        sws_freeContext(m_swsCtx);
        m_swsCtx = nullptr;
    }

    // 释放输出缓冲区
    if (m_outBuffer)
    {
        av_free(m_outBuffer);
        m_outBuffer = nullptr;
        m_outBufSize = 0;
    }

    // 关闭 FFmpeg 网络
    avformat_network_deinit();

    m_bInited = false;
    return SUCCESS;
}

#endif // FFMEDIA_PLATFORM_DUMMY