// ZapMate's software video decoding: FFmpeg's decoders drawn onto an Android Surface.
//
// Frames go three ways at once, so each has one owner. The decoder thread drains FFmpeg into `ready`. A frame
// taken for an output buffer moves to a numbered slot until Java releases it. The playback thread only
// reads a slot, to copy it to the window. A frame that is not plain 8-bit 4:2:0 with limited range is converted
// on the decoder thread, so the window copy is the same for every format. An interlaced frame is made whole
// there too.
#include <android/log.h>
#include <android/native_window.h>
#include <android/native_window_jni.h>
#include <dlfcn.h>
#include <jni.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include <sys/system_properties.h>

#include <deque>
#include <mutex>
#include <vector>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavutil/imgutils.h>
#include <libavutil/pixdesc.h>
#include <libswscale/swscale.h>
}

#define TAG "zap_video"
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, TAG, __VA_ARGS__)

#define VIDEO_FUNC(RETURN_TYPE, NAME, ...)                                           \
  extern "C" JNIEXPORT RETURN_TYPE JNICALL                                           \
      Java_androidx_media3_decoder_ffmpeg_FfmpegVideoDecoder_##NAME(JNIEnv* env,     \
                                                                    jobject thiz,    \
                                                                    ##__VA_ARGS__)

namespace {

constexpr int kFormatYV12 = 0x32315659;
constexpr jlong kNoFrame = INT64_MIN;
constexpr size_t kMaxReady = 32;
// ADATASPACE values of <android/data_space.h>: the matrix a window applies to turn YUV into colour.
constexpr int32_t kDataSpaceBt709 = 281083904;
constexpr int32_t kDataSpaceBt601 = 281149440;

typedef int32_t (*SetDataSpace)(ANativeWindow*, int32_t);

// ANativeWindow_setBuffersDataSpace exists from Android 9. This app runs from Android 7.
SetDataSpace FindSetDataSpace() {
  static SetDataSpace fn = reinterpret_cast<SetDataSpace>(dlsym(RTLD_DEFAULT, "ANativeWindow_setBuffersDataSpace"));
  return fn;
}

// Drawing on a window leaves it connected as "the CPU" until someone disconnects it, and while it is, a device
// decoder cannot connect: the channel after a software channel would fail to start. The NDK has no call for
// this. Android's own is native_window_api_disconnect, in system/window.h: window->perform(window, 14, 2).
// This is the start of that struct. Graphics drivers are compiled against it, so it holds on every Android
// version, and a window that is not what this expects is left alone.
struct WindowLayout {
  int magic;
  int version;
  void* reserved[4];
  void (*inc_ref)(void*);
  void (*dec_ref)(void*);
  uint32_t flags;
  int min_swap_interval;
  int max_swap_interval;
  float xdpi;
  float ydpi;
  intptr_t oem[4];
  void* set_swap_interval;
  void* dequeue_buffer;
  void* lock_buffer;
  void* queue_buffer;
  void* query;
  int (*perform)(ANativeWindow*, int, ...);
};

constexpr int kWindowMagic = ('_' << 24) | ('w' << 16) | ('n' << 8) | 'd';
constexpr int kPerformDisconnect = 14;  // NATIVE_WINDOW_API_DISCONNECT
constexpr int kApiCpu = 2;              // NATIVE_WINDOW_API_CPU

// An ARM build under translation, on a PC-based Android: the window's functions are the machine's own code, and
// calling one directly from translated code kills the app. Better a next channel that fails than that.
bool IsTranslated() {
#if defined(__arm__) || defined(__aarch64__)
  char bridge[PROP_VALUE_MAX] = {};
  return __system_property_get("ro.dalvik.vm.native.bridge", bridge) > 0 && strcmp(bridge, "0") != 0;
#else
  return false;
#endif
}

void DisconnectFromCpu(ANativeWindow* window) {
  if (IsTranslated()) return;
  const auto* layout = reinterpret_cast<const WindowLayout*>(window);
  const size_t needed = offsetof(WindowLayout, perform) + sizeof(layout->perform);
  if (layout->magic != kWindowMagic || layout->version < static_cast<int>(needed) || layout->perform == nullptr) {
    LOGE("unknown window layout, a device decoder may not be able to use it next");
    return;
  }
  layout->perform(window, kPerformDisconnect, kApiCpu);
}

struct VideoContext {
  AVCodecContext* codec = nullptr;
  AVPacket* packet = nullptr;
  AVFrame* decoded = nullptr;
  SwsContext* sws = nullptr;
  AVFrame* previous = nullptr;  // the last interlaced frame as decoded, to tell what moved since
  std::deque<AVFrame*> ready;

