//! # Simple yet efficient video encoding
//!
//! This is a lossless video codec, aiming at screen capture and anything else you'd want to format in PNG instead of JPEG. It is the equivalent of PNG for video.
//!
//! As PNG, there are two steps: filtering, and compressing.
//!
//! The filter transforms the pixels into the differences between adjacent pixels. This way adjacent pixels with the same color have the value 0, and linear gradients become constant series. The filtered image is now easier to compress with a conventional algorithm.
//!
//! The difference with PNG is that the filter also takes the difference across time, so that if the frames are almost the same, they will be full of zeros.
//!
//! ⚠ **unstable, not benchmarked, not fuzz-tested**... but promising 😎
//!
//! ```rust
//! let mut buf = Vec::new();
//! let mut encoder = syeve::Encoder::new((640, 480), 3, syeve::Compression::Brotli(4), 30);
//! let mut decoder = syeve::Decoder::new();
//!
//! for i in 0..10 {
//!     let mut frame: Vec<u8> = (0..640*480*3).map(|j| ((i+j)%256) as u8).collect(); // dummy frame
//!     let frame_bak = frame.clone();
//!
//!     encoder.encode(&mut frame, &mut buf).unwrap();
//!
//!     // [...] Transfer `buf` to decoder (e.g. UDP or file storage or avian carrier)
//!
//!     let frame = decoder.decode(&buf).unwrap().pixels;
//!
//!     assert_eq!(frame, frame_bak);
//!
//!     buf.clear();
//! }
//! ```

mod filter;

use std::{
	convert::TryInto,
	io::{Read, Write},
};

#[cfg(feature = "deflate")]
pub use deflate::Compression as Deflate;

/// Compression algo
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Compression {
	Brotli(
		/// From 0 (faster) to 11 (better)
		u32,
	),
	#[cfg(feature = "deflate")]
	Deflate(deflate::Compression),
}

impl Default for Compression {
	fn default() -> Self {
		Self::Brotli(4)
	}
}

#[derive(Debug)]
pub enum EncodeError {
	CompressError(std::io::Error),
	WriteError(std::io::Error),
}

#[derive(Clone)]
pub struct Encoder {
	compression: Compression,
	frame_number: u32,
	old_pixels: Vec<u8>,
	pixel_size: usize,
	seq_len: u32,
	seq_number: u32,
	size: (usize, usize),
}

impl Encoder {
	/// Creates a new encoder
	///
	/// `pixel_size` is the number of bytes per pixel.
	///
	/// Because each frame is encoded relatively to the previous frame, a lost frame can cause all the next frames to be wrong. Hence, we need to periodically send a frame that does not depend on the previous frame. `seq_len` (sequence length) is the number of frames between these "independent" frames. The higher, the longer glitches can remain. The lower, the more bandwidth will be used.
	pub fn new(
		size: (usize, usize),
		pixel_size: usize,
		compression: Compression,
		seq_len: u32,
	) -> Self {
		Self {
			compression,
			frame_number: 0,
			old_pixels: vec![0; size.0 * size.1 * pixel_size],
			pixel_size,
			seq_len,
			seq_number: 0,
			size,
		}
	}

	pub fn set_compression(&mut self, compression: Compression) {
		self.compression = compression;
	}

	pub fn set_frame_settings(&mut self, pixel_size: usize, size: (usize, usize)) {
		self.seq_number = self.seq_number.wrapping_add(1);
		self.frame_number = 0;
		self.size = size;
		self.pixel_size = pixel_size;
		self.old_pixels = vec![0; size.0 * size.1 * pixel_size];
	}

	pub fn set_seq_len(&mut self, seq_len: u32) {
		self.seq_len = seq_len;
	}

	/// If the decoder may have already read data from another encoder, then you should ensure to set a different sequence number than the last one received by the decoder from the previous encoder. Otherwise, the decoder may think it's the same sequence, and won't decode the frame (thus possibly the entire sequence) correctly. If you don't know, just use random.
	pub fn set_seq_number(&mut self, seq_number: u32) {
		self.seq_number = seq_number;
	}

