package com.thatgamerblue.jcodectest;

import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import org.jcodec.codecs.h264.H264Decoder;
import org.jcodec.codecs.h264.H264Utils;
import org.jcodec.codecs.h264.io.model.NALUnit;
import org.jcodec.codecs.h264.io.model.NALUnitType;
import org.jcodec.codecs.h264.io.model.SeqParameterSet;
import org.jcodec.common.AutoFileChannelWrapper;
import org.jcodec.common.Codec;
import org.jcodec.common.DemuxerTrackMeta;
import org.jcodec.common.SeekableDemuxerTrack;
import org.jcodec.common.model.ColorSpace;
import org.jcodec.common.model.Packet;
import org.jcodec.common.model.Picture;
import org.jcodec.common.model.Size;
import org.jcodec.containers.mkv.demuxer.MKVDemuxer;
import org.jcodec.javase.scale.AWTUtil;

/**
 * @author ThatGamerBlue
 * Licensed under WTFPL, see: http://www.wtfpl.net/
 */
public class MKVDecoder
{
	private JCodecMKVFrameGrab frameGrab;

	public MKVDecoder(final File file, final boolean repeating)
	{
		try
		{
			frameGrab = new JCodecMKVFrameGrab(new AutoFileChannelWrapper(file), repeating);
		}
		catch (IOException | IllegalArgumentException e)
		{
			System.err.printf("Failed to open file %s for reading MKV video%n", file.toString());
			e.printStackTrace();
		}
	}

	public BufferedImage getNextFrame()
	{
		try
		{
			return frameGrab.getFrame();
		}
		catch (IOException e)
		{
			System.err.printf("Failed to read frame from MKV video using JCodec%n");
			e.printStackTrace();
		}
		return null;
	}

	static class JCodecMKVFrameGrab
	{
		private static final int PRELOAD_FRAMES = 9;
			// covers the most common patterns of bframes and iframes interleaved

		private final SeekableDemuxerTrack videoTrack;
		private final H264Decoder decoder;
		private final TreeMap<Double, BufferedImage> preloadedFrames;
		private final TreeMap<Double, BufferedImage> initialFrames;
		private final ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
		private final boolean repeating;

		private byte[][] buffers;

		public JCodecMKVFrameGrab(AutoFileChannelWrapper afcw, boolean repeating) throws IOException
		{
			this.repeating = repeating;

			MKVDemuxer demuxer = new MKVDemuxer(afcw);

			if (demuxer.getVideoTracks().isEmpty())
			{
				throw new IllegalArgumentException("MKV file has no video track.");
			}

			videoTrack = (SeekableDemuxerTrack) demuxer.getVideoTracks().get(0);

			if (videoTrack.getMeta().getCodec() != Codec.H264)
			{
				throw new IllegalArgumentException("MKV file must be in H264 format.");
			}

			decoder = H264Decoder.createH264DecoderFromCodecPrivate(videoTrack.getMeta().getCodecPrivate());
			preloadedFrames = new TreeMap<>(Double::compare);
			initialFrames = new TreeMap<>(Double::compare);

			bufferFrames();
		}

		private void bufferFrames() throws IOException
		{
			boolean loadingFirst = false;
			if (initialFrames.isEmpty())
			{
				loadingFirst = true;
			}

			while (preloadedFrames.size() < PRELOAD_FRAMES)
			{
				Packet frame = videoTrack.nextFrame();
				if (frame == null)
				{
					break;
				}

				BufferedImage image = AWTUtil.toBufferedImage(decoder.decodeFrame(frame.getData(), getBuffer()));
				preloadedFrames.put(frame.getPtsD(), image);
				if (loadingFirst)
				{
					initialFrames.put(frame.getPtsD(), image);
				}
			}
		}

		public BufferedImage getFrame() throws IOException
		{
			// deals with interleaved frames, preventing jittery playback
			bufferFrames();

			Map.Entry<Double, BufferedImage> retn = preloadedFrames.firstEntry();
			if (retn == null)
			{
				// we're at the end of the buffer
				if (!repeating)
				{
					return null;
				}

				preloadedFrames.putAll(initialFrames);
				// skip initial frames when next buffering
				videoTrack.gotoFrame(initialFrames.size());
				// refresh loaded frame
				retn = preloadedFrames.firstEntry();
			}
			preloadedFrames.remove(retn.getKey());
			return retn.getValue();
		}

		private byte[][] getBuffer()
		{
			if (buffers == null)
			{
				Size size = calcBufferSize();
				buffers = Picture.create(size.getWidth(), size.getHeight(), ColorSpace.YUV444).getData();
			}
			return buffers;
		}

		private Size calcBufferSize()
		{
			// declare output variables
			int w = Integer.MIN_VALUE;
			int h = Integer.MIN_VALUE;

			// get track metadata
			DemuxerTrackMeta meta = videoTrack.getMeta();

			// allocate buffers to hold metadata
			ByteBuffer codecBuffer = meta.getCodecPrivate().duplicate();
			ByteBuffer nalBuffer;

			// search the metadata for the size units
			while ((nalBuffer = H264Utils.nextNALUnit(codecBuffer)) != null)
			{
				// read unit from the nalBuffer
				NALUnit unit = NALUnit.read(nalBuffer);

				// only interested in Sequence Parameter Set units
				if (unit.type != NALUnitType.SPS)
				{
					continue;
				}

				// read the sps from the buffer
				// luckily jcodec does 99% of the work for us
				SeqParameterSet sps = H264Utils.readSPS(nalBuffer);

				// read the width from the parameters
				int tempW = sps.picWidthInMbsMinus1 + 1;
				if (tempW > w)
				{
					w = tempW;
				}

				// read the height from the parameters
				int tempH = SeqParameterSet.getPicHeightInMbs(sps);
				if (tempH > h)
				{
					h = tempH;
				}
			}

			return new Size(w << 4, h << 4);
		}
	}
}