/**
 * RV1106 H.264 硬件编码器 — RKMPI (librockit) 实现
 *
 * 依据 Luckfox 官方 Wiki (RKMPI) 与 SDK samples:
 *   media/samples/example/common/sample_comm_venc.c
 *   media/rockit/rockit/mpi/example/mod/test_mpi_venc.cpp
 *
 * 数据通路:
 *   RK_MPI_SYS_Init
 *   -> RK_MPI_MB_CreatePool (输入 NV12 帧, DMA + cached)
 *   -> RK_MPI_VENC_CreateChn (H264 CBR)
 *   -> RK_MPI_VENC_StartRecvFrame
 *   [feed] RK_MPI_MB_GetMB -> memcpy -> MmzFlushCache -> RK_MPI_VENC_SendFrame(-1) -> ReleaseMB
 *   [取流线程] RK_MPI_VENC_GetStream(500ms) -> cb 回调 -> RK_MPI_VENC_ReleaseStream
 *   [销毁] StopRecvFrame -> DestroyChn -> DestroyPool -> RK_MPI_SYS_WaitFreeMB -> SYS_Exit
 *
 * 注意: RV1106 的 venc540c 是内核态编码器 (rockit.ko), 必须用 RKMPI 接口,
 *       不能用老 MPP 的 mpi->encode_put_frame 用户态接口 (会 check_is_mpp_packet 崩溃)。
 */

#include "mpp_encoder.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>
#include <unistd.h>

#include "rk_mpi_sys.h"
#include "rk_mpi_venc.h"
#include "rk_mpi_mb.h"

struct mpp_encoder {
    int width, height;
    int hor_stride, ver_stride;   /* NV12 输入 stride (16 对齐) */
    int fps;
    int bitrate_kbps;
    int gop;
    RK_U32 profile_idc;           /* 66/77/100 */
    char profile_name[16];

    VENC_CHN        chn;
    VENC_CHN_ATTR_S attr;         /* 保存用于动态改码率 */
    MB_POOL         pool;         /* 输入帧池 */
    size_t          frame_size;   /* 单帧 NV12 字节 */

    nalu_callback_t cb;
    void           *userdata;

    pthread_t       stream_th;
    volatile int    stopping;
    int             th_started;

    int sys_inited;               /* 本实例是否执行了 RK_MPI_SYS_Init */
    int pool_ok;                  /* MB 池创建成功 (经 GetMB 探测验证) */
    int chn_created;              /* VENC 通道创建成功 */
};

/* ---------- 工具 ---------- */

static RK_U32 profile_to_idc(const char *p){
    if(!p) return 100;
    if(!strcasecmp(p, "baseline")) return 66;
    if(!strcasecmp(p, "main"))     return 77;
    return 100; /* high */
}

/* ---------- 取流线程 (官方 test_mpi_venc 模式) ---------- */

