/*
 * This file is part of FFmpeg.
 *
 * FFmpeg is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * FFmpeg is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with FFmpeg; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 */


#include "config.h"
#include "decode.h"
#include "hevc/hevcdec.h"
#include "hwaccel_internal.h"
#include "hwconfig.h"
#include "internal.h"

#include "v4l2_fmt.h"
#include "v4l2_request_hevc.h"

#include "libavutil/hwcontext_drm.h"
#include "libavutil/mem.h"
#include "libavutil/pixdesc.h"

#include "v4l2_req_devscan.h"
#include "v4l2_req_dmabufs.h"
#include "v4l2_req_pollqueue.h"
#include "v4l2_req_media.h"
#include "v4l2_req_utils.h"

static size_t bit_buf_size(unsigned int w, unsigned int h, unsigned int bits_minus8)
{
    const size_t wxh = w * h;
    size_t bits_alloc;

    /* Annex A gives a min compression of 2 @ lvl 3.1
     * (wxh <= 983040) and min 4 thereafter but avoid
     * the odity of 983041 having a lower limit than
     * 983040.
     * Multiply by 3/2 for 4:2:0
     */
    bits_alloc = wxh < 983040 ? wxh * 3 / 4 :
        wxh < 983040 * 2 ? 983040 * 3 / 4 :
        wxh * 3 / 8;
    /* Allow for bit depth */
    bits_alloc += (bits_alloc * bits_minus8) / 8;
    /* Add a few bytes (16k) for overhead */
    bits_alloc += 0x4000;
    return bits_alloc;
}

int ff_v4l2_request_start_frame(AVCodecContext *avctx,
                                     av_unused const uint8_t *buffer,
                                     av_unused uint32_t size)
{
    V4L2RequestPrivHEVC * const priv = avctx->internal->hwaccel_priv_data;
    V4L2RequestContextHEVC *const ctx = priv->cctx;
    return ctx->fns->start_frame(avctx, ctx, buffer, size);
}

int ff_v4l2_request_decode_slice(AVCodecContext *avctx, const uint8_t *buffer, uint32_t size)
{
    V4L2RequestPrivHEVC * const priv = avctx->internal->hwaccel_priv_data;
    V4L2RequestContextHEVC *const ctx = priv->cctx;
    return ctx->fns->decode_slice(avctx, ctx, buffer, size);
}

int ff_v4l2_request_end_frame(AVCodecContext *avctx)
{
    V4L2RequestPrivHEVC * const priv = avctx->internal->hwaccel_priv_data;
    V4L2RequestContextHEVC *const ctx = priv->cctx;
    return ctx->fns->end_frame(avctx, ctx);
}

void ff_v4l2_request_abort_frame(AVCodecContext * const avctx)
{
    V4L2RequestPrivHEVC * const priv = avctx->internal->hwaccel_priv_data;
    V4L2RequestContextHEVC *const ctx = priv->cctx;
    ctx->fns->abort_frame(avctx, ctx);
}

int ff_v4l2_request_frame_params(AVCodecContext *avctx, AVBufferRef *hw_frames_ctx)
{
    V4L2RequestPrivHEVC * const priv = avctx->internal->hwaccel_priv_data;
    V4L2RequestContextHEVC *const ctx = priv->cctx;
    return ctx->fns->frame_params(avctx, ctx, hw_frames_ctx);
}

int ff_v4l2_request_alloc_frame(AVCodecContext * avctx, AVFrame *frame)
{
    V4L2RequestPrivHEVC * const priv = avctx->internal->hwaccel_priv_data;
    V4L2RequestContextHEVC *const ctx = priv->cctx;
    return ctx->fns->alloc_frame(avctx, ctx, frame);
}


static void
cctx_free(void * v, uint8_t * data)
{
    V4L2RequestContextHEVC *const ctx = (V4L2RequestContextHEVC *)data;

    if (ctx->fns && ctx->fns->uninit)
        ctx->fns->uninit(ctx);

    mediabufs_ctl_unref(&ctx->mbufs);
    media_pool_delete(&ctx->mpool);
    pollqueue_unref(&ctx->pq);
    dmabufs_ctl_unref(&ctx->dbufs);
    devscan_delete(&ctx->devscan);

    decode_q_uninit(&ctx->decode_q);

    av_free(ctx);
}

