/* GStreamer
 * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: LGPL-2.0-only
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library 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
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA 02110-1301, USA.
 */

#include <stdio.h>
#include <string.h>
#include <gst/video/video.h>
#include "gstnvjpegbufferpool.h"

GST_DEBUG_CATEGORY_STATIC (gst_nvjpeg_pool_debug);
#define GST_CAT_DEFAULT gst_nvjpeg_pool_debug

typedef struct _GstNvJpegMemoryAllocator GstNvJpegMemoryAllocator;
typedef struct _GstNvJpegMemoryAllocatorClass GstNvJpegMemoryAllocatorClass;

typedef struct GstNvJpegMemory {
  GstMemory mem;

  NvBufSurface *data;
} GstNvJpegMemory;

struct _GstNvJpegMemoryAllocator
{
  GstAllocator parent;
};

struct _GstNvJpegMemoryAllocatorClass
{
  GstAllocatorClass parent_class;
};

struct _GstNvJpegBufferPoolPrivate
{
  GstCaps *caps;
  GstVideoInfo info;
  GstAllocator *allocator;
  GstAllocationParams params;
  guint gpuId;
  guint batchSize;
  guint colorSpace;
  guint layout;
  guint size;
  NvBufSurfaceMemType memType;
  gboolean clearChroma;
  guint planeOrder;
  guint precision;
  gboolean isContiguous;
  NvBufSurfaceTag memtag;
  gboolean disablePitchPadding;
};

GType gst_nvjpeg_memory_allocator_get_type (void);

#define gst_nvjpeg_buffer_pool_parent_class parent_class
G_DEFINE_TYPE_WITH_CODE (GstNvJpegBufferPool, gst_nvjpeg_buffer_pool, GST_TYPE_BUFFER_POOL, G_ADD_PRIVATE (GstNvJpegBufferPool));

G_DEFINE_TYPE (GstNvJpegMemoryAllocator, gst_nvjpeg_memory_allocator,
    GST_TYPE_ALLOCATOR);


static NvBufSurfaceColorFormat
get_nvbuf_surface_format (GstVideoFormat format, GstVideoColorimetry *cinfo,
        gboolean planeOrder, guint precision)
{
    NvBufSurfaceColorFormat dstFormat = NVBUF_COLOR_FORMAT_INVALID;
    switch (format) {
    case GST_VIDEO_FORMAT_I420:
      dstFormat = NVBUF_COLOR_FORMAT_YUV420;
      break;
    case  GST_VIDEO_FORMAT_NV12:
      dstFormat = NVBUF_COLOR_FORMAT_NV12;
      break;
    case GST_VIDEO_FORMAT_Y444:
      dstFormat = NVBUF_COLOR_FORMAT_YUV444;
      break;
    case  GST_VIDEO_FORMAT_RGBA:
      dstFormat = NVBUF_COLOR_FORMAT_RGBA;
      break;
    case  GST_VIDEO_FORMAT_GRAY8:
      dstFormat = NVBUF_COLOR_FORMAT_GRAY8;
      break;
    case GST_VIDEO_FORMAT_Y42B:
      dstFormat =  NVBUF_COLOR_FORMAT_YUV422;
      break;
    case GST_VIDEO_FORMAT_YUY2:
      dstFormat = NVBUF_COLOR_FORMAT_YUYV;
      break;
    default:
      dstFormat = NVBUF_COLOR_FORMAT_INVALID;
      break;
  }
  return dstFormat;
}

static gpointer
gst_nvjpeg_memory_map (GstMemory *mem, gsize maxsize, GstMapFlags flags)
{
  GstNvJpegMemory *dsmem = (GstNvJpegMemory *) mem;
  return dsmem->data;
}

static void
gst_nvjpeg_memory_unmap (GstMemory *mem)
{
  // Nothing to do here.
}

