btt 0.5.0

use core::{array, marker::PhantomData};

use crate::{
	DecodeIntoExactError, Decoder, DynDecoder, DynEncoder, EncodeIntoExactError, Encoder,
	alphabet::{Alphabet, DecodingTable, DynAlphabet, DynDecodingTable},
	generic::{GenericDecoder, GenericEncoder},
	padding::Padding
};

#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub(crate) struct ChunkedEncoder<A, D, P, Imp, const LEN: usize, const I: usize, const O: usize>
where
	Imp: ChunkedEncoderImpl<I, O>
{
	encoder: GenericEncoder<A, D, P, LEN>,
	imp: PhantomData<Imp>
}

impl<A, D, P, Imp, const LEN: usize, const I: usize, const O: usize> AsRef<Alphabet<LEN>>
	for ChunkedEncoder<A, D, P, Imp, LEN, I, O>
where
	A: AsRef<Alphabet<LEN>>,
	Imp: ChunkedEncoderImpl<I, O>
{
	fn as_ref(&self) -> &Alphabet<LEN> {
		self.encoder.as_ref()
	}
}

impl<A, D, P, Imp, const LEN: usize, const I: usize, const O: usize> AsRef<DecodingTable<LEN>>
	for ChunkedEncoder<A, D, P, Imp, LEN, I, O>
where
	D: AsRef<DecodingTable<LEN>>,
	Imp: ChunkedEncoderImpl<I, O>
{
	fn as_ref(&self) -> &DecodingTable<LEN> {
		self.encoder.as_ref()
	}
}

impl<A, D, P, Imp, const LEN: usize, const I: usize, const O: usize>
	From<GenericEncoder<A, D, P, LEN>> for ChunkedEncoder<A, D, P, Imp, LEN, I, O>
where
	Imp: ChunkedEncoderImpl<I, O>
{
	fn from(encoder: GenericEncoder<A, D, P, LEN>) -> Self {
		Self { encoder, imp: PhantomData }
	}
}

impl<A, D, P, Imp, const LEN: usize, const I: usize, const O: usize> Encoder<LEN>
	for ChunkedEncoder<A, D, P, Imp, LEN, I, O>
where
	A: AsRef<Alphabet<LEN>> + Copy,
	D: AsRef<DecodingTable<LEN>> + Copy,
	P: Padding,
	Imp: ChunkedEncoderImpl<I, O>,
	Self: Copy // TODO
{
	fn alphabet(&self) -> &Alphabet<LEN> {
		self.as_ref()
	}
}

impl<A, D, P, Imp, const LEN: usize, const I: usize, const O: usize> DynEncoder
	for ChunkedEncoder<A, D, P, Imp, LEN, I, O>
where
	A: AsRef<Alphabet<LEN>> + Copy,
	D: AsRef<DecodingTable<LEN>> + Copy,
	P: Padding,
	Imp: ChunkedEncoderImpl<I, O>,
	Self: Copy // TODO
{
	fn dyn_alphabet(&self) -> &DynAlphabet {
		self.encoder.alphabet().as_ref().as_dyn_alphabet()
	}

	fn encoded_len(&self, len: usize) -> Option<usize> {
		Imp::encoded_len(len, self.encoder.padding().padding().is_some())
	}

	fn encode_into_exact<'a>(
		&self,
		src: &[u8],
		dst: &'a mut [u8]
	) -> Result<&'a str, EncodeIntoExactError> {
		if self.encoded_len(src.len()).ok_or(EncodeIntoExactError::LengthOverflow)? != dst.len() {
			return Err(EncodeIntoExactError::LengthMismatch);
		}

		let alphabet: &Alphabet<LEN> = self.encoder.as_ref();
		let padding = self.encoder.padding().padding();

		let mut src_chunks = src.chunks_exact(I);
		let mut dst_chunks = dst.chunks_mut(O);

		debug_assert_eq!(
			dst_chunks.len(),
			src_chunks.len() + usize::from(!src_chunks.remainder().is_empty())
		);

		src_chunks.by_ref().zip(&mut dst_chunks).for_each(|(src, dst)| {
			Imp::encode_chunk_raw(src.try_into().unwrap(), dst.try_into().unwrap());
			dst.iter_mut().for_each(|dst| *dst = alphabet[*dst as usize]);
		});

		debug_assert!(src_chunks.next().is_none());

		let src = src_chunks.remainder();
		let padding_len = if !src.is_empty() {
			let dst = dst_chunks.next().unwrap();

			let (dst, to_pad) = dst.split_at_mut(if padding.is_some() {
				Imp::padding_len(src.len()) // TODO
			} else {
				dst.len()
			});

			let mut tmp = [0; O];
			Imp::encode_last_chunk_raw(src, &mut tmp);
			dst.copy_from_slice(&tmp[..dst.len()]);
			dst.iter_mut().for_each(|dst| *dst = alphabet[*dst as usize]);

			if let Some(padding) = padding {
				to_pad.fill(padding);
			}

			to_pad.len()
		} else {
			0
		};

		debug_assert!(dst_chunks.next().is_none());
		let mut iter = dst.iter();
		debug_assert!(iter.by_ref().rev().take(padding_len).all(|&b| b == padding.unwrap()));
		// debug_assert!(iter.all(|b| alphabet.contains(b)));

		Ok(unsafe { str::from_utf8_unchecked(dst) })
	}
}