int ff_v4l2_request_uninit(AVCodecContext *avctx)
{
    V4L2RequestPrivHEVC * const priv = avctx->internal->hwaccel_priv_data;

    av_log(avctx, AV_LOG_DEBUG, "<<< %s\n", __func__);

    if (priv->cctx != NULL) {
        V4L2RequestContextHEVC *const ctx = priv->cctx;

        decode_q_wait(&ctx->decode_q, NULL);  // Wait for all other threads to be out of decode
        mediabufs_stream_wait_dst_done(ctx->mbufs);  // Now wait for processing to finish

        priv->cctx = NULL;
        av_buffer_unref(&priv->cctx_buf);
    }

//    if (avctx->hw_frames_ctx) {
//        AVHWFramesContext *hwfc = (AVHWFramesContext*)avctx->hw_frames_ctx->data;
//        av_buffer_pool_flush(hwfc->pool);
//    }
    return 0;
}

static int dst_fmt_accept_cb(void * v, const struct v4l2_fmtdesc *fmtdesc)
{
    const int bit_depth = *(int *)v;

    // SAND is currently confised as to whether it is s/w or hardware
    // *** Current usage is probably wrong - it shouldn't be a h/w fmt
    if (fmtdesc->pixelformat == V4L2_PIX_FMT_NV12_COL128 ||
        fmtdesc->pixelformat == V4L2_PIX_FMT_NV12_COL128M) {
        return bit_depth == 8;
    }
    if (fmtdesc->pixelformat == V4L2_PIX_FMT_NV12_10_COL128 ||
        fmtdesc->pixelformat == V4L2_PIX_FMT_NV12_10_COL128M) {
        return bit_depth == 10;
    }
    else {
        const enum AVPixelFormat fmt = ff_v4l2_format_v4l2_to_avfmt(fmtdesc->pixelformat, AV_CODEC_ID_RAWVIDEO);
        const AVPixFmtDescriptor * const desc = av_pix_fmt_desc_get(fmt);

        if (fmt == AV_PIX_FMT_NONE || desc == NULL)
            return 0;

        return bit_depth == desc->comp[0].depth;
    }
}

