// built on example from https://github.com/vzhn/ffmpeg-java-samples/blob/master/src/main/java/DemuxAndDecodeH264.java

!include once #1035133 // Video Libs [Include]

import java.time.Duration;
import java.time.temporal.ChronoUnit;

import org.bytedeco.javacpp.*;
import static org.bytedeco.javacpp.avcodec.*;
import static org.bytedeco.javacpp.avformat.*;
import static org.bytedeco.javacpp.avutil.*;
import org.bytedeco.javacpp.presets.avutil;

sclass ImageStreamFromVideoFile is Steppable, AutoCloseable {
  // image: video frame as an image. this image is a buffer, i.e.
  // don't hold on to it in the listener after it returns.
  // d: position in video
  event haveImage(BufferedImage image, Duration d);
  
  // set when you are no longer requested in more frames
  settable bool stopped;
  
  /** Matroska format context */
  private AVFormatContext avfmtCtx;

  /** Matroska video stream information  */
  private AVStream videoStream;

  /** matroska packet */
  private AVPacket avpacket;

  /** H264 Decoder ID */
  private AVCodec codec;

  /** H264 Decoder context */
  private AVCodecContext codecContext;

  /** yuv420 frame */
  private AVFrame yuv420Frame;

  /** RGB frame */
  private AVFrame rgbFrame;

  /** java RGB frame */
  private BufferedImage img;

  /** yuv420 to rgb converter */
  private swscale.SwsContext sws_ctx;

  /** number of frames decoded */
  settable long nFrames;
  
  // where are we in the video
  gettable Duration position;
  
  /* 1/1000 of second */
  private AVRational tb1000;
  
  Thread decodeThread;

  {
    tb1000 = new AVRational();
    tb1000.num(1);
    tb1000.den(1000);
  }
  
  void open(File videoFile) throws IOException {
    if (avfmtCtx != null) fail("Already opened a file");
    av_log_set_level(AV_LOG_VERBOSE);

    openInput(videoFile);
    findVideoStream();
    initDecoder();
    initRgbFrame();
    initYuv420Frame();
    getSwsContext();

    avpacket = new avcodec.AVPacket();
  }
  
  int stage;
  
  public bool step() ctex {
    if (stopped()) false;
    if (stage == 0) {
      if ((av_read_frame(avfmtCtx, avpacket)) >= 0) {
        if (avpacket.stream_index() == videoStream.index())
          processAVPacket(avpacket);
        av_packet_unref(avpacket);
      } else
        stage = 1;
      true;
    }
    if (stage == 1) {
      // process delayed frames
      processAVPacket(null);
      stage = 2;
      false;
    }
    false;
  }

  private AVFormatContext openInput(File videoFile) throws IOException {
      avfmtCtx = new AVFormatContext(null);
      BytePointer filePointer = new BytePointer(videoFile.getPath());
      int r = avformat.avformat_open_input(avfmtCtx, filePointer, null, null);
      filePointer.deallocate();
      if (r < 0) {
          avfmtCtx.close();
          throw new IOException("avformat_open_input error: " + r);
      }
      return avfmtCtx;
  }

  private void findVideoStream() throws IOException {
      int r = avformat_find_stream_info(avfmtCtx, (PointerPointer) null);
      if (r < 0) {
          avformat_close_input(avfmtCtx);
          avfmtCtx.close();
          throw new IOException("error: " + r);
      }

      PointerPointer<AVCodec> decoderRet = new PointerPointer<>(1);
      int videoStreamNumber = av_find_best_stream(avfmtCtx, AVMEDIA_TYPE_VIDEO, -1, -1, decoderRet, 0);
      if (videoStreamNumber < 0) {
          throw new IOException("failed to find video stream");
      }

      if (decoderRet.get(AVCodec.class).id() != AV_CODEC_ID_H264) {
          throw new IOException("failed to find h264 stream");
      }
      decoderRet.deallocate();
      videoStream =  avfmtCtx.streams(videoStreamNumber);
  }

  private void initDecoder() {
      codec = avcodec_find_decoder(AV_CODEC_ID_H264);
      codecContext = avcodec_alloc_context3(codec);
      if((codec.capabilities() & avcodec.AV_CODEC_CAP_TRUNCATED) != 0) {
          codecContext.flags(codecContext.flags() | avcodec.AV_CODEC_CAP_TRUNCATED);
      }
      avcodec_parameters_to_context(codecContext, videoStream.codecpar());
      if(avcodec_open2(codecContext, codec, (PointerPointer) null) < 0) {
          throw new RuntimeException("Error: could not open codec.\n");
      }
  }

  private void initYuv420Frame() {
      yuv420Frame = av_frame_alloc();
      if (yuv420Frame == null) {
          throw new RuntimeException("Could not allocate video frame\n");
      }
  }

  private void initRgbFrame() {
    rgbFrame = av_frame_alloc();
    rgbFrame.format(AV_PIX_FMT_BGR24);
    rgbFrame.width(codecContext.width());
    rgbFrame.height(codecContext.height());
    int ret = av_image_alloc(rgbFrame.data(),
            rgbFrame.linesize(),
            rgbFrame.width(),
            rgbFrame.height(),
            rgbFrame.format(),
            1);
    if (ret < 0) {
        throw new RuntimeException("could not allocate buffer!");
    }
    img = new BufferedImage(rgbFrame.width(), rgbFrame.height(), BufferedImage.TYPE_3BYTE_BGR);
  }

  private void getSwsContext() {
    sws_ctx = swscale.sws_getContext(
            codecContext.width(), codecContext.height(), codecContext.pix_fmt(),
            rgbFrame.width(), rgbFrame.height(), rgbFrame.format(),
            0, null, null, (DoublePointer) null);
  }

  private void processAVPacket(AVPacket avpacket) throws IOException {
    int ret = avcodec.avcodec_send_packet(codecContext, avpacket);
    if (ret < 0)
      fail("Error sending a packet for decoding");
    receiveFrames();
  }
  
  // TODO: do in multiple step()s
  private void receiveFrames() throws IOException {
    while (!stopped() && receiveFrame()) {}
  }
  
  bool receiveFrame() throws IOException {
    int ret = avcodec.avcodec_receive_frame(codecContext, yuv420Frame);
    if (ret == avutil.AVERROR_EAGAIN()
      || ret == org.bytedeco.javacpp.avutil.AVERROR_EOF())
      false;
    else if (ret < 0)
      fail("error during video decoding");
      
    swscale.sws_scale(sws_ctx, yuv420Frame.data(), yuv420Frame.linesize(), 0,
      yuv420Frame.height(), rgbFrame.data(), rgbFrame.linesize());

    rgbFrame.best_effort_timestamp(yuv420Frame.best_effort_timestamp());
    processFrame(rgbFrame);
    true;
  }

  private void processFrame(AVFrame rgbFrame) {
    DataBufferByte buffer = cast img.getRaster().getDataBuffer();
    rgbFrame.data(0).get(buffer.getData());
    
    long ptsMillis = av_rescale_q(rgbFrame.best_effort_timestamp(), videoStream.time_base(), tb1000);
    position = Duration.of(ptsMillis, ChronoUnit.MILLIS);

    ++nFrames;
    haveImage(img, position);
  }

  close {
    stopped(true);
    if (decodeThread == null)
      _actualClose();
  }
  
  void _actualClose {
    if (avpacket != null) av_packet_unref(avpacket);
    if (codecContext != null) {
      avcodec.avcodec_close(codecContext);
      avcodec.avcodec_free_context(codecContext);
    }

    swscale.sws_freeContext(sws_ctx);
    if (rgbFrame != null) av_frame_free(rgbFrame);
    if (yuv420Frame != null) av_frame_free(yuv420Frame);
    if (avfmtCtx != null) {
      avformat.avformat_close_input(avfmtCtx);
      avformat.avformat_free_context(avfmtCtx);
    }
  }
  
  void stop {
    stopped(true);
    interruptThread(decodeThread);
  }
  
  void startDecodeThread {
    if (decodeThread != null) ret;
    decodeThread = startThread("Decode video", r {
      stepAll(ImageStreamFromVideoFile.this);
      if (stopped()) _actualClose();
    });
  }
} // end of ImageStreamFromVideoFile