impl<A, D, P, Imp, const LEN: usize, const I: usize, const O: usize> Decoder<LEN>
	for ChunkedEncoder<A, D, P, Imp, LEN, I, O>
where
	A: AsRef<Alphabet<LEN>> + Copy,
	D: AsRef<DecodingTable<LEN>> + Copy,
	P: Padding,
	Imp: ChunkedEncoderImpl<I, O>,
	Self: Copy // TODO
{
	fn decoding_table(&self) -> &DecodingTable<LEN> {
		self.as_ref()
	}
}

impl<A, D, P, Imp, const LEN: usize, const I: usize, const O: usize> DynDecoder
	for ChunkedEncoder<A, D, P, Imp, LEN, I, O>
where
	A: AsRef<Alphabet<LEN>> + Copy,
	D: AsRef<DecodingTable<LEN>> + Copy,
	P: Padding,
	Imp: ChunkedEncoderImpl<I, O>,
	Self: Copy // TODO
{
	fn dyn_decoding_table(&self) -> &DynDecodingTable {
		self.encoder.decoding_table().as_ref().as_dyn_decoding_table()
	}

	fn decoded_len(&self, src: &[u8]) -> Option<usize> {
		Imp::decoded_len(src, self.encoder.padding().padding())
	}

	fn decode_into_exact<'a>(
		&self,
		src: &[u8],
		dst: &'a mut [u8]
	) -> Result<&'a [u8], DecodeIntoExactError> {
		ChunkedDecoder::from(*self).decode_into_exact(src, dst)
	}
}

#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
pub(crate) struct ChunkedDecoder<D, P, Imp, const LEN: usize, const I: usize, const O: usize>
where
	Imp: ChunkedEncoderImpl<I, O>
{
	decoder: GenericDecoder<D, P, LEN>,
	imp: PhantomData<Imp>
}

impl<A, D, P, Imp, const LEN: usize, const I: usize, const O: usize>
	From<ChunkedEncoder<A, D, P, Imp, LEN, I, O>> for ChunkedDecoder<D, P, Imp, LEN, I, O>
where
	Imp: ChunkedEncoderImpl<I, O>
{
	fn from(encoder: ChunkedEncoder<A, D, P, Imp, LEN, I, O>) -> Self {
		Self { decoder: encoder.encoder.into(), imp: encoder.imp }
	}
}

impl<D, P, Imp, const LEN: usize, const I: usize, const O: usize> From<GenericDecoder<D, P, LEN>>
	for ChunkedDecoder<D, P, Imp, LEN, I, O>
where
	Imp: ChunkedEncoderImpl<I, O>
{
	fn from(decoder: GenericDecoder<D, P, LEN>) -> Self {
		Self { decoder, imp: PhantomData }
	}
}

impl<D, P, Imp, const LEN: usize, const I: usize, const O: usize> Decoder<LEN>
	for ChunkedDecoder<D, P, Imp, LEN, I, O>
where
	D: AsRef<DecodingTable<LEN>> + Copy,
	P: Padding,
	Imp: ChunkedEncoderImpl<I, O>
{
	fn decoding_table(&self) -> &DecodingTable<LEN> {
		self.decoder.as_ref()
	}
}

impl<D, P, Imp, const LEN: usize, const I: usize, const O: usize> DynDecoder
	for ChunkedDecoder<D, P, Imp, LEN, I, O>