static GstMemory *
gst_nvjpeg_memory_share (GstMemory *mem, gssize offset, gssize size)
{
  /*
     Currently it won't be used because memory is non-shared.
   */
  g_assert_not_reached ();
  return NULL;
}

static void
gst_nvjpeg_memory_allocator_init (GstNvJpegMemoryAllocator *allocator)
{
  GstAllocator *parent = GST_ALLOCATOR_CAST (allocator);

  parent->mem_type = GST_NVJPEG_MEMORY_TYPE;
  parent->mem_map = gst_nvjpeg_memory_map;
  parent->mem_unmap = gst_nvjpeg_memory_unmap;
  parent->mem_share = gst_nvjpeg_memory_share;

  /* We want to use default implementation of ->mem_copy which uses
     default allocator for memory allocation and then do memcpy().
     We are using this approach because we haven't implemented
     ->alloc() funciton.
     Secondly, we want to copy NvMMBuffer structure only not the deep copy.
   */

  GST_OBJECT_FLAG_SET (allocator, GST_ALLOCATOR_FLAG_CUSTOM_ALLOC);
}

static GstMemory *
gst_nvjpeg_memory_allocator_alloc (GstMemoryFlags flags,
                                 GstNvJpegBufferPoolPrivate *priv)
{
  GstNvJpegMemory *mem = NULL;
  NvBufSurface *surf = NULL;
  NvBufSurfaceAllocateParams param = {0};
  int status = -1;

  mem = g_slice_new0 (GstNvJpegMemory);

  param.params.width = priv->info.width & 1 ? priv->info.width + 1 : priv->info.width;
  param.params.height = priv->info.height & 1 ? priv->info.height + 1 : priv->info.height;
  param.params.memType = priv->memType;
  param.params.gpuId = priv->gpuId;
  param.params.isContiguous = priv->isContiguous;
  param.memtag = priv->memtag;
  param.disablePitchPadding = priv->disablePitchPadding;

  param.params.colorFormat = get_nvbuf_surface_format (GST_VIDEO_INFO_FORMAT (&priv->info),
      &(priv->info.colorimetry), priv->planeOrder, priv->precision);

  if (priv->layout)
    param.params.layout =  NVBUF_LAYOUT_BLOCK_LINEAR;
  else
    param.params.layout =   NVBUF_LAYOUT_PITCH;

  status = NvBufSurfaceAllocate(&surf, priv->batchSize, &param);
  if (status < 0) {
    printf ("Error(%d) in buffer allocation\n", status);
    goto error;
  }
  surf->numFilled = 1;

  if (priv->clearChroma) {
    guint i, j;
    for (i = 0; i < surf->batchSize; i++) {
      for (j = 0; j < surf->surfaceList[i].planeParams.num_planes; j++) {
        /* Memset 0th surface to 0 and other surfaces to 128 for black color */
        status = NvBufSurfaceMemSet (surf, i, j, ((j == 0) ? 0 : 128));
        if (status < 0) {
          printf ("Error(%d) in buffer memset\n", status);
          goto error;
        }
      }
    }
  }


  flags = (GstMemoryFlags) (flags | GST_MEMORY_FLAG_NO_SHARE);

  gst_memory_init (GST_MEMORY_CAST (mem), flags, priv->allocator, NULL,
                   sizeof (NvBufSurface), 0, 0, sizeof (NvBufSurface));

  mem->data = surf;

  return GST_MEMORY_CAST (mem);

error:
  if (surf && surf->surfaceList) {
    NvBufSurfaceDestroy (surf);
  }
  g_slice_free (GstNvJpegMemory, mem);

  return NULL;
}

static GstMemory *
gst_nvjpeg_memory_allocator_alloc_dummy (GstAllocator *allocator,
        gsize size, GstAllocationParams *params)
{
  /*
     We are not using this function for memory allocation because
     to allocate NvBufSurface we need additional information which
     can't be passed here. We have defined custom function
     gst_nvjpeg_memory_allocator_alloc for this.
   */
  g_assert_not_reached ();
  return NULL;
}

