---

#include <gst.h>
#include <stdio.h>
#include <unistd.h>
#include <stdint.h>
#include <stdbool.h>

#include <fcntl.h>
#include "common.h"
#include "gstimx_ai_aecnr.h"
#include "gst_buffers_helpers.h"
#include "imx_ai_aecnr_core.h"


G_DEFINE_TYPE (Gstimx_ai_aecnr, gst_imx_ai_aecnr, GST_TYPE_ELEMENT);
GST_ELEMENT_REGISTER_DEFINE (imx_ai_aecnr, "imx_ai_aecnr", GST_RANK_NONE, GST_TYPE_IMX_AI_AECNR);

enum
{
  PROP_0,
  PROP_BYPASS,       /* bypass AI_AECNR for debug purpose. */
  PROP_RX_PADDING,   /* Fill Rx with 0 if no data received.*/
  PROP_MODEL
};


/* Plugin pads definition. */

static GstStaticPadTemplate sink_remote_pad_template = GST_STATIC_PAD_TEMPLATE ("sink_remote",
    GST_PAD_SINK,
    GST_PAD_ALWAYS,
    GST_STATIC_CAPS ("audio/x-raw,channels=1,rate=16000,format=F32LE")
    );
static GstStaticPadTemplate sink_mic_pad_template = GST_STATIC_PAD_TEMPLATE ("sink_mic",
    GST_PAD_SINK,
    GST_PAD_ALWAYS,
    GST_STATIC_CAPS ("audio/x-raw,channels=1,rate=16000,format=F32LE")
    );
static GstStaticPadTemplate src_remote_pad_template = GST_STATIC_PAD_TEMPLATE ("src_remote",
    GST_PAD_SRC,
    GST_PAD_ALWAYS,
    GST_STATIC_CAPS ("audio/x-raw,channels=1,rate=16000,format=F32LE")
    );

static GstElementClass *parent_class;

/* Input buffer from remote
 * - Add it in a queue for processing later with buffer from mic
 */
static GstFlowReturn gst_imx_ai_aecnr_chain_remote (GstPad * pad, GstObject * parent, GstBuffer * buf)
{
  Gstimx_ai_aecnr *ai_aecnr = GST_IMX_AI_AECNR (parent);

  if (ai_aecnr->bypass) {
    /* RxBuf not needed in that case: drop it. */
    gst_buffer_unref(buf);
    return GST_FLOW_OK;
  }

  GST_OBJECT_LOCK(ai_aecnr);
  buffer_list_append(&ai_aecnr->rx_buf_list, buf, AI_AECNR_SAMPLE_RATE);
  GST_OBJECT_UNLOCK(ai_aecnr);

  return GST_FLOW_OK;
}


/* Input buffer from microphone
 * - retrieve remote buffer from queue
 * - generate an empty buffer if nothing available from remote
 * - remove occasional glitches because of UAC stack
 * - insert both remote and microphone data into their respective circular buffer
 * - retrieve data from previous iteration
 * - TfLite inference
 * - store output tensors for next iteration
 * - create output gstreamer buffer from output tensor
 */