static void *stream_thread(void *arg){
    struct mpp_encoder *e = (struct mpp_encoder *)arg;
    VENC_STREAM_S st;
    memset(&st, 0, sizeof(st));
    st.pstPack = (VENC_PACK_S *)malloc(sizeof(VENC_PACK_S));
    if(!st.pstPack){
        fprintf(stderr, "[venc] allocate stream pack failed\n");
        return NULL;
    }
    st.u32PackCount = 1;

    while(!e->stopping){
        st.u32PackCount = 1;
        RK_S32 ret = RK_MPI_VENC_GetStream(e->chn, &st, 500); /* 500ms 超时, 周期检查退出标志 */
        if(ret != 0){
            continue; /* 超时或暂无数据 */
        }
        RK_U32 cnt = st.u32PackCount;
        if(cnt > 1) cnt = 1; /* 只分配了 1 个 pack 结构 */
        for(RK_U32 i = 0; i < cnt; i++){
            VENC_PACK_S *pk = &st.pstPack[i];
            if(!pk->pMbBlk || pk->u32Len == 0) continue;

            /* VENC writes the MB from the device side. Invalidate the CPU cache
             * before reading it. On RV1106 pMbBlk already addresses this pack;
             * u32Offset is the position in VENC's ring buffer, not an offset to
             * add to RK_MPI_MB_Handle2VirAddr(pMbBlk). */
            RK_MPI_SYS_MmzFlushCache(pk->pMbBlk, RK_TRUE);
            RK_U64 mb_size = RK_MPI_MB_GetSize(pk->pMbBlk);
            if((RK_U64)pk->u32Len > mb_size){
                fprintf(stderr,
                        "[venc] invalid pack bounds seq=%u offset=%u len=%u mb=%llu\n",
                        st.u32Seq, pk->u32Offset, pk->u32Len,
                        (unsigned long long)mb_size);
                continue;
            }
            const uint8_t *data =
                (const uint8_t *)RK_MPI_MB_Handle2VirAddr(pk->pMbBlk);
            if(data && e->cb){
                int is_key = (pk->DataType.enH264EType == H264E_NALU_IDRSLICE);
                e->cb(data, pk->u32Len, pk->u64PTS, is_key, e->userdata);
            }
        }
        RK_MPI_VENC_ReleaseStream(e->chn, &st);
    }
    free(st.pstPack);
    return NULL;
}

/* ---------- 生命周期 ---------- */