static void
gst_nvjpeg_memory_allocator_free (GstAllocator *allocator, GstMemory *memory)
{
  GstNvJpegMemory *dsmem = (GstNvJpegMemory *) memory;
  NvBufSurface *surf = dsmem->data;

  NvBufSurfaceDestroy (surf);
  g_slice_free (GstNvJpegMemory, dsmem);
}

static void
gst_nvjpeg_memory_allocator_class_init (GstNvJpegMemoryAllocatorClass *klass)
{
  GstAllocatorClass *allocator_class;

  allocator_class = GST_ALLOCATOR_CLASS (klass);

  allocator_class->alloc = gst_nvjpeg_memory_allocator_alloc_dummy;
  allocator_class->free = gst_nvjpeg_memory_allocator_free;
}

static void
gst_nvjpeg_buffer_pool_finalize (GObject *object)
{
  GstNvJpegBufferPool *pool = GST_NVJPEG_BUFFER_POOL (object);
  GstNvJpegBufferPoolPrivate *priv = pool->priv;

  if (priv->caps)
    gst_caps_unref (priv->caps);
  priv->caps = NULL;

  if (priv->allocator)
    gst_object_unref (priv->allocator);
  priv->allocator = NULL;

  G_OBJECT_CLASS (parent_class)->finalize (object);
}


static const gchar ** gst_nvjpeg_buffer_pool_get_options (GstBufferPool *bpool)
{
  static const gchar *pool_options[] = {
      NULL };

  return pool_options;
}

static gboolean
gst_nvjpeg_buffer_pool_set_config (GstBufferPool *bpool, GstStructure *config)
{
  GstNvJpegBufferPool *pool = GST_NVJPEG_BUFFER_POOL (bpool);
  GstNvJpegBufferPoolPrivate *priv = pool->priv;
  GstCaps *caps = NULL;
  GstAllocator *allocator = NULL;
  GstAllocationParams params;
  GstVideoInfo info;
  guint size;
  gchar *capsString = NULL;

  GST_OBJECT_LOCK (pool);

  if (!gst_buffer_pool_config_get_params (config, &caps, &size, NULL, NULL))
    goto wrong_config;

  if (caps == NULL)
    goto no_caps;

  capsString = gst_caps_to_string (caps);
  GST_DEBUG_OBJECT (pool, "caps : %s\n", capsString);
  g_free(capsString);

  if (!gst_video_info_from_caps (&info, caps))
    goto wrong_video_caps;

  priv->info = info;
  priv->size = size;

  if (!priv->size)
    priv->size = info.size;

  if (priv->caps)
    gst_caps_unref (priv->caps);
  priv->caps = gst_caps_ref (caps);

  if (!gst_buffer_pool_config_get_allocator (config, &allocator, &params))
    goto wrong_config;

  if (allocator) {
    // Update user provided allocator
    if (priv->allocator)
      gst_object_unref (priv->allocator);
    if ((priv->allocator = allocator))
      gst_object_ref (allocator);

    priv->params = params;
  }

  if (!gst_structure_get_uint (config, "memtype", (guint*)(&priv->memType)))
    priv->memType = NVBUF_MEM_DEFAULT;

  if (!gst_structure_get_uint (config, "gpu-id", &priv->gpuId))
    priv->gpuId = 0;

  if (!gst_structure_get_uint (config, "memtag", (guint*)(&priv->memtag)))
    priv->memtag = NvBufSurfaceTag_NONE;

  if (!gst_structure_get_uint (config, "batch-size", &priv->batchSize) ||
      !priv->batchSize) {
    GST_DEBUG_OBJECT (pool, "batch-size not set, using 1 as default value.");
    priv->batchSize = 1;
  }

  if (!gst_structure_get_boolean (config, "clear-chroma", &priv->clearChroma))
    priv->clearChroma = FALSE;

  //Check RGB plane ordering for planar RGB format by default, RGB else BGR
  if (!gst_structure_get_uint (config,"plane-order", &priv->planeOrder))
    priv->planeOrder = 0;

  // for RGB planar format, precision=0 u8, precision=1 fp32
  if (!gst_structure_get_uint (config,"precision", &priv->precision))
    priv->precision = 0;

  if (!gst_structure_get_boolean (config,"disable-pitch-padding", &priv->disablePitchPadding))
    priv->disablePitchPadding = 0;

  if (!allocator) {
    // Allocator is not set, use private allocator.
    gst_buffer_pool_config_set_allocator (config, priv->allocator,
                                          &priv->params);
  }

  // Check block linear layout flag
  if (!gst_structure_get_uint (config,"bl-output", &priv->layout))
    priv->layout = 0;

  if (!gst_structure_get_boolean (config, "contiguous-alloc", &priv->isContiguous))
    priv->isContiguous = FALSE;

  GST_OBJECT_UNLOCK (pool);

  return GST_BUFFER_POOL_CLASS (parent_class)->set_config (bpool, config);

  /* ERRORS */
wrong_config:
  {
    GST_OBJECT_UNLOCK (pool);
    GST_WARNING_OBJECT (pool, "invalid config");
    return FALSE;
  }
no_caps:
  {
    GST_OBJECT_UNLOCK (pool);
    GST_WARNING_OBJECT (pool, "no caps in config");
    return FALSE;
  }
wrong_video_caps:
  {
    GST_OBJECT_UNLOCK (pool);
    GST_WARNING_OBJECT (pool,
        "failed getting video info from caps %" GST_PTR_FORMAT, caps);
    return FALSE;
  }
}