where
	D: AsRef<DecodingTable<LEN>> + Copy,
	P: Padding,
	Imp: ChunkedEncoderImpl<I, O>
{
	fn dyn_decoding_table(&self) -> &DynDecodingTable {
		self.decoder.decoding_table().as_ref().as_dyn_decoding_table()
	}

	fn decoded_len(&self, src: &[u8]) -> Option<usize> {
		Imp::decoded_len(src, self.decoder.padding().padding())
	}

	fn decode_into_exact<'a>(
		&self,
		src: &[u8],
		dst: &'a mut [u8]
	) -> Result<&'a [u8], DecodeIntoExactError> {
		if self.decoded_len(src).ok_or(DecodeIntoExactError::InvalidLength)? != dst.len() {
			return Err(DecodeIntoExactError::LengthMismatch);
		}

		let decoding_table: &DecodingTable<LEN> = self.decoder.as_ref();
		let padding = self.decoder.padding().padding();
		let mut src_chunks = src.chunks(O).enumerate().map(|(i, src)| (i * O, src));
		let mut dst_chunks = dst.chunks_mut(I);

		let last_chunk =
			if !src.len().is_multiple_of(O) || src.last().is_some_and(|&b| Some(b) == padding) {
				src_chunks.next_back()
			} else {
				None
			};

		let mut decoded = [0; O];

		src_chunks.by_ref().zip(&mut dst_chunks).try_for_each(|((index, src), dst)| {
			let src: &[u8; O] = src.try_into().unwrap();
			let dst: &mut [u8; I] = dst.try_into().unwrap();

			for i in 0..O {
				decoded[i] = decoding_table
					.decode(src[i])
					.ok_or_else(|| DecodeIntoExactError::InvalidCharacter(index + i))?;
			}

			Imp::decode_raw_chunk(&decoded, dst)
				.map_err(|()| DecodeIntoExactError::InvalidChunk { index, len: O })
		})?;

		debug_assert!(src_chunks.next().is_none());

		if let Some((index, src)) = last_chunk {
			debug_assert!(!src.is_empty());
			let dst = dst_chunks.next().unwrap();

			let padding_len = if let Some(padding) = padding {
				debug_assert_eq!(*src.last().unwrap(), padding);
				src.iter().rev().take_while(|&&b| b == padding).count()
			} else {
				0
			};

			let (encoded, src_padding) = src.split_at(src.len() - padding_len);

			debug_assert!(padding.is_none_or(|p| encoded.last().copied() != Some(p)));
			debug_assert!(src_padding.iter().all(|&b| b == padding.unwrap()));

			for i in 0..encoded.len() {
				decoded[i] = decoding_table
					.decode(encoded[i])
					.ok_or_else(|| DecodeIntoExactError::InvalidCharacter(index + i))?;
			}

			let raw_encoded = &decoded[..encoded.len()];
			let mut decoded = [0; I];

			Imp::decode_raw_last_chunk(raw_encoded, &mut decoded)
				.map_err(|()| DecodeIntoExactError::InvalidChunk { index, len: O })?;

			dst.copy_from_slice(&decoded[..dst.len()]);

			let mut tmp = [0; O];
			Imp::encode_last_chunk_raw(dst, &mut tmp);

			if raw_encoded != &tmp[..raw_encoded.len()] {
				return Err(DecodeIntoExactError::NonCanonical);
			}
		}

		debug_assert!(dst_chunks.next().is_none());
		Ok(dst)
	}
}

pub(crate) trait ChunkedEncoderImpl<const I: usize, const O: usize> {
	fn padding_len(len: usize) -> usize;
	fn encoded_len(len: usize, padded: bool) -> Option<usize>;
	fn decoded_len(src: &[u8], padding: Option<u8>) -> Option<usize>;
	fn encode_chunk_raw(src: &[u8; I], dst: &mut [u8; O]);
	fn decode_raw_chunk(src: &[u8; O], dst: &mut [u8; I]) -> Result<(), ()>;

	fn encode_last_chunk_raw(src: &[u8], dst: &mut [u8; O]) {
		Self::encode_chunk_raw(&array::from_fn(|i| src.get(i).copied().unwrap_or(0)), dst);
	}

	fn decode_raw_last_chunk(src: &[u8], dst: &mut [u8; I]) -> Result<(), ()> {
		Self::decode_raw_chunk(&array::from_fn(|i| src.get(i).copied().unwrap_or(0)), dst)
	}
}