#include "mpp_h264_source.h"

#include "webrtcpush_config.h"
#include "webrtcpush_log.h"

#include <gst/app/gstappsink.h>
#include <gst/video/video-event.h>
#include <string.h>

#define MPP_VIDEO_WIDTH  1280
#define MPP_VIDEO_HEIGHT 720
#define MPP_VIDEO_GOP    WEBRTCPUSH_H264_GOP
#define MPP_VIDEO_START_BPS 1400000u
#define MPP_VIDEO_MIN_BPS   WEBRTCPUSH_MIN_BITRATE
#define MPP_VIDEO_MAX_BPS   WEBRTCPUSH_MAX_BITRATE

#define WRTC_V4L2_CAP_VIDEO_CAPTURE         (1u << 0)
#define WRTC_V4L2_CAP_VIDEO_CAPTURE_MPLANE  (1u << 12)

struct MppH264Source {
  GstElement *pipeline;
  GstElement *enc;
  GstElement *parse;
  GstElement *appsink;
  gboolean use_mpp;
  guint8 *frame_buf;
  size_t frame_size;
  size_t frame_capacity;
  guint keyframe_req;
};

static gchar *pick_v4l2_capture_device(void)
{
  for (int i = 0; i < 32; i++) {
    gchar *path = g_strdup_printf("/dev/video%d", i);
    if (!g_file_test(path, G_FILE_TEST_EXISTS)) {
      g_free(path);
      continue;
    }

    gchar *sys_cap = g_strdup_printf("/sys/class/video4linux/video%d/device_caps", i);
    gboolean use = FALSE;

    if (g_file_test(sys_cap, G_FILE_TEST_EXISTS)) {
      gchar *contents = NULL;
      if (g_file_get_contents(sys_cap, &contents, NULL, NULL)) {
        guint64 caps = g_ascii_strtoull(g_strstrip(contents), NULL, 0);
        use = (caps & WRTC_V4L2_CAP_VIDEO_CAPTURE) != 0 ||
              (caps & WRTC_V4L2_CAP_VIDEO_CAPTURE_MPLANE) != 0;
        g_free(contents);
      }
    } else {
      use = TRUE;
    }

    g_free(sys_cap);
    if (use)
      return path;
    g_free(path);
  }
  return NULL;
}

static void configure_leaky_queue(GstElement *q)
{
  if (!q)
    return;
  g_object_set(q,
      "max-size-buffers", 1,
      "max-size-bytes", 0,
      "max-size-time", (guint64)0,
      "leaky", 2,
      NULL);
}

static void configure_v4l2_low_latency(GstElement *src)
{
  if (!src)
    return;
  g_object_set(src, "do-timestamp", TRUE, NULL);
}

static GstElement *make_h264_encoder(gboolean prefer_mpp, gboolean use_test,
                                     gboolean *use_mpp_out)
{
  GstElement *enc = NULL;

  if (prefer_mpp && !use_test) {
    enc = gst_element_factory_make("mpph264enc", "enc");
    if (enc) {
      *use_mpp_out = TRUE;
      g_object_set(enc,
          "gop", MPP_VIDEO_GOP,
          "rc-mode", 1,
          "bps", MPP_VIDEO_START_BPS,
          "bps-min", (guint)(MPP_VIDEO_START_BPS * 0.80),
          "bps-max", (guint)(MPP_VIDEO_START_BPS * 1.08),
          "header-mode", 1,
          "profile", 66,
          "level", 31,
          "qos", TRUE,
          "zero-copy-pkt", TRUE,
          NULL);
      my_zlog_info("mpp_h264_source: mpph264enc %dx%d@%dfps GOP=%d",
                   MPP_VIDEO_WIDTH, MPP_VIDEO_HEIGHT, WEBRTCPUSH_H264_FPS,
                   MPP_VIDEO_GOP);
      return enc;
    }
    my_zlog_warn("mpp_h264_source: mpph264enc unavailable");
  }

  enc = gst_element_factory_make("x264enc", "enc");
  if (enc) {
    g_object_set(enc,
        "tune", "zerolatency",
        "speed-preset", "ultrafast",
        "bitrate", MPP_VIDEO_START_BPS / 1000,
        "key-int-max", MPP_VIDEO_GOP,
        "byte-stream", TRUE,
        "profile", "baseline",
        NULL);
    my_zlog_info("mpp_h264_source: x264enc software fallback");
    return enc;
  }

  return NULL;
}