static gboolean
gst_nvjpeg_buffer_pool_start (GstBufferPool *bpool)
{
  GstNvJpegBufferPool *pool = GST_NVJPEG_BUFFER_POOL (bpool);

  GST_DEBUG_OBJECT (pool, "start");


  return GST_BUFFER_POOL_CLASS (parent_class)->start (bpool);
}

static gboolean gst_nvjpeg_buffer_pool_stop (GstBufferPool *bpool)
{
  GstNvJpegBufferPool *pool = GST_NVJPEG_BUFFER_POOL (bpool);

  GST_DEBUG_OBJECT (pool, "stop");

  return GST_BUFFER_POOL_CLASS (parent_class)->stop (bpool);
}

static GstFlowReturn gst_nvjpeg_buffer_pool_acquire_buffer (
    GstBufferPool *bpool, GstBuffer **buffer,
    GstBufferPoolAcquireParams *params)
{
  GstFlowReturn ret;
  GstNvJpegBufferPool *pool = GST_NVJPEG_BUFFER_POOL (bpool);

  GST_DEBUG_OBJECT (pool, "acquire_buffer");
  g_atomic_int_inc (&pool->num_queued);

  /*
   * Currently, we are using base class version because we are not attaching
   * anything extra to buffers. If we need to do so We will modify this function
   * accordingly.
   */
  ret = GST_BUFFER_POOL_CLASS (parent_class)->acquire_buffer (bpool, buffer, params);

  return ret;
}

static GstFlowReturn
gst_nvjpeg_buffer_pool_alloc_buffer (GstBufferPool *bpool,
    GstBuffer **buffer, GstBufferPoolAcquireParams *params)
{
  GstNvJpegBufferPool *pool = GST_NVJPEG_BUFFER_POOL (bpool);
  GstNvJpegBufferPoolPrivate *priv = pool->priv;
  GstBuffer *buf = NULL;
  GstMemory *mem = NULL;

  GST_DEBUG_OBJECT (pool, "alloc_buffer");

  if (!g_strcmp0(priv->allocator->mem_type, GST_NVJPEG_MEMORY_TYPE)) {
    mem = gst_nvjpeg_memory_allocator_alloc ((GstMemoryFlags) 0, priv);
    g_return_val_if_fail (mem, GST_FLOW_ERROR);

    buf = gst_buffer_new ();
    gst_buffer_append_memory (buf, mem);
  } else {
    buf = gst_buffer_new_allocate (priv->allocator, priv->size,
                                   &priv->params);

    g_return_val_if_fail (buf, GST_FLOW_ERROR);
  }

  *buffer = buf;

  return GST_FLOW_OK;
}