static GstFlowReturn gst_imx_ai_aecnr_chain_mic (GstPad * pad, GstObject * parent, GstBuffer * micBuf)
{
  Gstimx_ai_aecnr *ai_aecnr = GST_IMX_AI_AECNR (parent);
  GstBuffer *buf, *out_buf;
  GstMapInfo out_map;
  float mic_in[AI_AECNR_IO_SAMPLES];
  float rx_in[AI_AECNR_IO_SAMPLES];
  float out[AI_AECNR_IO_SAMPLES];

  if (!ai_aecnr->ai_aecnr_core) {
    return GST_FLOW_ERROR;
  }

  /* Bypass AEC-ML if option is set */
  if (ai_aecnr->bypass) {
    /* reset internal buffers. */
    GST_OBJECT_LOCK(ai_aecnr);
    init_buf_list(&ai_aecnr->rx_buf_list);
    init_buf_list(&ai_aecnr->mic_buf_list);
    GST_OBJECT_UNLOCK(ai_aecnr);
    ai_aecnr->stream_pts = 0;

    /* Push directly mic buffer to remote. */
    gst_pad_push(ai_aecnr->src_remote_pad, micBuf);

    return GST_FLOW_OK;
  }

  /* Add incoming buffer to input list.
   */
  GST_OBJECT_LOCK(ai_aecnr);
  buffer_list_append(&ai_aecnr->mic_buf_list, micBuf, AI_AECNR_SAMPLE_RATE);

  /* Initialize output timestamp to the first received buffer. */
  if (ai_aecnr->stream_pts == 0) {
    ai_aecnr->stream_pts = micBuf->pts;
  }

  /* Enough data in queue to trigger processing. */
  while ((ai_aecnr->mic_buf_list.unprocessed_samples >= AI_AECNR_IO_SAMPLES)
         && ((ai_aecnr->rx_buf_list.unprocessed_samples >= AI_AECNR_IO_SAMPLES) || (ai_aecnr->rx_padding)))  {

    /* If not enough data received from Rx then insert zeros. */
    if (ai_aecnr->rx_buf_list.unprocessed_samples < AI_AECNR_IO_SAMPLES) {
      buf = gst_buffer_new_allocate(NULL, AI_AECNR_IO_SAMPLES*sizeof(float), NULL);
      if (buf) {
        gst_buffer_memset(buf, 0, 0, AI_AECNR_IO_SAMPLES*sizeof(float));
        GST_BUFFER_PTS (buf) = 0;
        GST_BUFFER_DURATION (buf) = SAMPLES_TO_NS(AI_AECNR_IO_SAMPLES, AI_AECNR_SAMPLE_RATE);

        buffer_list_append(&ai_aecnr->rx_buf_list, buf, AI_AECNR_SAMPLE_RATE);
      }
    }

    /* Extract data from gst buffers to input buffers */
    copy_float_from_gst_buffers(&ai_aecnr->mic_buf_list, mic_in, AI_AECNR_IO_SAMPLES, 1, AI_AECNR_SAMPLE_RATE);
    copy_float_from_gst_buffers(&ai_aecnr->rx_buf_list, rx_in, AI_AECNR_IO_SAMPLES, 1, AI_AECNR_SAMPLE_RATE);


    GST_OBJECT_UNLOCK(ai_aecnr); /* Release lock now that accesses to buffer lists are done. */

    /* Call processing function */
    if (imx_ai_aecnr_core_process(ai_aecnr->ai_aecnr_core, mic_in, rx_in, out) != 0) {
      GST_ERROR_OBJECT(ai_aecnr, "license expired - disabling imx_ai_aecnr processing\n");
      return GST_FLOW_ERROR;
    }

    /* Create a new output buffer and copy processed data into it. */
    out_buf = gst_buffer_new_allocate(NULL, sizeof(float)*AI_AECNR_IO_SAMPLES, NULL);
    if (out_buf) {
      GST_BUFFER_PTS (out_buf) = ai_aecnr->stream_pts;
      ai_aecnr->stream_pts += SAMPLES_TO_NS(AI_AECNR_IO_SAMPLES, AI_AECNR_SAMPLE_RATE);
      GST_BUFFER_DURATION (out_buf) = SAMPLES_TO_NS(AI_AECNR_IO_SAMPLES, AI_AECNR_SAMPLE_RATE);
      gst_buffer_map (out_buf, &out_map, GST_MAP_WRITE);

      memcpy(out_map.data, out, sizeof(float)*AI_AECNR_IO_SAMPLES);

      gst_buffer_unmap(out_buf, &out_map);
      gst_pad_push(ai_aecnr->src_remote_pad, out_buf);
    }

    GST_OBJECT_LOCK(ai_aecnr); /* Lock for accessing the buffer list in while condition. */
  }
  GST_OBJECT_UNLOCK(ai_aecnr);

  return GST_FLOW_OK;
}


static gboolean gst_imx_ai_aecnr_sink_event (GstPad * pad, GstObject * parent, GstEvent * event)
{
  Gstimx_ai_aecnr *ai_aecnr = GST_IMX_AI_AECNR (parent);
  gboolean res = TRUE;

  switch (GST_EVENT_TYPE (event)) {
    case GST_EVENT_EOS:
      /* Consider EOS only from mic pad since remote input is automatically replaced by 0 when no data available. */
      if (pad == ai_aecnr->sink_mic_pad ) {
        res = gst_pad_push_event (ai_aecnr->src_remote_pad, event);
      } else {
        gst_event_unref (event);
      }
      break;
    default:
      res = gst_pad_push_event (ai_aecnr->src_remote_pad, event);
      break;
  }

  return res;
}