  std::mutex slots_lock;
  std::vector<AVFrame*> slots;

  ANativeWindow* window = nullptr;
  jobject surface = nullptr;  // a global reference
  int window_width = 0;
  int window_height = 0;
  int32_t window_space = -1;
  bool window_drawn = false;

  void FreeReady() {
    for (AVFrame* f : ready) av_frame_free(&f);
    ready.clear();
    av_frame_free(&previous);
  }
};

VideoContext* From(jlong handle) { return reinterpret_cast<VideoContext*>(handle); }

bool IsPlain(const AVFrame* f) { return f->format == AV_PIX_FMT_YUV420P && f->color_range != AVCOL_RANGE_JPEG; }

// A new frame in plain 8-bit 4:2:0, or null when conversion fails.
AVFrame* ToPlain(VideoContext* c, const AVFrame* in) {
  AVFrame* out = av_frame_alloc();
  if (out == nullptr) return nullptr;
  out->format = AV_PIX_FMT_YUV420P;
  out->width = in->width;
  out->height = in->height;
  if (av_frame_get_buffer(out, 0) < 0) {
    av_frame_free(&out);
    return nullptr;
  }
  c->sws = sws_getCachedContext(c->sws, in->width, in->height, static_cast<AVPixelFormat>(in->format), out->width,
                                out->height, AV_PIX_FMT_YUV420P, SWS_FAST_BILINEAR, nullptr, nullptr, nullptr);
  if (c->sws == nullptr) {
    av_frame_free(&out);
    return nullptr;
  }
  const int* table = sws_getCoefficients(in->colorspace == AVCOL_SPC_UNSPECIFIED ? SWS_CS_DEFAULT : in->colorspace);
  sws_setColorspaceDetails(c->sws, table, in->color_range == AVCOL_RANGE_JPEG || in->format == AV_PIX_FMT_YUVJ420P,
                           table, 0, 0, 1 << 16, 1 << 16);
  sws_scale(c->sws, in->data, in->linesize, 0, in->height, out->data, out->linesize);
  av_frame_copy_props(out, in);
  out->color_range = AVCOL_RANGE_MPEG;
  return out;
}

// A broadcast frame can be two half-pictures taken 1/50 s apart, on alternate lines. Shown as one, whatever moved
// between them has comb teeth. The lines of the first half-picture are kept. A pixel of the other is kept where
// nothing moved since the frame before, so a still picture keeps all its lines, and where it does not stick
// out from the lines above and below; elsewhere it becomes the mean of those two.
constexpr int kMoved = 8;   // levels of change since the last frame that count as movement
constexpr int kComb = 60;   // (above - this) * (below - this): both off the same way, by about 8 levels each

inline int Apart(int a, int b) { return a > b ? a - b : b - a; }

void DeinterlacePlane(const uint8_t* in, int in_stride, const uint8_t* prev, int prev_stride, uint8_t* out,
                      int out_stride, int width, int height) {
  for (int y = 0; y < height; y++) {
    const uint8_t* cur = in + y * in_stride;
    uint8_t* to = out + y * out_stride;
    if ((y & 1) == 0) {
      memcpy(to, cur, width);
      continue;
    }
    const uint8_t* above = cur - in_stride;
    const uint8_t* below = y + 1 < height ? cur + in_stride : above;
    if (prev == nullptr) {
      for (int x = 0; x < width; x++) {
        const int a = above[x], b = cur[x], d = below[x];
        to[x] = static_cast<uint8_t>((a - b) * (d - b) > kComb ? (a + d + 1) >> 1 : b);
      }
      continue;
    }
    const uint8_t* was = prev + y * prev_stride;
    const uint8_t* was_above = was - prev_stride;
    const uint8_t* was_below = y + 1 < height ? was + prev_stride : was_above;
    for (int x = 0; x < width; x++) {
      const int a = above[x], b = cur[x], d = below[x];
      const int moved = (Apart(b, was[x]) >= kMoved) | (Apart(a, was_above[x]) >= kMoved) | (Apart(d, was_below[x]) >= kMoved);
      const int comb = (a - b) * (d - b) > kComb;
      to[x] = static_cast<uint8_t>((moved & comb) ? (a + d + 1) >> 1 : b);
    }
  }
}

// A new whole frame from an interlaced one in plain 4:2:0, or null when out of memory. Chroma lines alternate
// between the half-pictures the same way, so each plane is treated alike.
AVFrame* Deinterlace(VideoContext* c, const AVFrame* in) {
  AVFrame* out = av_frame_alloc();
  if (out == nullptr) return nullptr;
  out->format = AV_PIX_FMT_YUV420P;
  out->width = in->width;
  out->height = in->height;
  if (av_frame_get_buffer(out, 0) < 0) {
    av_frame_free(&out);
    return nullptr;
  }
  av_frame_copy_props(out, in);
  out->interlaced_frame = 0;
  const AVFrame* prev = c->previous;
  if (prev != nullptr && (prev->width != in->width || prev->height != in->height)) prev = nullptr;
  for (int p = 0; p < 3; p++) {
    const int width = p == 0 ? in->width : (in->width + 1) / 2;
    const int height = p == 0 ? in->height : (in->height + 1) / 2;
    DeinterlacePlane(in->data[p], in->linesize[p], prev != nullptr ? prev->data[p] : nullptr,
                     prev != nullptr ? prev->linesize[p] : 0, out->data[p], out->linesize[p], width, height);
  }
  return out;
}

// Everything FFmpeg has finished goes into `ready`, in display order.
void Drain(VideoContext* c) {
  while (avcodec_receive_frame(c->codec, c->decoded) == 0) {
    AVFrame* kept = nullptr;
    if (IsPlain(c->decoded)) {
      kept = av_frame_alloc();
      if (kept != nullptr) av_frame_move_ref(kept, c->decoded);
    } else {
      kept = ToPlain(c, c->decoded);
    }
    av_frame_unref(c->decoded);
    if (kept == nullptr) continue;
    if (kept->interlaced_frame) {
      AVFrame* whole = Deinterlace(c, kept);
      av_frame_free(&c->previous);
      c->previous = kept;
      kept = whole;
      if (kept == nullptr) continue;
    } else {
      av_frame_free(&c->previous);
    }
    kept->pts = kept->best_effort_timestamp != AV_NOPTS_VALUE ? kept->best_effort_timestamp : kept->pts;
    c->ready.push_back(kept);
    while (c->ready.size() > kMaxReady) {
      AVFrame* old = c->ready.front();
      c->ready.pop_front();
      av_frame_free(&old);
    }
  }
}

// One plane to the window, line by line: the two strides rarely agree.
void CopyPlane(const uint8_t* from, int from_stride, uint8_t* to, int to_stride, int width, int height) {
  for (int y = 0; y < height; y++) {
    memcpy(to, from, width);
    from += from_stride;
    to += to_stride;
  }
}

void ReleaseWindow(JNIEnv* env, VideoContext* c) {
  if (c->window != nullptr) {
    if (c->window_drawn) DisconnectFromCpu(c->window);
    ANativeWindow_release(c->window);
  }
  if (c->surface != nullptr) env->DeleteGlobalRef(c->surface);
  c->window = nullptr;
  c->surface = nullptr;
  c->window_width = c->window_height = 0;
  c->window_space = -1;
  c->window_drawn = false;
}

}  // namespace

