/* Thin C wrapper around libx264 for the Parrot Tools video encoder (compiled to WebAssembly, single thread).
 * x264 is GPL-2.0-or-later (https://www.videolan.org/developers/x264.html); this wrapper is released under the same licence.
 * Calling convention: pt_init(preset,tune) -> pt_set(key,value)* -> pt_start(w,h,fps_num,fps_den) -> loop { write I420 into pt_input(); pt_encode(pts, idr) }
 * -> pt_flush() until it returns 0 -> pt_close(). Output is length-prefixed (MP4 style) NAL units; pt_headers() returns SPS+PPS. */
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <emscripten.h>
#include "x264.h"

static x264_t *h = NULL;
static x264_param_t param;
static x264_picture_t pin, pout;
static int W = 0, H = 0;
static uint8_t *outbuf = NULL;
static int outcap = 0, outlen = 0, lastkey = 0;
static double lastpts = 0, lastdts = 0;
static int inited = 0;

static void ensure_out(int n) {
  if (n > outcap) { outcap = n * 2; outbuf = (uint8_t *)realloc(outbuf, outcap); }
}

EMSCRIPTEN_KEEPALIVE int pt_init(const char *preset, const char *tune) {
  if (x264_param_default_preset(&param, preset && *preset ? preset : "medium", tune && *tune ? tune : NULL) < 0) return -1;
  param.i_threads = 1;
  param.i_lookahead_threads = 0;
  param.b_sliced_threads = 0;
  param.b_annexb = 0;            /* length-prefixed NAL units, ready for MP4 */
  param.b_repeat_headers = 0;    /* SPS/PPS only in pt_headers() */
  param.b_stitchable = 1;        /* identical headers in every segment, so segments can be joined */
  param.b_vfr_input = 0;
  param.i_log_level = X264_LOG_NONE;
  inited = 1;
  return 0;
}
EMSCRIPTEN_KEEPALIVE int pt_set(const char *k, const char *v) { return x264_param_parse(&param, k, v); }
EMSCRIPTEN_KEEPALIVE int pt_profile(const char *name) { return x264_param_apply_profile(&param, name); }   /* call after pt_set, before pt_start */

EMSCRIPTEN_KEEPALIVE int pt_start(int w, int hh, int fps_num, int fps_den) {
  if (!inited) return -1;
  W = w; H = hh;
  param.i_width = w; param.i_height = hh;
  param.i_csp = X264_CSP_I420;
  param.i_fps_num = fps_num; param.i_fps_den = fps_den;
  param.i_timebase_num = fps_den; param.i_timebase_den = fps_num;   /* pts counts frames */
  param.vui.b_fullrange = 0;
  h = x264_encoder_open(&param);
  if (!h) return -2;
  x264_picture_alloc(&pin, X264_CSP_I420, w, hh);
  return 0;
}

EMSCRIPTEN_KEEPALIVE uint8_t *pt_input(void) { return pin.img.plane[0]; }   /* planes are contiguous: Y (w*h), U, V (w*h/4 each) */

static int collect(x264_nal_t *nals, int n) {
  int total = 0;
  for (int i = 0; i < n; i++) total += nals[i].i_payload;
  ensure_out(total);
  /* x264 returns the NAL units of one access unit back to back in memory */
  if (n) memcpy(outbuf, nals[0].p_payload, total);
  outlen = total;
  return total;
}

EMSCRIPTEN_KEEPALIVE int pt_headers(void) {
  x264_nal_t *nals; int n;
  int r = x264_encoder_headers(h, &nals, &n);
  if (r < 0) return -1;
  return collect(nals, n);
}

static int do_encode(x264_picture_t *p) {
  x264_nal_t *nals; int n = 0;
  int r = x264_encoder_encode(h, &nals, &n, p, &pout);
  if (r < 0) return -1;
  if (r == 0) { outlen = 0; return 0; }
  lastkey = pout.b_keyframe ? 1 : 0;
  lastpts = (double)pout.i_pts; lastdts = (double)pout.i_dts;
  return collect(nals, n);
}
EMSCRIPTEN_KEEPALIVE int pt_encode(double pts, int idr) {
  pin.i_pts = (int64_t)pts;
  pin.i_type = idr ? X264_TYPE_IDR : X264_TYPE_AUTO;
  return do_encode(&pin);
}
EMSCRIPTEN_KEEPALIVE int pt_flush(void) {
  if (x264_encoder_delayed_frames(h) == 0) return 0;
  return do_encode(NULL);
}
EMSCRIPTEN_KEEPALIVE uint8_t *pt_out(void) { return outbuf; }
EMSCRIPTEN_KEEPALIVE int pt_key(void) { return lastkey; }
EMSCRIPTEN_KEEPALIVE double pt_pts(void) { return lastpts; }
EMSCRIPTEN_KEEPALIVE double pt_dts(void) { return lastdts; }
EMSCRIPTEN_KEEPALIVE void pt_close(void) {
  if (h) { x264_encoder_close(h); h = NULL; x264_picture_clean(&pin); }
  inited = 0;
}