/* Initialization of the plugin */
static void gst_imx_ai_aecnr_init (Gstimx_ai_aecnr * ai_aecnr)
{
  ai_aecnr->sink_remote_pad = gst_pad_new_from_static_template (&sink_remote_pad_template, "sink_remote");
  gst_pad_set_chain_function (ai_aecnr->sink_remote_pad, GST_DEBUG_FUNCPTR (gst_imx_ai_aecnr_chain_remote));
  gst_pad_set_event_function (ai_aecnr->sink_remote_pad, GST_DEBUG_FUNCPTR (gst_imx_ai_aecnr_sink_event));
  gst_pad_use_fixed_caps (ai_aecnr->sink_remote_pad);
  gst_element_add_pad (GST_ELEMENT (ai_aecnr), ai_aecnr->sink_remote_pad);

  ai_aecnr->sink_mic_pad = gst_pad_new_from_static_template (&sink_mic_pad_template, "sink_mic");
  gst_pad_set_chain_function (ai_aecnr->sink_mic_pad, GST_DEBUG_FUNCPTR (gst_imx_ai_aecnr_chain_mic));
  gst_pad_set_event_function (ai_aecnr->sink_mic_pad, GST_DEBUG_FUNCPTR (gst_imx_ai_aecnr_sink_event));
  gst_pad_use_fixed_caps (ai_aecnr->sink_mic_pad);
  gst_element_add_pad (GST_ELEMENT (ai_aecnr), ai_aecnr->sink_mic_pad);

  ai_aecnr->src_remote_pad = gst_pad_new_from_static_template (&src_remote_pad_template, "src_remote");
  gst_pad_use_fixed_caps (ai_aecnr->src_remote_pad);
  gst_element_add_pad (GST_ELEMENT (ai_aecnr), ai_aecnr->src_remote_pad);

  ai_aecnr->bypass = FALSE;
  ai_aecnr->rx_padding = TRUE;
  /* by default using small model. */
  ai_aecnr->use_small_model = 1;
}


/* Plugin state management.
   Used to open and close internal library.
*/
static GstStateChangeReturn gst_imx_ai_aecnr_change_state (GstElement *element, GstStateChange transition)
{
  Gstimx_ai_aecnr *ai_aecnr = GST_IMX_AI_AECNR (element);
  GstStateChangeReturn ret = GST_STATE_CHANGE_SUCCESS;
  imx_ai_aecnr_config core_config;

  /* plugin starting. */
  if (transition == GST_STATE_CHANGE_NULL_TO_READY) {
    core_config.use_small_model = ai_aecnr->use_small_model;
    core_config.version = NULL;

    ai_aecnr->ai_aecnr_core = imx_ai_aecnr_core_open(&core_config);
    if (!ai_aecnr->ai_aecnr_core) {
      GST_ERROR_OBJECT(ai_aecnr, "Failed to initialize");
      return GST_STATE_CHANGE_FAILURE;
    }

    /* Print version information from the core library */
    GST_INFO_OBJECT(ai_aecnr, "plugin version: %s", IMX_VOICE_PLUGINS_VERSION);
    if (core_config.version != NULL) {
      GST_INFO_OBJECT(ai_aecnr, "%s", core_config.version);
    } else {
      GST_WARNING_OBJECT(ai_aecnr, "AI AECNR Core Library version information not available");
    }

    /* Initialize input buffers lists. */
    GST_OBJECT_LOCK(ai_aecnr);
    init_buf_list(&ai_aecnr->rx_buf_list);
    init_buf_list(&ai_aecnr->mic_buf_list);
    GST_OBJECT_UNLOCK(ai_aecnr);

    ai_aecnr->stream_pts = 0;
  }

  /* Parent class state management.*/
  ret = GST_ELEMENT_CLASS (parent_class)->change_state (element, transition);
  if (ret == GST_STATE_CHANGE_FAILURE)
    return ret;

  /* Plugin closing. */
  if (transition == GST_STATE_CHANGE_READY_TO_NULL) {
    imx_ai_aecnr_core_close(ai_aecnr->ai_aecnr_core);
    ret = GST_STATE_CHANGE_SUCCESS;
  }

  return ret;
}