VIDEO_FUNC(jlong, nativeInit, jstring jname, jbyteArray jextra, jint threads) {
  const char* name = env->GetStringUTFChars(jname, nullptr);
  const AVCodec* decoder = avcodec_find_decoder_by_name(name);
  env->ReleaseStringUTFChars(jname, name);
  if (decoder == nullptr) return 0;

  auto* c = new VideoContext();
  c->codec = avcodec_alloc_context3(decoder);
  c->packet = av_packet_alloc();
  c->decoded = av_frame_alloc();
  if (c->codec == nullptr || c->packet == nullptr || c->decoded == nullptr) goto fail;

  c->codec->thread_count = threads;
  c->codec->thread_type = FF_THREAD_FRAME | FF_THREAD_SLICE;
  c->codec->pkt_timebase = AVRational{1, 1000000};
  if (jextra != nullptr) {
    const jsize size = env->GetArrayLength(jextra);
    if (size > 0) {
      c->codec->extradata = static_cast<uint8_t*>(av_mallocz(size + AV_INPUT_BUFFER_PADDING_SIZE));
      if (c->codec->extradata == nullptr) goto fail;
      c->codec->extradata_size = size;
      env->GetByteArrayRegion(jextra, 0, size, reinterpret_cast<jbyte*>(c->codec->extradata));
    }
  }
  if (avcodec_open2(c->codec, decoder, nullptr) < 0) goto fail;
  return reinterpret_cast<jlong>(c);

fail:
  avcodec_free_context(&c->codec);
  av_packet_free(&c->packet);
  av_frame_free(&c->decoded);
  delete c;
  return 0;
}