static void gst_nvjpeg_buffer_pool_release_buffer (GstBufferPool *bpool,
        GstBuffer *buffer)
{
  GstNvJpegBufferPool *pool = GST_NVJPEG_BUFFER_POOL (bpool);
  guint num_queued;

  GST_DEBUG_OBJECT (pool, "release_buffer");

  num_queued = g_atomic_int_get (&pool->num_queued);

  if (num_queued > 0) {
    g_atomic_int_add (&pool->num_queued, -1);
  }

  GST_BUFFER_POOL_CLASS (parent_class)->release_buffer (bpool, buffer);
}

static void gst_nvjpeg_buffer_pool_free_buffer (GstBufferPool *bpool,
        GstBuffer *buffer)
{
  GstNvJpegBufferPool *pool = GST_NVJPEG_BUFFER_POOL (bpool);

  GST_DEBUG_OBJECT (pool, "free_buffer");

  GST_BUFFER_POOL_CLASS (parent_class)->free_buffer (bpool, buffer);
}

static void gst_nvjpeg_buffer_pool_init (GstNvJpegBufferPool *pool)
{
  GstNvJpegBufferPool *self = GST_NVJPEG_BUFFER_POOL (pool);
  pool->priv = (GstNvJpegBufferPoolPrivate *) gst_nvjpeg_buffer_pool_get_instance_private (self);
  memset (pool->priv, 0, sizeof(GstNvJpegBufferPoolPrivate));

  pool->priv->allocator = (GstAllocator *) g_object_new (
            gst_nvjpeg_memory_allocator_get_type (), NULL);

  pool->priv->memType = NVBUF_MEM_DEFAULT;
  pool->priv->gpuId = 0;
  pool->priv->batchSize = 1;
  pool->priv->memtag = NvBufSurfaceTag_NONE;

  gst_allocation_params_init (&pool->priv->params);
}

static void
gst_nvjpeg_buffer_pool_class_init (GstNvJpegBufferPoolClass *klass)
{
  GObjectClass *gobject_class = (GObjectClass *) klass;
  GstBufferPoolClass *gstbufferpool_class = (GstBufferPoolClass *) klass;

  gobject_class->finalize = gst_nvjpeg_buffer_pool_finalize;
  gstbufferpool_class->start = gst_nvjpeg_buffer_pool_start;
  gstbufferpool_class->stop = gst_nvjpeg_buffer_pool_stop;
  gstbufferpool_class->get_options = gst_nvjpeg_buffer_pool_get_options;
  gstbufferpool_class->set_config = gst_nvjpeg_buffer_pool_set_config;
  gstbufferpool_class->alloc_buffer = gst_nvjpeg_buffer_pool_alloc_buffer;
  gstbufferpool_class->free_buffer = gst_nvjpeg_buffer_pool_free_buffer;
  gstbufferpool_class->acquire_buffer = gst_nvjpeg_buffer_pool_acquire_buffer;
  gstbufferpool_class->release_buffer = gst_nvjpeg_buffer_pool_release_buffer;

  GST_DEBUG_CATEGORY_INIT (gst_nvjpeg_pool_debug, "nvjpegpool", 0,
        "nvjpeg buffer pool object");
}

GstBufferPool * gst_nvjpeg_buffer_pool_new (void)
{
  GstNvJpegBufferPool *pool;
  pool = (GstNvJpegBufferPool *) g_object_new (GST_TYPE_NVJPEG_BUFFER_POOL, NULL);

  return GST_BUFFER_POOL (pool);
}