int ff_v4l2_request_init(AVCodecContext *avctx,
                         const struct v4l2_req_decode_fns * const * const try_fns,
                         const int width, const int height, const int bit_depth,
                         const int dst_buffers)
{
    V4L2RequestPrivHEVC * const priv = avctx->internal->hwaccel_priv_data;
    V4L2RequestContextHEVC * ctx;
    int ret;
    const struct decdev * decdev;
    const uint32_t src_pix_fmt = try_fns[0]->src_pix_fmt_v4l2;  // Assuming constant for all APIs but avoiding V4L2 includes
    size_t src_size;
    enum mediabufs_memory src_memtype;
    enum mediabufs_memory dst_memtype;

    av_log(avctx, AV_LOG_DEBUG, "<<< %s\n", __func__);

    if ((ctx = av_mallocz(sizeof(*ctx))) == NULL) {
        av_log(avctx, AV_LOG_ERROR, "Unable to allocate context");
        return AVERROR(ENOMEM);
    }
    if ((priv->cctx_buf = av_buffer_create((uint8_t*)ctx, sizeof(*ctx), cctx_free, NULL, 0)) == NULL) {
        av_log(avctx, AV_LOG_ERROR, "Unable to allocate context buffer");
        av_free(ctx);
        return AVERROR(ENOMEM);
    }
    priv->cctx = ctx;

    decode_q_init(&ctx->decode_q);

    if ((ret = devscan_build(avctx, &ctx->devscan)) != 0) {
        av_log(avctx, AV_LOG_WARNING, "Failed to find any V4L2 devices\n");
        ret = AVERROR(-ret);
        goto fail0;
    }
    ret = AVERROR(ENOMEM);  // Assume mem fail by default for these

    if ((decdev = devscan_find(ctx->devscan, src_pix_fmt)) == NULL)
    {
        av_log(avctx, AV_LOG_WARNING, "Failed to find a device for %s\n", try_fns[0]->name);
        ret = AVERROR(ENODEV);
        goto fail0;
    }
    av_log(avctx, AV_LOG_DEBUG, "Trying V4L2 devices: %s,%s\n",
           decdev_media_path(decdev), decdev_video_path(decdev));

    if ((ctx->pq = pollqueue_new()) == NULL) {
        av_log(avctx, AV_LOG_ERROR, "Unable to create pollqueue\n");
        goto fail1;
    }

    if ((ctx->mpool = media_pool_new(decdev_media_path(decdev), ctx->pq, 4)) == NULL) {
        av_log(avctx, AV_LOG_ERROR, "Unable to create media pool\n");
        goto fail2;
    }

    if ((ctx->mbufs = mediabufs_ctl_new(avctx, decdev_video_path(decdev), ctx->pq)) == NULL) {
        av_log(avctx, AV_LOG_ERROR, "Unable to create media controls\n");
        goto fail3;
    }

    // Version test for functional Pi5 HEVC iommu.
    // rpivid kernel patch was merged in 6.1.57
    // *** Remove when it is unlikely that there are any broken kernels left
    if (mediabufs_ctl_driver_version(ctx->mbufs) >= MEDIABUFS_DRIVER_VERSION(6,1,57))
        ctx->dbufs = dmabufs_ctl_new_vidbuf_cached();
    else
        ctx->dbufs = dmabufs_ctl_new();

    if (ctx->dbufs == NULL) {
        av_log(avctx, AV_LOG_DEBUG, "Unable to open dmabufs - try mmap buffers\n");
        src_memtype = MEDIABUFS_MEMORY_MMAP;
        dst_memtype = MEDIABUFS_MEMORY_MMAP;
    }
    else {
        av_log(avctx, AV_LOG_DEBUG, "Dmabufs opened - try dmabuf buffers\n");
        src_memtype = MEDIABUFS_MEMORY_DMABUF;
        dst_memtype = MEDIABUFS_MEMORY_DMABUF;
    }

    // Ask for an initial bitbuf size of max size / 4
    // We will realloc if we need more
    // Must use sps->h/w as avctx contains cropped size
retry_src_memtype:
    src_size = bit_buf_size(width, height, bit_depth - 8);
    if (src_memtype == MEDIABUFS_MEMORY_DMABUF && mediabufs_src_resizable(ctx->mbufs))
        src_size /= 4;
    // Kludge for conformance tests which break Annex A limits
    else if (src_size < 0x40000)
        src_size = 0x40000;

    if (mediabufs_src_fmt_set(ctx->mbufs, decdev_src_type(decdev), src_pix_fmt,
                              width, height, src_size)) {
        char tbuf1[5];
        av_log(avctx, AV_LOG_ERROR, "Failed to set source format: %s %dx%d\n", strfourcc(tbuf1, src_pix_fmt), width, height);
        goto fail4;
    }

    if (mediabufs_src_chk_memtype(ctx->mbufs, src_memtype)) {
        if (src_memtype == MEDIABUFS_MEMORY_DMABUF) {
            src_memtype = MEDIABUFS_MEMORY_MMAP;
            goto retry_src_memtype;
        }
        av_log(avctx, AV_LOG_ERROR, "Failed to get src memory type\n");
        goto fail4;
    }

    for (const struct v4l2_req_decode_fns * const * tf = try_fns; (ctx->fns = *tf) != NULL; ++tf)
        if (ctx->fns->probe(avctx, ctx) == 0)
            break;
    if (ctx->fns == NULL) {
        av_log(avctx, AV_LOG_ERROR, "No %s version probed successfully\n", try_fns[0]->name);
        ret = AVERROR(EINVAL);
        goto fail4;
    }
    av_log(avctx, AV_LOG_DEBUG, "%s probed successfully: driver v %#x\n",
           ctx->fns->name, mediabufs_ctl_driver_version(ctx->mbufs));

    if (mediabufs_dst_fmt_set(ctx->mbufs, width, height, dst_fmt_accept_cb, (void*)&bit_depth)) {
        av_log(avctx, AV_LOG_ERROR, "Failed to set destination format: %dx%d %dbit\n", width, height, bit_depth);
        goto fail4;
    }

    if (mediabufs_src_pool_create(ctx->mbufs, ctx->dbufs, 6, src_memtype)) {
        av_log(avctx, AV_LOG_ERROR, "Failed to create source pool\n");
        goto fail4;
    }

    {
        unsigned int dst_slots = dst_buffers +
            avctx->thread_count + (avctx->extra_hw_frames > 0 ? avctx->extra_hw_frames : 6);
        av_log(avctx, AV_LOG_DEBUG, "Slots=%d: Reordering=%d, threads=%d, hw+=%d\n", dst_slots,
               dst_buffers,
               avctx->thread_count, avctx->extra_hw_frames);

        if (mediabufs_dst_chk_memtype(ctx->mbufs, dst_memtype)) {
            if (dst_memtype != MEDIABUFS_MEMORY_DMABUF) {
                av_log(avctx, AV_LOG_ERROR, "Failed to get dst memory type\n");
                goto fail4;
            }
            av_log(avctx, AV_LOG_DEBUG, "Dst DMABUF not supported - trying mmap\n");
            dst_memtype = MEDIABUFS_MEMORY_MMAP;
        }

        // extra_hw_frames is -1 if unset
        if (mediabufs_dst_slots_create(ctx->mbufs, dst_slots, (avctx->extra_hw_frames > 0), dst_memtype)) {
            av_log(avctx, AV_LOG_ERROR, "Failed to create destination slots\n");
            goto fail4;
        }
    }

    if (mediabufs_stream_on(ctx->mbufs)) {
        av_log(avctx, AV_LOG_ERROR, "Failed stream on\n");
        goto fail4;
    }

    if ((ret = ff_decode_get_hw_frames_ctx(avctx, AV_HWDEVICE_TYPE_DRM)) != 0) {
        av_log(avctx, AV_LOG_ERROR, "Failed to create frame ctx\n");
        goto fail4;
    }

    if ((ret = ctx->fns->set_controls(avctx, ctx)) != 0) {
        av_log(avctx, AV_LOG_ERROR, "Failed set controls\n");
        goto fail5;
    }

    // Set our s/w format
    avctx->sw_pix_fmt = ((AVHWFramesContext *)avctx->hw_frames_ctx->data)->sw_format;

    av_log(avctx, AV_LOG_INFO, "Hwaccel %s; devices: %s,%s; buffers: src %s, dst %s; swfmt=%s\n",
           ctx->fns->name,
           decdev_media_path(decdev), decdev_video_path(decdev),
           mediabufs_memory_name(src_memtype), mediabufs_memory_name(dst_memtype),
           av_get_pix_fmt_name(avctx->sw_pix_fmt));

    return 0;

fail5:
    av_buffer_unref(&avctx->hw_frames_ctx);
fail4:
fail3:
fail2:
fail1:
fail0:
    priv->cctx = NULL;
    av_buffer_unref(&priv->cctx_buf);
    return ret;
}