mpp_encoder_t *mpp_encoder_create(const mpp_encoder_cfg_t *cfg){
    if(!cfg) return NULL;

    struct mpp_encoder *e = (struct mpp_encoder *)calloc(1, sizeof(*e));
    if(!e) return NULL;

    e->width  = cfg->width;
    e->height = cfg->height;
    e->fps    = cfg->fps > 0 ? cfg->fps : 24;
    e->bitrate_kbps = cfg->bitrate_kbps > 0 ? cfg->bitrate_kbps : 1500;
    e->gop    = cfg->gop > 0 ? cfg->gop : (e->fps * 2);
    e->profile_idc = profile_to_idc(cfg->profile);
    snprintf(e->profile_name, sizeof(e->profile_name), "%s",
             e->profile_idc == 66 ? "baseline" : e->profile_idc == 77 ? "main" : "high");

    /* 16 对齐 stride (上层按此布局 NV12) */
    e->hor_stride = (e->width  + 15) & ~15;
    e->ver_stride = (e->height + 15) & ~15;
    e->frame_size = (size_t)e->hor_stride * e->ver_stride * 3 / 2;

    /* 1. RKMPI 系统初始化 (rockit.ko 必须已加载) */
    RK_S32 ret = RK_MPI_SYS_Init();
    if(ret != 0){
        fprintf(stderr, "[venc] RK_MPI_SYS_Init fail %d (检查 rockit.ko: lsmod | grep rockit)\n", ret);
        free(e);
        return NULL;
    }
    e->sys_inited = 1;

    /* 2. 输入帧 MB 池 (DMA; 与官方 test_mpi_venc 一致, 不设 remap mode)
     *    注意: MB_POOL 是 RK_U32, 池 id 0 可能合法, 不能用 if(!pool) 判失败
     *    -> 用试取一块 MB 探测, 拿得到才算创建成功 */
    MB_POOL_CONFIG_S poolCfg;
    memset(&poolCfg, 0, sizeof(poolCfg));
    poolCfg.u64MBSize     = e->frame_size;
    poolCfg.u32MBCnt      = 3;
    poolCfg.enAllocType   = MB_ALLOC_TYPE_DMA;
    poolCfg.bPreAlloc     = RK_TRUE;
    e->pool = RK_MPI_MB_CreatePool(&poolCfg);
    MB_BLK probe = RK_MPI_MB_GetMB(e->pool, e->frame_size, RK_FALSE);
    if(!probe){
        fprintf(stderr, "[venc] MB_CreatePool/GetMB fail (pool=%u)\n", (unsigned)e->pool);
        goto fail;
    }
    RK_MPI_MB_ReleaseMB(probe);
    e->pool_ok = 1;

    /* 3. VENC 通道属性 (H264 CBR, 样板 sample_comm_venc.c) */
    memset(&e->attr, 0, sizeof(e->attr));
    e->attr.stVencAttr.enType          = RK_VIDEO_ID_AVC;
    e->attr.stVencAttr.u32MaxPicWidth  = e->width;
    e->attr.stVencAttr.u32MaxPicHeight = e->height;
    e->attr.stVencAttr.u32PicWidth     = e->width;
    e->attr.stVencAttr.u32PicHeight    = e->height;
    e->attr.stVencAttr.u32VirWidth     = e->hor_stride;
    e->attr.stVencAttr.u32VirHeight    = e->ver_stride;
    e->attr.stVencAttr.enPixelFormat   = RK_FMT_YUV420SP;        /* NV12 */
    e->attr.stVencAttr.u32Profile      = e->profile_idc;
    e->attr.stVencAttr.bByFrame        = RK_TRUE;
    e->attr.stVencAttr.u32StreamBufCnt = 3;
    e->attr.stVencAttr.u32BufSize      = (RK_U32)e->frame_size;

    e->attr.stRcAttr.enRcMode = VENC_RC_MODE_H264CBR;
    e->attr.stRcAttr.stH264Cbr.u32Gop               = e->gop;
    e->attr.stRcAttr.stH264Cbr.u32BitRate           = e->bitrate_kbps;  /* 单位 kbps */
    e->attr.stRcAttr.stH264Cbr.u32SrcFrameRateNum   = e->fps;
    e->attr.stRcAttr.stH264Cbr.u32SrcFrameRateDen   = 1;
    e->attr.stRcAttr.stH264Cbr.fr32DstFrameRateNum  = e->fps;
    e->attr.stRcAttr.stH264Cbr.fr32DstFrameRateDen  = 1;
    e->attr.stRcAttr.stH264Cbr.u32StatTime          = 3;

    e->chn = 0;
    ret = RK_MPI_VENC_CreateChn(e->chn, &e->attr);
    if(ret != 0){
        fprintf(stderr, "[venc] CreateChn fail %d\n", ret);
        goto fail;
    }
    e->chn_created = 1;

    /* 4. 启动收帧 */
    VENC_RECV_PIC_PARAM_S recv;
    memset(&recv, 0, sizeof(recv));
    recv.s32RecvPicNum = -1;   /* 不限帧数 */
    ret = RK_MPI_VENC_StartRecvFrame(e->chn, &recv);
    if(ret != 0){
        fprintf(stderr, "[venc] StartRecvFrame fail %d\n", ret);
        goto fail;
    }

    /* 5. 取流线程 */
    e->stopping = 0;
    if(pthread_create(&e->stream_th, NULL, stream_thread, e) != 0){
        fprintf(stderr, "[venc] stream thread create fail\n");
        goto fail;
    }
    e->th_started = 1;

    fprintf(stderr, "[venc] encoder ready %dx%d@%d %dkbps profile=%s gop=%d stride=%dx%d (RKMPI)\n",
            e->width, e->height, e->fps, e->bitrate_kbps, e->profile_name, e->gop,
            e->hor_stride, e->ver_stride);
    return e;

fail:
    if(e->th_started){ e->stopping = 1; pthread_join(e->stream_th, NULL); }
    /* 只清理真正创建过的资源。
     * 之前对未创建的通道调 DestroyChn(0) 触发 mpp 内部读垃圾内存:
     * "h265:strg_intra_smth_disable to -13xxxxx" + "chan_id 0 is already valid" */
    if(e->chn_created){ RK_MPI_VENC_StopRecvFrame(e->chn); RK_MPI_VENC_DestroyChn(e->chn); }
    if(e->pool_ok)   { RK_MPI_MB_DestroyPool(e->pool); }
    if(e->sys_inited){ RK_MPI_SYS_WaitFreeMB(); RK_MPI_SYS_Exit(); }
    free(e);
    return NULL;
}