/* Plugin properties management */
static void gst_imx_ai_aecnr_set_property (GObject * object, guint prop_id, const GValue * value, GParamSpec * pspec)
{
  Gstimx_ai_aecnr *ai_aecnr = GST_IMX_AI_AECNR (object);
  gchar *model_name = NULL;

  switch (prop_id) {
    case PROP_BYPASS:
      ai_aecnr->bypass = g_value_get_boolean (value);
      break;
    case PROP_RX_PADDING:
      ai_aecnr->rx_padding = g_value_get_boolean (value);
      break; 
    case PROP_MODEL:
      model_name = g_value_dup_string (value);
      if (g_str_equal(model_name, "small") || (model_name == NULL)) {
        GST_INFO_OBJECT(ai_aecnr, "configuring small model\n");
        ai_aecnr->use_small_model = true;
      } else if (g_str_equal(model_name, "large")) {
        GST_INFO_OBJECT(ai_aecnr, "configuring large model\n");
        ai_aecnr->use_small_model = false;
      } else {
        GST_INFO_OBJECT(ai_aecnr, "unknown model identifier %s : using small model by default\n", model_name);
        ai_aecnr->use_small_model = true;
      }
      g_free(model_name);
      break;
    default:
      GST_ERROR_OBJECT(ai_aecnr, "unsupported property ID: %d", prop_id);
      break;
  }
}

static void gst_imx_ai_aecnr_get_property (GObject * object, guint prop_id, GValue * value, GParamSpec * pspec)
{
  Gstimx_ai_aecnr *ai_aecnr = GST_IMX_AI_AECNR (object);

  switch (prop_id) {
    case PROP_BYPASS:
      g_value_set_boolean (value, ai_aecnr->bypass);
      break;
    case PROP_RX_PADDING:
      g_value_set_boolean (value, ai_aecnr->rx_padding);
      break;
    default:
      GST_ERROR_OBJECT(ai_aecnr, "unsupported property ID: %d", prop_id);
      break;
  }
}

/* class initialization */
static void gst_imx_ai_aecnr_class_init (Gstimx_ai_aecnrClass * ai_aecnrclass)
{
  GObjectClass *gobject_class = (GObjectClass *)ai_aecnrclass;
  GstElementClass *gstelement_class = (GstElementClass *)ai_aecnrclass;

  GST_DEBUG_CATEGORY_INIT (imx_ai_aecnr, "imx_ai_aecnr", 0, "imx_ai_aecnr plugin");

  gobject_class->set_property = gst_imx_ai_aecnr_set_property;
  gobject_class->get_property = gst_imx_ai_aecnr_get_property;

  parent_class = (GstElementClass *) g_type_class_peek_parent (ai_aecnrclass);
  gstelement_class->change_state = gst_imx_ai_aecnr_change_state;

  g_object_class_install_property (gobject_class,
                                   PROP_BYPASS,
                                   g_param_spec_boolean ("bypass", "Bypass", "Bypass AI AECNR processing.",
                                                         FALSE, G_PARAM_READWRITE));
  g_object_class_install_property (gobject_class,
                                   PROP_RX_PADDING,
                                   g_param_spec_boolean ("rx_padding",
                                                         "RxPadding",
                                                         "Fill Rx with 0 if no data received.",
                                                         TRUE,
                                                         G_PARAM_READWRITE | GST_PARAM_MUTABLE_READY));
  g_object_class_install_property (gobject_class, PROP_MODEL,
                                   g_param_spec_string ("model",
                                                        "model",
                                                        "selection of the model [small,large]",
                                                        "small",
                                                        (GParamFlags)(G_PARAM_WRITABLE | GST_PARAM_MUTABLE_READY | G_PARAM_STATIC_STRINGS)));

  gst_element_class_set_details_simple (gstelement_class,
      "imx_ai_aecnr",
      "filter/audio/ai_aecnr",
      "NXP AI-based acoustic echo canceller and noise suppressor",
      IMX_VOICE_PLUGINS_CONTACT);

  gst_element_class_add_pad_template (gstelement_class, gst_static_pad_template_get (&src_remote_pad_template));
  gst_element_class_add_pad_template (gstelement_class, gst_static_pad_template_get (&sink_remote_pad_template));
  gst_element_class_add_pad_template (gstelement_class, gst_static_pad_template_get (&sink_mic_pad_template));
}


/* Plugin creation with platform check */
static gboolean imx_ai_aecnr_init (GstPlugin * imx_ai_aecnr)
{
  return GST_ELEMENT_REGISTER (imx_ai_aecnr, imx_ai_aecnr);
}

GST_PLUGIN_DEFINE (GST_VERSION_MAJOR,
    GST_VERSION_MINOR,
    imx_ai_aecnr,
    "imx_ai_aecnr",
    imx_ai_aecnr_init,
    IMX_VOICE_PLUGINS_VERSION,
    IMX_VOICE_PLUGINS_LICENSE,
    IMX_VOICE_PLUGINS_BINARY_PACKAGE,
    IMX_VOICE_PLUGINS_URL)