VIDEO_FUNC(void, nativeClose, jlong handle) {
  VideoContext* c = From(handle);
  if (c == nullptr) return;
  c->FreeReady();
  {
    std::lock_guard<std::mutex> hold(c->slots_lock);
    for (AVFrame*& f : c->slots) av_frame_free(&f);
    c->slots.clear();
  }
  ReleaseWindow(env, c);
  sws_freeContext(c->sws);
  avcodec_free_context(&c->codec);
  av_packet_free(&c->packet);
  av_frame_free(&c->decoded);
  delete c;
}

VIDEO_FUNC(void, nativeFlush, jlong handle) {
  VideoContext* c = From(handle);
  avcodec_flush_buffers(c->codec);
  c->FreeReady();
}

// A broken packet is normal on a broadcast, or on joining one mid-picture: it is skipped, and only running
// out of memory is an error.
VIDEO_FUNC(jint, nativeDecode, jlong handle, jobject jdata, jint size, jlong time_us) {
  VideoContext* c = From(handle);
  auto* data = static_cast<uint8_t*>(env->GetDirectBufferAddress(jdata));
  if (av_new_packet(c->packet, size) < 0) return -1;
  memcpy(c->packet->data, data, size);
  c->packet->pts = time_us;
  int sent = avcodec_send_packet(c->codec, c->packet);
  if (sent == AVERROR(EAGAIN)) {
    Drain(c);
    sent = avcodec_send_packet(c->codec, c->packet);
  }
  av_packet_unref(c->packet);
  Drain(c);
  return sent == AVERROR(ENOMEM) ? -1 : 0;
}

VIDEO_FUNC(jlong, nativeNextTime, jlong handle) {
  VideoContext* c = From(handle);
  return c->ready.empty() ? kNoFrame : c->ready.front()->pts;
}

VIDEO_FUNC(void, nativeDropNext, jlong handle) {
  VideoContext* c = From(handle);
  if (c->ready.empty()) return;
  AVFrame* f = c->ready.front();
  c->ready.pop_front();
  av_frame_free(&f);
}

