//! Container for switching between un/encrypted stream 
use super::*;

use std::mem;
use chacha20stream::{
    AsyncSink,
    Key, IV,
};
use tokio::io::AsyncWrite;
use std::{
    pin::Pin,
    task::{Context, Poll},
    io,
    marker::Unpin,
};

bool_type!(pub Encrypted; "Is the value encrypted?");
bool_type!(pub Encryption; "What way are we en/decrypting?" => Encrypt, Decrypt);

/// A wrapper `AsyncWrite` stream that allows switching between encrypted (chacha20stream) and plain stream writing.
///
/// # Polymorphic dispatching
/// While this type implements `AsyncWrite` itself, it is recommended to use the polymorphic mutable reference returned by `as_dyn_unpin_mut()` (or `as_dyn_mut()` for non-`Unpin` values of `S`) for writing if lots of `AsyncWrite` methods will be dispatched on the instance.
/// This will prevent the need to check the discriminant of the enum each time the `DualStream`'s `AsyncWrite` methods are polled.
/// The type implements `AsRef/Mut` for streams `S` that are both `Unpin` and not `Unpin` for convenience.
#[derive(Debug)]
pub enum DualStream<S>
{
    /// If there is a panic while switching modes, the stream is left in this invariant state.
    ///
    /// Stream `S` is dropped if the instance is in this state.
    Poisoned,

    /// Stream `S` is being written to through a chacha20stream `AsyncSink` stream cipher.
    Encrypted(AsyncSink<S>),
    /// Stream `S` is being written to directly.
    Plain(S),
}