void mpp_encoder_destroy(mpp_encoder_t *e){
    if(!e) return;
    e->stopping = 1;
    /* StopRecvFrame also wakes a GetStream call that is waiting in the
     * output thread, so Ctrl+C does not have to wait for the driver timeout. */
    if(e->chn_created){
        RK_MPI_VENC_StopRecvFrame(e->chn);
    }
    if(e->th_started){
        pthread_join(e->stream_th, NULL);
    }
    if(e->chn_created){
        RK_MPI_VENC_DestroyChn(e->chn);
    }
    if(e->pool_ok){
        RK_MPI_MB_DestroyPool(e->pool);
    }
    /* RV1106 必须等 MB 全部释放再退出 */
    fprintf(stderr, "[venc] waiting for media buffers\n");
    RK_MPI_SYS_WaitFreeMB();
    RK_MPI_SYS_Exit();
    fprintf(stderr, "[venc] encoder stopped\n");
    free(e);
}

void mpp_encoder_set_nalu_callback(mpp_encoder_t *e, nalu_callback_t cb, void *userdata){
    if(!e) return;
    e->cb = cb;
    e->userdata = userdata;
}

int mpp_encoder_encode(mpp_encoder_t *e, const uint8_t *frame, size_t size, uint64_t pts_us){
    if(!e || !frame) return -1;
    if(size < e->frame_size){
        fprintf(stderr, "[venc] frame too small %zu < %zu\n", size, e->frame_size);
        return -1;
    }

    /* 取 MB -> 拷贝 -> flush cache -> 送编码 (阻塞) -> 归还 */
    MB_BLK blk = RK_MPI_MB_GetMB(e->pool, e->frame_size, RK_TRUE);
    if(!blk){
        fprintf(stderr, "[venc] GetMB fail\n");
        return -1;
    }
    void *dst = RK_MPI_MB_Handle2VirAddr(blk);
    memcpy(dst, frame, e->frame_size);
    RK_MPI_SYS_MmzFlushCache(blk, RK_FALSE);

    VIDEO_FRAME_INFO_S info;
    memset(&info, 0, sizeof(info));
    info.stVFrame.pMbBlk        = blk;
    info.stVFrame.u32Width      = e->width;
    info.stVFrame.u32Height     = e->height;
    info.stVFrame.u32VirWidth   = e->hor_stride;
    info.stVFrame.u32VirHeight  = e->ver_stride;
    info.stVFrame.enPixelFormat = RK_FMT_YUV420SP;
    info.stVFrame.u64PTS        = pts_us;

    RK_S32 ret = RK_MPI_VENC_SendFrame(e->chn, &info, -1);  /* 阻塞送帧 */
    RK_MPI_MB_ReleaseMB(blk);
    if(ret != 0){
        fprintf(stderr, "[venc] SendFrame fail %d\n", ret);
        return -1;
    }
    return 0;
}

int mpp_encoder_request_idr(mpp_encoder_t *e){
    if(!e) return -1;
    /* RK_FALSE = 下一帧出 IDR (官方测试程序用法) */
    RK_S32 ret = RK_MPI_VENC_RequestIDR(e->chn, RK_FALSE);
    return ret == 0 ? 0 : -1;
}

int mpp_encoder_set_bitrate(mpp_encoder_t *e, int kbps){
    if(!e || kbps <= 0) return -1;
    e->bitrate_kbps = kbps;
    e->attr.stRcAttr.stH264Cbr.u32BitRate = (RK_U32)kbps;
    RK_S32 ret = RK_MPI_VENC_SetChnAttr(e->chn, &e->attr);
    if(ret != 0){
        fprintf(stderr, "[venc] SetChnAttr(bitrate=%d) fail %d\n", kbps, ret);
        return -1;
    }
    return 0;
}

void mpp_encoder_get_stride(mpp_encoder_t *e, int *hor_stride, int *ver_stride){
    if(!e) return;
    if(hor_stride) *hor_stride = e->hor_stride;
    if(ver_stride) *ver_stride = e->ver_stride;
}