	pub fn get_compression(&self) -> &Compression {
		&self.compression
	}

	pub fn get_frame_number(&self) -> u32 {
		self.frame_number
	}

	pub fn get_pixel_size(&self) -> usize {
		self.pixel_size
	}

	pub fn get_seq_len(&self) -> u32 {
		self.seq_len
	}

	pub fn get_seq_number(&self) -> u32 {
		self.seq_number
	}

	pub fn get_size(&self) -> (usize, usize) {
		self.size
	}

	/// Encode raw frame `pixels` into `buf`
	pub fn encode<W: Write>(&mut self, pixels: &mut [u8], mut buf: W) -> Result<(), EncodeError> {
		assert_eq!(pixels.len(), self.old_pixels.len());

		if self.frame_number >= self.seq_len {
			self.frame_number = 0;
			self.seq_number = self.seq_number.wrapping_add(1);
			self.old_pixels.fill(0);
		}

		#[cfg(feature = "simd")]
		{
			if self.pixel_size % 4 == 0 {
				filter::filter_simd_4(self.size, self.pixel_size, pixels, &mut self.old_pixels);
			} else if self.pixel_size % 3 == 0 {
				filter::filter_simd_3(self.size, self.pixel_size, pixels, &mut self.old_pixels);
			} else {
				filter::filter(self.size, self.pixel_size, pixels, &mut self.old_pixels);
			}
		}
		#[cfg(not(feature = "simd"))]
		filter::filter(self.size, self.pixel_size, pixels, &mut self.old_pixels);

		buf.write(&self.seq_number.to_be_bytes())
			.map_err(EncodeError::WriteError)?;
		buf.write(&self.frame_number.to_be_bytes())
			.map_err(EncodeError::WriteError)?;
		buf.write(&(self.size.0 as u32).to_be_bytes())
			.map_err(EncodeError::WriteError)?;
		buf.write(&(self.size.1 as u32).to_be_bytes())
			.map_err(EncodeError::WriteError)?;
		buf.write(&[self.pixel_size as u8])
			.map_err(EncodeError::WriteError)?;

		match self.compression {
			Compression::Brotli(level) => {
				buf.write(&[1]).map_err(EncodeError::WriteError)?;
				let mut compressor = brotli2::write::BrotliEncoder::new(buf, level);
				compressor
					.write_all(pixels)
					.map_err(EncodeError::CompressError)?;
				compressor.flush().map_err(EncodeError::CompressError)?;
			}
			#[cfg(feature = "deflate")]
			Compression::Deflate(level) => {
				buf.write(&[2]).map_err(EncodeError::WriteError)?;
				let mut compressor = deflate::write::DeflateEncoder::new(buf, level);
				compressor
					.write_all(pixels)
					.map_err(EncodeError::CompressError)?;
				compressor.finish().map_err(EncodeError::CompressError)?;
			}
		}

		self.frame_number = self.frame_number.wrapping_add(1);

		Ok(())
	}
}

#[derive(Debug)]
pub enum DecodeError {
	BadDataLength,
	BadHeader,
	DecompressError(std::io::Error),
	UnknownCompression,
}

/// Describes the trust we can have in the correctness of a decoded frame
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DecodedFrameStatus {
	/// Some frames must have been lost. The frame may be incorrect, but the error may be little, especially if the image is static.
	MissedFrame,
	/// The frame seems correct
	Ok,
	/// The encoder has probably done something against the protocol. The frame is probably incorrect.
	Wtf,
}

#[derive(Clone, Debug)]
pub struct DecodedFrame {
	pub frame_number: u32,
	pub pixel_size: usize,
	pub pixels: Vec<u8>,
	pub seq_number: u32,
	pub size: (usize, usize),
	pub status: DecodedFrameStatus,
}