static gboolean buffer_is_idr(const guint8 *data, size_t size)
{
  size_t i = 0;

  while (i + 3 < size) {
    size_t sc = 0;
    if (i + 4 <= size && data[i] == 0 && data[i + 1] == 0 &&
        data[i + 2] == 0 && data[i + 3] == 1)
      sc = 4;
    else if (data[i] == 0 && data[i + 1] == 0 && data[i + 2] == 1)
      sc = 3;
    else {
      i++;
      continue;
    }

    i += sc;
    if (i >= size)
      break;
    if ((data[i] & 0x1f) == 5)
      return TRUE;
  }
  return FALSE;
}

static gboolean send_force_key_unit_on_pad(GstPad *pad)
{
  GstEvent *event;

  if (!pad)
    return FALSE;
  event = gst_video_event_new_upstream_force_key_unit(
      GST_CLOCK_TIME_NONE, TRUE, 0);
  return gst_pad_send_event(pad, event);
}

static void request_keyframe(MppH264Source *src)
{
  gboolean ok = FALSE;

  if (!src)
    return;

  src->keyframe_req++;

  /* Upstream force-key-unit 必须交给编码器的 src pad 处理。 */
  if (src->enc) {
    GstPad *pad = gst_element_get_static_pad(src->enc, "src");
    ok = send_force_key_unit_on_pad(pad);
    if (pad)
      gst_object_unref(pad);
  }

  if (ok)
    my_zlog_info("mpp_h264_source: requested IDR (%u)", src->keyframe_req);
  else
    my_zlog_warn("mpp_h264_source: IDR request not handled");
}

void mpp_h264_source_request_keyframe(MppH264Source *src)
{
  request_keyframe(src);
}

static gboolean copy_frame(MppH264Source *src, GstBuffer *buffer)
{
  GstMapInfo map;

  if (!gst_buffer_map(buffer, &map, GST_MAP_READ))
    return FALSE;

  if (map.size > src->frame_capacity) {
    src->frame_buf = g_realloc(src->frame_buf, map.size);
    src->frame_capacity = map.size;
  }
  memcpy(src->frame_buf, map.data, map.size);
  src->frame_size = map.size;
  gst_buffer_unmap(buffer, &map);
  return TRUE;
}

