// Copyright 2019 Parity Technologies
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::{Encode, Error, Decode, Compact, CompactLen};

use core::{iter::ExactSizeIterator, mem};
use crate::alloc::vec::Vec;

/// Trait that allows to append items to an encoded representation without
/// decoding all previous added items.
pub trait EncodeAppend {
	/// The item that will be appended.
	type Item: Encode;

	/// Append all items in `iter` to the given `self_encoded` representation.
	///
	/// # Example
	///
	/// ```
	///# use tetsy_scale_codec::EncodeAppend;
	///
	/// // Some encoded data
	/// let data = Vec::new();
	///
	/// let item = 8u32;
	/// let encoded = <Vec<u32> as EncodeAppend>::append(data, std::iter::once(&item)).expect("Adds new element");
	///
	/// // Add multiple element
	/// <Vec<u32> as EncodeAppend>::append(encoded, &[700u32, 800u32, 10u32]).expect("Adds new elements");
	/// ```
	fn append<'a, I: IntoIterator<Item=&'a Self::Item>>(
		self_encoded: Vec<u8>,
		iter: I,
	) -> Result<Vec<u8>, Error> where Self::Item: 'a, I::IntoIter: ExactSizeIterator;
}

fn extract_length_data(data: &[u8], input_len: usize) -> Result<(u32, usize, usize), Error> {
	let len = u32::from(Compact::<u32>::decode(&mut &data[..])?);
	let new_len = len
		.checked_add(input_len as u32)
		.ok_or_else(|| "New vec length greater than `u32::max_value()`.")?;

	let encoded_len = Compact::<u32>::compact_len(&len);
	let encoded_new_len = Compact::<u32>::compact_len(&new_len);

	Ok((new_len, encoded_len, encoded_new_len))
}

impl<T: Encode + 'static> EncodeAppend for Vec<T> {
	type Item = T;

	fn append<'a, I: IntoIterator<Item=&'a Self::Item>>(
		mut self_encoded: Vec<u8>,
		iter: I,
	) -> Result<Vec<u8>, Error> where Self::Item: 'a, I::IntoIter: ExactSizeIterator {
		let iter = iter.into_iter();
		let input_len = iter.len();

		// No data present, just encode the given input data.
		if self_encoded.is_empty() {
			crate::codec::compact_encode_len_to(&mut self_encoded, iter.len())?;
			iter.for_each(|e| e.encode_to(&mut self_encoded));
			return Ok(self_encoded);
		}

		let (new_len, encoded_len, encoded_new_len) = extract_length_data(&self_encoded, input_len)?;

		let replace_len = |dest: &mut Vec<u8>| {
			Compact(new_len).using_encoded(|e| {
				dest[..encoded_new_len].copy_from_slice(e);
			})
		};

		let append_new_elems = |dest: &mut Vec<u8>| iter.for_each(|a| a.encode_to(dest));

		// If old and new encoded len is equal, we don't need to copy the
		// already encoded data.
		if encoded_len == encoded_new_len {
			replace_len(&mut self_encoded);
			append_new_elems(&mut self_encoded);

			Ok(self_encoded)
		} else {
			let size = encoded_new_len + self_encoded.len() - encoded_len;

			let mut res = Vec::with_capacity(size + input_len * mem::size_of::<T>());
			unsafe { res.set_len(size); }

			// Insert the new encoded len, copy the already encoded data and
			// add the new element.
			replace_len(&mut res);
			res[encoded_new_len..size].copy_from_slice(&self_encoded[encoded_len..]);
			append_new_elems(&mut res);

			Ok(res)
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::{Input, EncodeLike};

	#[test]
	fn vec_encode_append_works() {
		let max_value = 1_000_000;

		let encoded = (0..max_value).fold(Vec::new(), |encoded, v| {
			<Vec<u32> as EncodeAppend>::append(encoded, std::iter::once(&v)).unwrap()
		});

		let decoded = Vec::<u32>::decode(&mut &encoded[..]).unwrap();
		assert_eq!(decoded, (0..max_value).collect::<Vec<_>>());
	}

	#[test]
	fn vec_encode_append_multiple_items_works() {
		let max_value = 1_000_000u32;

		let encoded = (0..max_value).fold(Vec::new(), |encoded, v| {
			<Vec<u32> as EncodeAppend>::append(encoded, &[v, v, v, v]).unwrap()
		});

		let decoded = Vec::<u32>::decode(&mut &encoded[..]).unwrap();
		let expected = (0..max_value).fold(Vec::new(), |mut vec, i| {
			vec.append(&mut vec![i, i, i, i]);
			vec
		});
		assert_eq!(decoded, expected);
	}

	#[test]
	fn append_non_copyable() {
		#[derive(Eq, PartialEq, Debug)]
		struct NoCopy { data: u32 };

		impl EncodeLike for NoCopy {}

		impl Encode for NoCopy {
			fn encode(&self) -> Vec<u8> {
				self.data.encode()
			}
		}

		impl Decode for NoCopy {
			fn decode<I: Input>(input: &mut I) -> Result<Self, Error> {
				u32::decode(input).map(|data| Self { data })
			}
		}

		let append = NoCopy { data: 100 };
		let data = Vec::new();
		let encoded = <Vec<NoCopy> as EncodeAppend>::append(data, std::iter::once(&append)).unwrap();

		let decoded = <Vec<NoCopy>>::decode(&mut &encoded[..]).unwrap();
		assert_eq!(vec![append], decoded);
	}
}