#[derive(Clone)]
pub struct Decoder {
	frame_number: u32,
	old_pixels: Vec<u8>,
	seq_number: u32,
}

impl Decoder {
	pub fn new() -> Self {
		Self {
			frame_number: 0,
			old_pixels: Vec::new(),
			seq_number: u32::MAX,
		}
	}

	pub fn decode(&mut self, data: &[u8]) -> Result<DecodedFrame, DecodeError> {
		if data.len() < 18 {
			return Err(DecodeError::BadHeader);
		}

		let seq_number = u32::from_be_bytes(data[0..4].try_into().unwrap()) as u32;
		let frame_number = u32::from_be_bytes(data[4..8].try_into().unwrap()) as u32;
		let size = (
			u32::from_be_bytes(data[8..12].try_into().unwrap()) as usize,
			u32::from_be_bytes(data[12..16].try_into().unwrap()) as usize,
		);
		let pixel_size = data[16] as usize;
		let compression_algo = data[17];
		let pixels_compressed = &data[18..];

		let resize = size.0 * size.1 * pixel_size != self.old_pixels.len();
		if resize {
			self.frame_number = 0;
			self.old_pixels = vec![0; size.0 * size.1 * pixel_size];
		} else if self.seq_number != seq_number {
			self.old_pixels.fill(0);
		}

		let mut pixels = Vec::with_capacity(size.0 * size.1 * pixel_size);

		match compression_algo {
			1 => {
				let mut decompressor = brotli2::read::BrotliDecoder::new(pixels_compressed);
				decompressor
					.read_to_end(&mut pixels)
					.map_err(DecodeError::DecompressError)?;
			}
			#[cfg(feature = "inflate")]
			2 => {
				let mut decompressor = inflate::InflateWriter::new(&mut pixels);
				decompressor
					.write(pixels_compressed)
					.map_err(DecodeError::DecompressError)?;
				decompressor
					.finish()
					.map_err(DecodeError::DecompressError)?;
			}
			_ => return Err(DecodeError::UnknownCompression),
		}

		if pixels.len() != size.0 * size.1 * pixel_size {
			return Err(DecodeError::BadDataLength);
		}

		#[cfg(feature = "simd")]
		{
			if pixel_size % 4 == 0 {
				filter::unfilter_simd_4(size, pixel_size, &mut pixels, &mut self.old_pixels);
			} else if pixel_size % 3 == 0 {
				filter::unfilter_simd_3(size, pixel_size, &mut pixels, &mut self.old_pixels);
			} else {
				filter::unfilter(size, pixel_size, &mut pixels, &mut self.old_pixels);
			}
		}
		#[cfg(not(feature = "simd"))]
		filter::unfilter(size, pixel_size, &mut pixels, &mut self.old_pixels);

		let ret = Ok(DecodedFrame {
			seq_number,
			frame_number,
			pixel_size,
			size,
			status: if frame_number == 0 {
				if self.seq_number == seq_number {
					if self.frame_number.wrapping_add(1) == frame_number {
						if resize {
							DecodedFrameStatus::Wtf
						} else {
							DecodedFrameStatus::Ok
						}
					} else {
						DecodedFrameStatus::MissedFrame
					}
				} else {
					DecodedFrameStatus::Ok
				}
			} else if self.seq_number == seq_number {
				if resize {
					DecodedFrameStatus::Wtf
				} else if self.frame_number.wrapping_add(1) == frame_number {
					DecodedFrameStatus::Ok
				} else {
					DecodedFrameStatus::MissedFrame
				}
			} else {
				DecodedFrameStatus::MissedFrame
			},
			pixels: pixels.to_vec(),
		});
		self.seq_number = seq_number;
		self.frame_number = frame_number;
		ret
	}
}

impl Default for Decoder {
	fn default() -> Self {
		Self::new()
	}
}

#[cfg(test)]
mod test {
	use super::*;