MppH264Source *mpp_h264_source_start(const char *video_device, char **error_message)
{
  MppH264Source *src;
  GstElement *pipe, *vsrc, *caps_jpeg = NULL, *jpegparse = NULL, *jpegdec = NULL;
  GstElement *vrate, *conv, *caps1, *enc, *parse, *q, *caps_h264, *sink;
  gchar *auto_v4l = NULL;
  const gchar *vdev = video_device;
  gboolean use_test = FALSE;
  gboolean use_auto = FALSE;
  gboolean use_mjpg = FALSE;
  gboolean use_mpp = FALSE;
  gboolean use_mpp_jpeg = FALSE;
  GstStateChangeReturn ret;

  if (!gst_is_initialized())
    gst_init(NULL, NULL);

  src = g_new0(MppH264Source, 1);

  use_test = (vdev && g_strcmp0(vdev, "test") == 0);
  use_auto = (!vdev || g_strcmp0(vdev, "auto") == 0);

  if (use_auto && !use_test) {
    auto_v4l = pick_v4l2_capture_device();
    if (auto_v4l) {
      my_zlog_info("mpp_h264_source: auto V4L2 %s", auto_v4l);
      vdev = auto_v4l;
    } else {
      my_zlog_warn("mpp_h264_source: no V4L2 device, videotestsrc");
      use_test = TRUE;
    }
  }

  pipe = gst_pipeline_new("mpp-h264-pipe");
  if (use_test) {
    vsrc = gst_element_factory_make("videotestsrc", "vsrc");
    g_object_set(vsrc, "is-live", TRUE, NULL);
  } else {
    vsrc = gst_element_factory_make("v4l2src", "vsrc");
    g_object_set(vsrc, "device", vdev, NULL);
    configure_v4l2_low_latency(vsrc);
    caps_jpeg = gst_element_factory_make("capsfilter", "caps_jpeg");
    jpegparse = gst_element_factory_make("jpegparse", "jpegparse");
    jpegdec = gst_element_factory_make("mppjpegdec", "jpegdec");
    if (jpegdec) {
      use_mpp_jpeg = TRUE;
      g_object_set(jpegdec, "format", 23, "fast-mode", TRUE, NULL);
    } else {
      jpegdec = gst_element_factory_make("jpegdec", "jpegdec");
    }
    if (caps_jpeg && jpegparse && jpegdec) {
      gchar *jc_str = g_strdup_printf(
          "image/jpeg,width=%d,height=%d,framerate=30/1",
          MPP_VIDEO_WIDTH, MPP_VIDEO_HEIGHT);
      GstCaps *jc = gst_caps_from_string(jc_str);
      g_object_set(caps_jpeg, "caps", jc, NULL);
      gst_caps_unref(jc);
      g_free(jc_str);
      use_mjpg = TRUE;
      my_zlog_info("mpp_h264_source: JPEG decoder=%s",
                   use_mpp_jpeg ? "mppjpegdec" : "jpegdec software fallback");
    }
  }

  conv = gst_element_factory_make("videoconvert", "conv");
  vrate = gst_element_factory_make("videorate", "vrate");
  caps1 = gst_element_factory_make("capsfilter", "caps1");
  q = gst_element_factory_make("queue", "vq");
  enc = make_h264_encoder(TRUE, use_test, &use_mpp);
  parse = gst_element_factory_make("h264parse", "parse");
  caps_h264 = gst_element_factory_make("capsfilter", "caps_h264");
  sink = gst_element_factory_make("appsink", "sink");

  if (!pipe || !vsrc || !conv || !vrate || !caps1 || !q || !enc || !parse ||
      !caps_h264 || !sink || (!use_test && !use_mjpg)) {
    g_set_error_literal(NULL, 0, 0, "GStreamer element create failed");
    if (error_message)
      *error_message = g_strdup("GStreamer element create failed");
    g_free(auto_v4l);
    g_free(src);
    return NULL;
  }

  g_object_set(vrate, "max-rate", WEBRTCPUSH_H264_FPS, "drop-only", TRUE, NULL);
  configure_leaky_queue(q);
  g_object_set(parse, "config-interval", -1, NULL);

  {
    gchar *cap_str = g_strdup_printf(
        "video/x-raw,format=%s,width=%d,height=%d,framerate=%d/1",
        use_mpp ? "NV12" : "I420",
        use_test ? 640 : MPP_VIDEO_WIDTH,
        use_test ? 480 : MPP_VIDEO_HEIGHT,
        WEBRTCPUSH_H264_FPS);
    GstCaps *caps = gst_caps_from_string(cap_str);
    g_object_set(caps1, "caps", caps, NULL);
    gst_caps_unref(caps);
    g_free(cap_str);
  }

  {
    GstCaps *hc = gst_caps_from_string(
        "video/x-h264,stream-format=byte-stream,alignment=au");
    g_object_set(caps_h264, "caps", hc, NULL);
    gst_caps_unref(hc);
  }

  g_object_set(sink,
      "emit-signals", FALSE,
      "sync", FALSE,
      "max-buffers", 1,
      "drop", TRUE,
      NULL);

  if (use_mjpg) {
    gst_bin_add_many(GST_BIN(pipe), vsrc, caps_jpeg, jpegparse, jpegdec, conv, vrate,
                     caps1, q, enc, parse, caps_h264, sink, NULL);
    if (!gst_element_link_many(vsrc, caps_jpeg, jpegparse, jpegdec, conv, vrate,
                               caps1, q, enc, parse, caps_h264, sink, NULL)) {
      if (error_message)
        *error_message = g_strdup("failed to link MJPG capture chain");
      gst_object_unref(pipe);
      g_free(auto_v4l);
      g_free(src);
      return NULL;
    }
  } else {
    gst_bin_add_many(GST_BIN(pipe), vsrc, conv, vrate, caps1, q, enc, parse,
                     caps_h264, sink, NULL);
    if (!gst_element_link_many(vsrc, conv, vrate, caps1, q, enc, parse,
                               caps_h264, sink, NULL)) {
      if (error_message)
        *error_message = g_strdup("failed to link video chain");
      gst_object_unref(pipe);
      g_free(auto_v4l);
      g_free(src);
      return NULL;
    }
  }

  src->pipeline = pipe;
  src->enc = enc;
  src->parse = parse;
  src->appsink = sink;
  src->use_mpp = use_mpp;

  gst_pipeline_set_latency(GST_PIPELINE(pipe), 0);
  ret = gst_element_set_state(pipe, GST_STATE_PLAYING);
  if (ret == GST_STATE_CHANGE_FAILURE) {
    if (error_message)
      *error_message = g_strdup("pipeline failed to reach PLAYING");
    mpp_h264_source_stop(src);
    g_free(auto_v4l);
    return NULL;
  }

  my_zlog_info("mpp_h264_source: live encoder started (mpp=%d)", use_mpp ? 1 : 0);
  g_free(auto_v4l);
  return src;
}