// Moves the next frame into a slot. `size` gets its width as shown (a broadcast's non-square pixels widened)
// and its height, which is what the player sizes the picture by.
VIDEO_FUNC(jint, nativeTake, jlong handle, jintArray jsize) {
  VideoContext* c = From(handle);
  if (c->ready.empty()) return -1;
  AVFrame* f = c->ready.front();
  c->ready.pop_front();

  AVRational sar = f->sample_aspect_ratio.num > 0 ? f->sample_aspect_ratio : c->codec->sample_aspect_ratio;
  int shown = f->width;
  if (sar.num > 0 && sar.den > 0 && sar.num != sar.den) {
    shown = static_cast<int>((static_cast<int64_t>(f->width) * sar.num + sar.den / 2) / sar.den);
  }
  const jint size[2] = {shown, f->height};
  env->SetIntArrayRegion(jsize, 0, 2, size);

  std::lock_guard<std::mutex> hold(c->slots_lock);
  for (size_t i = 0; i < c->slots.size(); i++) {
    if (c->slots[i] == nullptr) {
      c->slots[i] = f;
      return static_cast<jint>(i);
    }
  }
  c->slots.push_back(f);
  return static_cast<jint>(c->slots.size() - 1);
}

VIDEO_FUNC(void, nativeRelease, jlong handle, jint id) {
  VideoContext* c = From(handle);
  std::lock_guard<std::mutex> hold(c->slots_lock);
  if (id >= 0 && static_cast<size_t>(id) < c->slots.size()) av_frame_free(&c->slots[id]);
}

// A window that goes away mid-frame (the preview handing over to full screen) costs one frame and no error.
VIDEO_FUNC(void, nativeRender, jlong handle, jobject jsurface, jint id) {
  VideoContext* c = From(handle);
  if (c->surface == nullptr || !env->IsSameObject(c->surface, jsurface)) {
    ReleaseWindow(env, c);
    c->window = ANativeWindow_fromSurface(env, jsurface);
    if (c->window == nullptr) return;
    c->surface = env->NewGlobalRef(jsurface);
  }

  std::lock_guard<std::mutex> hold(c->slots_lock);
  if (id < 0 || static_cast<size_t>(id) >= c->slots.size() || c->slots[id] == nullptr) return;
  const AVFrame* f = c->slots[id];

  if (c->window_width != f->width || c->window_height != f->height) {
    if (ANativeWindow_setBuffersGeometry(c->window, f->width, f->height, kFormatYV12) != 0) return;
    c->window_width = f->width;
    c->window_height = f->height;
  }
  // HD is BT.709 and SD is BT.601 unless the stream says otherwise. The wrong matrix shifts every colour.
  const bool hd = f->colorspace == AVCOL_SPC_BT709 || (f->colorspace == AVCOL_SPC_UNSPECIFIED && f->height >= 720);
  const int32_t space = hd ? kDataSpaceBt709 : kDataSpaceBt601;
  if (space != c->window_space) {
    SetDataSpace set = FindSetDataSpace();
    if (set != nullptr) set(c->window, space);
    c->window_space = space;
  }

  ANativeWindow_Buffer out;
  if (ANativeWindow_lock(c->window, &out, nullptr) != 0) return;
  c->window_drawn = true;
  if (out.bits == nullptr) {
    ANativeWindow_unlockAndPost(c->window);
    return;
  }
  const int width = f->width < out.width ? f->width : out.width;
  const int height = f->height < out.height ? f->height : out.height;
  const int uv_width = (width + 1) / 2;
  const int uv_height = (height + 1) / 2;
  // Some graphics drivers (the Android emulator's) report a stride of 0 for YV12. The format itself fixes
  // the layout then: lines aligned to 16 bytes, chroma half as wide and aligned to 16 again.
  const int y_stride = out.stride > 0 ? out.stride : (out.width + 15) & ~15;
  const int uv_stride = ((y_stride / 2) + 15) & ~15;
  auto* y = static_cast<uint8_t*>(out.bits);
  uint8_t* v = y + y_stride * out.height;                // YV12 keeps V before U
  uint8_t* u = v + uv_stride * ((out.height + 1) / 2);
  CopyPlane(f->data[0], f->linesize[0], y, y_stride, width, height);
  CopyPlane(f->data[2], f->linesize[2], v, uv_stride, uv_width, uv_height);
  CopyPlane(f->data[1], f->linesize[1], u, uv_stride, uv_width, uv_height);
  ANativeWindow_unlockAndPost(c->window);
}