	use rand::{
		distributions::{Distribution, Uniform},
		Fill,
	};

	#[test]
	fn test_image_reversibility() {
		let maxpixels = 100;

		let mut pixels = vec![0; maxpixels * maxpixels * 4];
		let mut pixels_save = pixels.clone();
		let size_range = Uniform::new(2, maxpixels);
		let pixel_size_range = Uniform::new(1, 5);
		let mut rng = rand::thread_rng();

		let mut compressions = Vec::new();
		compressions.push(Compression::Brotli(3));
		#[cfg(feature = "inflate")]
		compressions.push(Compression::Deflate(Deflate::Fast));

		for compression in compressions {
			for _ in 0..50 {
				let size = (size_range.sample(&mut rng), size_range.sample(&mut rng));
				let pixel_size = pixel_size_range.sample(&mut rng);

				let mut encoder = Encoder::new(size, pixel_size, compression.clone(), 25);
				let mut decoder = Decoder::new();

				for _ in 0..100 {
					pixels[0..size.0 * size.1 * pixel_size]
						.try_fill(&mut rng)
						.unwrap();
					pixels_save[0..size.0 * size.1 * pixel_size]
						.copy_from_slice(&pixels[0..size.0 * size.1 * pixel_size]);
					let mut data = Vec::new();
					encoder
						.encode(&mut pixels[0..size.0 * size.1 * pixel_size], &mut data)
						.unwrap();
					let dec_frame = decoder.decode(&mut data).unwrap();
					assert_eq!(dec_frame.status, DecodedFrameStatus::Ok);
					assert_eq!(dec_frame.size, size);
					assert_eq!(dec_frame.pixel_size, pixel_size);
					assert_eq!(dec_frame.pixels.len(), size.0 * size.1 * pixel_size);
					assert_eq!(
						&dec_frame.pixels,
						&pixels_save[0..size.0 * size.1 * pixel_size]
					);
				}
			}
		}
	}

	#[test]
	fn test_missed_frames() {
		let mut encoder = Encoder::new((1, 1), 1, Compression::Brotli(3), 2);
		let mut decoder = Decoder::new();
		let mut buf = Vec::new();

		encoder.encode(&mut [0], &mut buf).unwrap();
		assert_eq!(
			dbg!(decoder.decode(&mut buf).unwrap()).status,
			DecodedFrameStatus::Ok
		);
		buf.clear();
		encoder.encode(&mut [0], &mut buf).unwrap();
		assert_eq!(
			dbg!(decoder.decode(&mut buf).unwrap()).status,
			DecodedFrameStatus::Ok
		);
		buf.clear();
		encoder.encode(&mut [0], &mut buf).unwrap();
		assert_eq!(
			dbg!(decoder.decode(&mut buf).unwrap()).status,
			DecodedFrameStatus::Ok
		);
		buf.clear();
		encoder.encode(&mut [0], &mut buf).unwrap();
		assert_eq!(
			dbg!(decoder.decode(&mut buf).unwrap()).status,
			DecodedFrameStatus::Ok
		);
		buf.clear();
		encoder.encode(&mut [0], &mut buf).unwrap();
		buf.clear();
		encoder.encode(&mut [0], &mut buf).unwrap();
		assert_eq!(
			dbg!(decoder.decode(&mut buf).unwrap()).status,
			DecodedFrameStatus::MissedFrame
		);
		buf.clear();
		encoder.encode(&mut [0], &mut buf).unwrap();
		assert_eq!(
			dbg!(decoder.decode(&mut buf).unwrap()).status,
			DecodedFrameStatus::Ok
		);
		buf.clear();
		encoder.encode(&mut [0], &mut buf).unwrap();
		buf.clear();
		encoder.encode(&mut [0], &mut buf).unwrap();
		assert_eq!(
			dbg!(decoder.decode(&mut buf).unwrap()).status,
			DecodedFrameStatus::Ok
		);
	}
}