static gboolean pull_sample(MppH264Source *src, gboolean require_idr,
                            gint timeout_ms, const uint8_t **data,
                            size_t *size, gboolean *is_idr,
                            guint64 *pts_ns)
{
  GstAppSink *appsink;
  GstSample *sample;
  GstBuffer *buffer;
  gboolean idr;
  GstClockTime pts;
  gint64 deadline;
  gint attempts = 0;

  if (!src || !src->appsink || !data || !size || timeout_ms < 1)
    return FALSE;

  appsink = GST_APP_SINK(src->appsink);
  deadline = g_get_monotonic_time() + (gint64)timeout_ms * 1000;

  while (attempts < 40) {
    gint wait_ms = (gint)((deadline - g_get_monotonic_time()) / 1000);

    if (wait_ms < 1)
      break;

    sample = gst_app_sink_try_pull_sample(appsink, (guint64)wait_ms * GST_MSECOND);
    if (!sample) {
      if (g_get_monotonic_time() >= deadline)
        break;
      continue;
    }

    buffer = gst_sample_get_buffer(sample);
    if (!buffer || !copy_frame(src, buffer)) {
      gst_sample_unref(sample);
      attempts++;
      continue;
    }

    pts = GST_BUFFER_PTS(buffer);
    idr = buffer_is_idr(src->frame_buf, src->frame_size);
    gst_sample_unref(sample);

    if (require_idr && !idr) {
      attempts++;
      continue;
    }

    *data = src->frame_buf;
    *size = src->frame_size;
    if (is_idr)
      *is_idr = idr;
    if (pts_ns)
      *pts_ns = GST_CLOCK_TIME_IS_VALID(pts) ? (guint64)pts : G_MAXUINT64;
    return TRUE;
  }

  return FALSE;
}

gboolean mpp_h264_source_pull(MppH264Source *src, gboolean force_idr,
                              const uint8_t **data, size_t *size,
                              gboolean *is_idr, guint64 *pts_ns,
                              gint timeout_ms)
{
  gint idr_timeout;

  if (!src || !src->appsink || !data || !size)
    return FALSE;

  if (timeout_ms < 1)
    timeout_ms = 80;

  if (force_idr) {
    idr_timeout = timeout_ms < 2000 ? 2000 : timeout_ms;
    request_keyframe(src);
    if (pull_sample(src, TRUE, idr_timeout, data, size, is_idr, pts_ns))
      return TRUE;
    my_zlog_warn("mpp_h264_source: IDR wait timeout, fallback to latest frame");
    return pull_sample(src, FALSE, 200, data, size, is_idr, pts_ns);
  }

  return pull_sample(src, FALSE, timeout_ms, data, size, is_idr, pts_ns);
}

void mpp_h264_source_set_bitrate(MppH264Source *src, guint bitrate_bps)
{
  GstElementFactory *f;
  const gchar *name;

  if (!src || !src->enc || !bitrate_bps)
    return;

  if (bitrate_bps < MPP_VIDEO_MIN_BPS)
    bitrate_bps = MPP_VIDEO_MIN_BPS;
  if (bitrate_bps > MPP_VIDEO_MAX_BPS)
    bitrate_bps = MPP_VIDEO_MAX_BPS;

  f = gst_element_get_factory(src->enc);
  name = f ? gst_plugin_feature_get_name(GST_PLUGIN_FEATURE(f)) : NULL;
  if (name && g_str_has_prefix(name, "mpph264enc")) {
    guint bmin = (guint)(bitrate_bps * 0.80);
    guint bmax = (guint)(bitrate_bps * 1.08);
    if (bmin < MPP_VIDEO_MIN_BPS)
      bmin = MPP_VIDEO_MIN_BPS;
    if (bmax > MPP_VIDEO_MAX_BPS)
      bmax = MPP_VIDEO_MAX_BPS;
    g_object_set(src->enc, "bps", bitrate_bps, "bps-min", bmin, "bps-max", bmax, NULL);
  } else if (name && g_str_has_prefix(name, "x264enc")) {
    g_object_set(src->enc, "bitrate", bitrate_bps / 1000, NULL);
  }
}

gboolean mpp_h264_source_is_mpp_encoder(const MppH264Source *src)
{
  return src && src->use_mpp;
}

void mpp_h264_source_stop(MppH264Source *src)
{
  if (!src)
    return;

  if (src->pipeline) {
    gst_element_set_state(src->pipeline, GST_STATE_NULL);
    gst_object_unref(src->pipeline);
    src->pipeline = NULL;
  }
  src->enc = NULL;
  src->parse = NULL;
  src->appsink = NULL;
  g_free(src->frame_buf);
  g_free(src);
}