int
ff_v4l2_request_update_thread_context(AVCodecContext *dst, const AVCodecContext *src)
{
    V4L2RequestPrivHEVC * const spriv = src->internal->hwaccel_priv_data;
    V4L2RequestPrivHEVC * const dpriv = dst->internal->hwaccel_priv_data;
    int rv;

    av_log(dst, AV_LOG_DEBUG, "<<< %s (%s)\n", __func__, dpriv->cctx_buf ? "old" : "new");

    if ((rv = av_buffer_replace(&dpriv->cctx_buf, spriv->cctx_buf)) != 0)
        return rv;

    dpriv->cctx = spriv->cctx;
    return 0;
}

void
ff_v4l2_request_free_frame_priv(FFRefStructOpaque hwctx, void *data)
{
}

static int v4l2_request_hevc_init(AVCodecContext *avctx)
{
    const HEVCContext *h = avctx->priv_data;
    const HEVCPPS * const pps = h->pps;
    const HEVCSPS * const sps = pps->sps;

    const struct v4l2_req_decode_fns * const try_fns[] = {
        &V2(ff_v4l2_req_hevc, 4),
#if CONFIG_V4L2_REQ_HEVC_VX
        &V2(ff_v4l2_req_hevc, 3),
        &V2(ff_v4l2_req_hevc, 2),
        &V2(ff_v4l2_req_hevc, 1),
#endif
        NULL
    };

    // Give up immediately if this is something that we have no code to deal with
    if (sps->chroma_format_idc != 1) {
        av_log(avctx, AV_LOG_WARNING, "chroma_format_idc(%d) != 1: Not implemented\n", sps->chroma_format_idc);
        return AVERROR_PATCHWELCOME;
    }
    if (!(sps->bit_depth == 10 || sps->bit_depth == 8) ||
        sps->bit_depth != sps->bit_depth_chroma) {
        av_log(avctx, AV_LOG_WARNING, "Bit depth Y:%d C:%d: Not implemented\n", sps->bit_depth, sps->bit_depth_chroma);
        return AVERROR_PATCHWELCOME;
    }

    return ff_v4l2_request_init(avctx, try_fns, sps->width, sps->height, sps->bit_depth,
                             sps->temporal_layer[sps->max_sub_layers - 1].max_dec_pic_buffering);
}

const FFHWAccel ff_hevc_v4l2request_hwaccel = {
    .p = {
        .name           = "hevc_v4l2request",
        .type           = AVMEDIA_TYPE_VIDEO,
        .id             = AV_CODEC_ID_HEVC,
        .pix_fmt        = AV_PIX_FMT_DRM_PRIME,
    },
    .alloc_frame    = ff_v4l2_request_alloc_frame,
    .start_frame    = ff_v4l2_request_start_frame,
    .decode_slice   = ff_v4l2_request_decode_slice,
    .end_frame      = ff_v4l2_request_end_frame,
    .abort_frame    = ff_v4l2_request_abort_frame,
    .init           = v4l2_request_hevc_init,
    .uninit         = ff_v4l2_request_uninit,
    .free_frame_priv = ff_v4l2_request_free_frame_priv,
    .frame_priv_data_size  = 128,
    .update_thread_context = ff_v4l2_request_update_thread_context,
    .priv_data_size = sizeof(V4L2RequestPrivHEVC),
    .frame_params   = ff_v4l2_request_frame_params,
    .caps_internal  = HWACCEL_CAP_ASYNC_SAFE | HWACCEL_CAP_THREAD_SAFE,
};