// We can use dynamic dispatching to prevent the need to check the enum's discriminant each time.
// We have `AsRef/Mut`s for normal and `Unpin` polymorphic `AsyncWrite`s 
impl<'a, S: AsyncWrite + Unpin> AsRef<dyn AsyncWrite + Unpin + 'a> for DualStream<S>
where S: 'a
{
    fn as_ref(&self) -> &(dyn AsyncWrite + Unpin + 'a)
    {
	self.as_dyn_unpin()
    }
}

impl<'a, S: AsyncWrite + Unpin> AsMut<dyn AsyncWrite + Unpin + 'a> for DualStream<S>
where S: 'a
{
    fn as_mut(&mut self) -> &mut (dyn AsyncWrite + Unpin + 'a)
    {
	self.as_dyn_unpin_mut()
    }
}


impl<'a, S: AsyncWrite> AsRef<dyn AsyncWrite + 'a> for DualStream<S>
where S: 'a
{
    fn as_ref(&self) -> &(dyn AsyncWrite + 'a)
    {
	self.as_dyn()
    }
}

impl<'a, S: AsyncWrite> AsMut<dyn AsyncWrite + 'a> for DualStream<S>
where S: 'a
{
    fn as_mut(&mut self) -> &mut (dyn AsyncWrite + 'a)
    {
	self.as_dyn_mut()
    }
}

impl<'a, S: AsyncWrite + Unpin> DualStream<S>
where S: 'a
{
    /// Convert this stream into an encrypted one.
    ///
    /// # Notes
    /// This method makes sure to `flush()` the encrypted stream before dropping the cipher.
    ///
    /// # Panics
    /// If initialising the cipher fails, this method will panic.
    /// # Poisons
    /// If this method panics, then the inner stream will be dropped and this instance will be set to `Poisoned`. This is (currently) irrecoverable.
    pub async fn to_plain(&mut self) -> io::Result<()>
    {
	if !self.is_encrypted() {
	    // No need to do anything
	    return Ok(());
	}
	use tokio::prelude::*;
	self.flush().await?;
	self.to_plain_now();
	Ok(())
    }
    
    /// Convert this stream into an encrypted one.
    ///
    /// # Notes
    /// This method makes sure to `flush()` the encrypted stream before constructing the cipher.
    ///
    /// # Panics
    /// If initialising the cipher fails, this method will panic.
    /// # Poisons
    /// If this method panics, then the inner stream will be dropped and this instance will be set to `Poisoned`. This is (currently) irrecoverable.
    pub async fn to_crypt(&mut self, enc: Encryption, key: Key, iv: IV) -> io::Result<()>
    {
	// NOTE: We can't skip this like `to_plain()` does by checking if the instance is already at `Encrypted` as the key/IV may differ.
	
	use tokio::prelude::*;
	self.flush().await?;
	self.to_crypt_now(enc, key, iv);
	Ok(())
    }
}

impl<'a, S: AsyncWrite> DualStream<S>
where S: Unpin + 'a
{
    /// As an immutable dynamic object for `Unpin` streams.
    pub fn as_dyn_unpin(&self) -> &(dyn AsyncWrite + Unpin + 'a)
    {
	match self {
	    Self::Plain(p) => p,
	    Self::Encrypted(e) => e,
	    _ => panic!("Poisoned")
	}
    }

    /// As a mutable dynamic object for `Unpin` streams
    pub fn as_dyn_unpin_mut(&mut self) -> &mut (dyn AsyncWrite + Unpin + 'a)
    {
	match self {
	    Self::Plain(p) => p,
	    Self::Encrypted(e) => e,
	    _ => panic!("Poisoned")
	}
    }
}

impl<'a, S: AsyncWrite> DualStream<S>
where S: 'a
{
    /// Create a transparent wrapper
    ///
    /// Identical to constructing the enum variant `Self::Plain` manually.
    #[inline(always)] pub fn plain(stream: S) -> Self
    {
	Self::Plain(stream)
    }

    /// Construct an encrypting wrapper over this stream
    ///
    /// # Panics
    /// If constructing the cipher fails
    #[inline] pub fn encrypt(stream: S, key: Key, iv: IV) -> Self
    {
	Self::Encrypted(AsyncSink::encrypt(stream, key, iv).expect("Initialising cipher failed"))
    }

    /// Construct a decrypting wrapper over this stream
    ///
    /// # Panics
    /// If constructing the cipher fails
    #[inline] pub fn decrypt(stream: S, key: Key, iv: IV) -> Self
    {
	Self::Encrypted(AsyncSink::decrypt(stream, key, iv).expect("Initialising cipher failed"))
    }

    /// Construct an encrypting or decrypting wrapper over this stream
    ///
    /// # Panics
    /// If constructing the cipher fails
    #[inline(always)] pub fn crypt(stream: S, enc: Encryption, key: Key, iv: IV) -> Self
    {
	match enc {
	    Encryption::Encrypt => Self::encrypt(stream, key, iv),
	    Encryption::Decrypt => Self::decrypt(stream, key, iv),
	}
    }
    
    /// Is this stream set to encrypted?
    #[inline] pub fn is_encrypted(&self) -> bool
    {
	if let Self::Encrypted(_) = self {
	    true
	} else {
	    false
	}
    }
    /// Is this stream in an invalid state?
    #[inline(always)] pub fn is_poisoned(&self) -> bool {
	
	if let Self::Poisoned = self {
	    true
	} else {
	    false
	}
    }

    /// Move out of self, and in turn poison this insatance until self is assigned a proper value again.
    #[inline(always)] fn poison(&mut self) -> Self
    {
	mem::replace(self, Self::Poisoned)
    }
    
    /// Immediately convert this stream into an encrypted one.
    ///
    /// # Notes
    /// Make sure to `flush()` the stream **before** calling this or data may be lost.
    ///
    /// # Panics
    /// If initialising the cipher fails, this method will panic.
    /// # Poisons
    /// If this method panics, then the inner stream will be dropped and this instance will be set to `Poisoned`. This is (currently) irrecoverable.
    //TODO: in chacha20stream: Add `try_encrypt()`, `try_decrypt()` which both return `Result<Self, (S, Error)>`, to not lose the inner stream if the cipher fails to initialise.
    #[inline] pub fn to_crypt_now(&mut self, enc: Encryption, key: Key, iv: IV)
    {
	let inner = match self.poison() {
	    Self::Encrypted(enc) => enc.into_inner(),
	    Self::Plain(inner) => inner,
	    _ => panic!("Poisoned"),
	};
	*self = Self::Encrypted(match enc {
	    Encryption::Encrypt => AsyncSink::encrypt(inner, key, iv),
	    Encryption::Decrypt => AsyncSink::decrypt(inner, key, iv),
	}.expect("Initialising cipher failed"));
    }

    /// Immediately convert this stream into a plain one
    ///
    /// # Notes
    /// Make sure to `flush()` the stream **before** calling this or encrypted data may be lost.
    ///
    /// # Panics
    /// If dropping the cipher fails, this method will panic.
    /// # Poisons
    /// If this method panics, then the inner stream will be dropped and this instance will be set to `Poisoned`. This is (currently) irrecoverable.
    #[inline] pub fn to_plain_now(&mut self)
    {
	*self = Self::Plain(match self.poison() {
	    Self::Plain(p) => p,
	    Self::Encrypted(e) => e.into_inner(),
	    _ => panic!("Poisoned"),
	});
    }

    /// A mutable reference to the inner (plain) stream, whether this instance is set to encrypted or not.
    pub fn inner_plain_mut(&mut self) -> &mut S
    {
	match self {
	    Self::Plain(p) => p,
	    Self::Encrypted(e) => e.inner_mut(),
	    _ => panic!("Poisoned")
	}
    }
    /// A reference to the inner (plain) stream, whether this instance is set to encrypted or not.
    pub fn inner_plain(&self) -> &S
    {
	match self {
	    Self::Plain(p) => p,
	    Self::Encrypted(e) => e.inner(),
	    _ => panic!("Poisoned")
	}
    }

    /// Consume into the inner (plain) stream
    #[inline] pub fn into_inner(self) -> S
    {
	match self {
	    Self::Plain(p) => p,
	    Self::Encrypted(e) => e.into_inner(),
	    _ => panic!("Poisoned")
	}
    }
    
    /// As an immutable dynamic object
    pub fn as_dyn(&self) -> &(dyn AsyncWrite + 'a)
    {
	match self {
	    Self::Plain(p) => p,
	    Self::Encrypted(e) => e,
	    _ => panic!("Poisoned")
	}
    }

    /// As a mutable dynamic object 
    pub fn as_dyn_mut(&mut self) -> &mut (dyn AsyncWrite + 'a)
    {
	match self {
	    Self::Plain(p) => p,
	    Self::Encrypted(e) => e,
	    _ => panic!("Poisoned")
	}
    }
}

impl<S: AsyncWrite> AsyncWrite for DualStream<S>
{
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, io::Error>> {
	let obj = unsafe {
	    self.map_unchecked_mut(|this| this.as_dyn_mut())		
	};
	obj.poll_write(cx, buf)
    }
    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
	let obj = unsafe {
	    self.map_unchecked_mut(|this| this.as_dyn_mut())		
	};
	obj.poll_flush(cx)
    }
    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
	let obj = unsafe {
	    self.map_unchecked_mut(|this| this.as_dyn_mut())		
	};
	obj.poll_shutdown(cx)
    }
}

#[cfg(test)]
mod tests
{
    use tokio::prelude::*;
    use chacha20stream::keygen;
    #[tokio::test]
    /// Write the whole `input` buffer encrypted, switch to plain, then write the first 5 bytes of `input` again.
    async fn wrapper_construct()
    {
	let input = "Hello world!";
	
	let backing = Vec::new();
	let (key, iv) = keygen();

	let written = {
	    let mut wrapper = super::DualStream::Plain(backing);

	    // Encrypting
	    wrapper.to_crypt(super::Encryption::Encrypt, key, iv).await.unwrap();
	    wrapper.write_all(input.as_bytes()).await.unwrap();

	    // Plain
	    wrapper.to_plain().await.unwrap();
	    wrapper.write_all(&input.as_bytes()[..5]).await.unwrap();

	    // Shutdown the stream and consume it.
	    wrapper.flush().await.unwrap();
	    wrapper.shutdown().await.unwrap();

	    wrapper.into_inner()
	};

	eprintln!("Output bytes: {:?}", written);
	eprintln!("Output attempted string: {:?}", String::from_utf8_lossy(&written[..]));
    }
    
    // TODO: Write a test using Tokio's `duplex` to read and write encrypted bytes on 2 tasks, then compare the output to the input afterwards

}