//! Serialisation helpers & traits
use super::*;
use tokio::io::{AsyncWrite,AsyncBufRead,AsyncRead,};
use std::{
    marker::Unpin,
    fmt,
};
use crypto::{
    password::{
	Salt,
	Password,
    },
    aes::{self, AesKey,},
};
use once_cell::sync::OnceCell;

/// Used to lazily generate a salt
pub struct SaltGen
{
    salt: OnceCell<Salt>,
}

impl SaltGen
{
    /// Create a new salt generator
    pub fn new() -> Self
    {
	Self {salt:OnceCell::new()}
    }
    /// Get or generate the salt
    #[inline] pub fn get(&self) -> &Salt
    {	
	self.salt.get_or_init(|| Salt::random().expect("Failed to generate random salt")) //is panicking here desired? I think so 
    }

    /// Is this instancce initialised?
    #[inline] pub fn is_set(&self) -> bool
    {
	self.salt.get().is_some()
    }

    /// Get value if it is inisialised
    #[inline] pub fn try_get(&self) -> Option<&Salt>
    {
	self.salt.get()
    }

    /// Set the salt to a specific value. Will panic if already set.
    pub fn set(&self, with: Salt) -> &Salt
    {
	self.salt.set(with).unwrap_or_else(|e| panic!("Tried to set `SaltGen` salt multiple times. New value was {:?}", e));
	unsafe {
	    self.salt.get_unchecked() // is initialised at this point
	}
    }

    /// Consume into an instance of `Salt` if it has been initialised yet
    #[inline] pub fn into_inner(self) -> Option<Salt>
    {
	self.salt.into_inner()
    }
}

impl AsRef<Salt> for SaltGen
{
    #[inline] fn as_ref(&self) -> &Salt
    {
	self.get()
    }
}

impl std::ops::Deref for SaltGen
{
    type Target = Salt;
    #[inline] fn deref(&self) -> &Self::Target {
	self.get()
    }
}

///// The callback to get a derrived password 
//pub trait PasswordFn: for<'r> FnOnce(&'r Salt) -> Option<Password>{}
//impl<T: for<'r> FnOnce(&'r Salt) -> Option<Password>> PasswordFn for T{}
// ^ -- super dumb awful unhelpful error messages trying to do this, so fuck it copy & paste it is

pub trait TextSerialiseable: Sized
{
    fn serialise_text<'a, 'b: 'a, T: AsyncWrite+Unpin+?Sized, F: for<'r> FnOnce(&'r SaltGen) -> Option<Password> + 'b>(&'a self, out: &'b mut T, passwd: F) -> LocalBoxFuture<'a, Result<usize, eyre::Report>>;
    fn deserialise_text<'a, T: AsyncBufRead+Unpin+?Sized, F: for<'r> FnOnce(&'r Salt) -> Option<Password> + 'a>(input: &'a mut T, passwd: F) -> LocalBoxFuture<'a, Result<Self, eyre::Report>>
    where Self: 'a;
    #[inline] fn size_hint(&self) -> (usize, Option<usize>)
    {
	(0,None)
    }
}


pub trait BinarySerialisable: Sized
{
    fn serialise_bytes<'a, 'b, T: AsyncWrite+Unpin+?Sized,F: for<'r> FnOnce(&'r SaltGen) -> Option<Password> + 'b>(&'a self, out: &'b mut T, passwd: F) -> LocalBoxFuture<'a, Result<usize, eyre::Report>>
    where 'b: 'a;
    fn deserialise_bytes<'b , T: AsyncRead+Unpin+?Sized, F: for<'r> FnOnce(&'r Salt) -> Option<Password> + 'b>(input: &'b mut T, passwd: F) -> LocalBoxFuture<'b, Result<Self, eyre::Report>>
    where Self: 'b;
    
    #[inline] fn size_hint(&self) -> (usize, Option<usize>)
    {
	(0,None)
    }
}

pub trait Serialisable: BinarySerialisable + TextSerialiseable
{
    /// Attempt to immediately serialise this object into memory.
    ///
    /// Returns a newly allocated `Vec<u8>` containing the bytes.
    /// Attempts to use the appropriate `size_hint` to allocate vector capacity if possible.
    ///
    /// # Notes
    /// This function will fail if any `await`ed future within the `serialise_*` method cannot complete immediately.
    #[instrument(skip(self, passwd), fields(ty = ?std::any::type_name::<Self>()))]
    fn into_memory<F: FnOnce(&SaltGen) -> Option<Password>>(&self, mode: Mode, passwd: F) -> Result<Vec<u8>, eyre::Report>
    {
	let mut output = match mode.size_hint(self)
	{
	    (0, None) | (_, Some(0)) => Vec::new(),
	    (_, Some(x)) | (x, None)  => Vec::with_capacity(x),
	};
	trace!("Cap for membuf is {} bytes", output.capacity());
	let note = || format!("While serialising {} into memory", std::any::type_name::<Self>());
	let written = match mode {
	    Mode::Text => self.serialise_text(&mut output, passwd).now_or_never(),
	    Mode::Binary => self.serialise_bytes(&mut output, passwd).now_or_never(),
	}
	.ok_or_else(|| eyre::eyre!("Failed to complete {} write syncronously. A future returned `Poll::Pending`.", mode))
	    .with_note(note.clone())?
	    .wrap_err_with(|| eyre::eyre!("Failed to write {} into memory", mode))
	    .with_note(note)?;
	if written != output.len()
	{
	    error!("Function reported {} bytes written, but actually wrote {}", written, output.len());
	}
	Ok(output)
    }

    /// Attempt to immediately deserialise this object from memory
    ///
    /// Returns the object and the number of bytes read from `buf`.
    ///
    /// # Notes
    /// This function will fail if any `await`ed future within the `deserialise_*` method cannot complete immediately.
    #[instrument(skip(buf, passwd),fields(ty = ?std::any::type_name::<Self>(), buf_len = ?buf.as_ref().len()))]
    fn from_memory<T: AsRef<[u8]>, F: FnOnce(&Salt) -> Option<Password>>(buf: T, mode: Mode, passwd: F) -> Result<(Self, usize), eyre::Report>
    {
	let buf = buf.as_ref();
	let mut reader = &buf[..];

	trace!("Attempting read of up to {} bytes in {} mode to instantiate type {}", buf.len(), mode, std::any::type_name::<Self>());
	let note = || format!("While deserialising {} from memory", std::any::type_name::<Self>());
	let bytes = || format!("{}", buf.fmt_view()).header("Binary buffer was");
	let text = ||  String::from_utf8_lossy(buf).into_owned().header("Text buffer was");
	let value = match mode {
	    Mode::Binary => Self::deserialise_bytes(&mut reader, passwd).now_or_never(),
	    Mode::Text => Self::deserialise_text(&mut reader, passwd).now_or_never(),
	}
	.ok_or_else(|| eyre::eyre!("Failed to complete {} read syncronously. A future returned `Poll::Pending`.", mode))
	    .with_note(note.clone())?
	    .wrap_err_with(|| eyre::eyre!("Failed to read {} from memory", mode))
	    .with_note(note)
	    .with_section(bytes)
	    .with_section(text)?;

	let diff = (&reader[0] as *const u8 as usize) - (&buf[0] as *const u8 as usize);

	Ok((value, diff))
    }
}

impl<T> Serialisable for T
where T: BinarySerialisable + TextSerialiseable{}

/// Serialisation mode for `Serialisable` types
#[derive(Debug, Clone, Copy)]
pub enum Mode
{
    /// Serialise as text
    Text,
    /// Serialise as binary
    Binary,
}

impl fmt::Display for Mode
{
    #[inline] fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
    {
	match self {
	    Self::Text => write!(f, "text"),
	    Self::Binary => write!(f, "binary"),
	}
    }
}


impl Mode
{
    #[inline(always)] fn size_hint<T: Serialisable>(&self, other: &T) -> (usize, Option<usize>)
    {
	match self {
	    Self::Text => TextSerialiseable::size_hint(other),
	    Self::Binary => BinarySerialisable::size_hint(other),
	}
    }
}

/// Serialise this object into base64 text
///
/// # TODO
/// Line wrapping
#[instrument(err, skip(value), fields(value_type=?std::any::type_name::<T>()))]
pub fn into_text<T: Serialize>(value: &T) -> Result<String, Error>
{
    let bytes = serde_cbor::to_vec(value).map_err(Error::Serialise)?;
    Ok(base64::encode(bytes))
}


/// Serialise this object into encrypted base64 text
///
/// # TODO
/// Line wrapping
#[instrument(err, skip(value), fields(value_type=?std::any::type_name::<T>()))]
pub async fn into_text_with_key_async<T: Serialize>(value: &T, k: &AesKey) -> Result<String, Error>
{
    let bytes = serde_cbor::to_vec(value).map_err(Error::Serialise)?;
    let mut output = Vec::with_capacity(bytes.len());
    trace!("Encrypting stream");
    let b = aes::encrypt_stream(k, &mut &bytes[..], &mut output).await.map_err(Error::Encryption)?;
    trace!("Encrypted {} bytes (from {} -> {})", b, bytes.len(), output.len());
    Ok(base64::encode(output))
}

#[derive(Debug)]
pub enum Error
{
    Base64(base64::DecodeError),
    Serialise(serde_cbor::Error),
    Deserialise(serde_cbor::Error),
    Encryption(aes::Error),
    Decryption(aes::Error),
}

impl std::error::Error for Error
{
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
	match &self {
	    Self::Base64(b) => Some(b),
	    Self::Serialise(b) => Some(b),
	    Self::Deserialise(b) => Some(b),
	    Self::Encryption(b) => Some(b),
	    Self::Decryption(b) => Some(b),
	}
    }
}

impl std::fmt::Display for Error
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result
    {
	match self {
	    Self::Base64(_) => write!(f, "base64 decode failed"),
	    Self::Serialise(_) => write!(f, "serialisation failed"),
	    Self::Deserialise(_) => write!(f, "deserialisation failed"),
	    Self::Encryption(_) => write!(f, "encryption failed"),
	    Self::Decryption(_) => write!(f, "decryption failed"),
	}
    }
}

/// Try to deserialise this object from text.
#[instrument(err, skip(string), fields(string = ?string.as_ref()))]
pub fn from_text<T: serde::de::DeserializeOwned>(string: impl AsRef<[u8]>) -> Result<T, Error>
{
    let bytes = base64::decode(string).map_err(Error::Base64)?;
    serde_cbor::from_reader(&bytes[..]).map_err(Error::Deserialise)
}


/// Try to deserialise this encrypted object from text.
#[instrument(err, skip(string), fields(string = ?string.as_ref()))]
pub async fn from_text_with_key_async<T: serde::de::DeserializeOwned>(string: impl AsRef<[u8]>, k: &AesKey) -> Result<T, Error>
{
    let bytes = base64::decode(string).map_err(Error::Base64)?;
    let mut output = Vec::with_capacity(bytes.len());
    trace!("Decrypting stream");
    let b = aes::decrypt_stream(k, &mut &bytes[..], &mut output).await.map_err(Error::Decryption)?;
    trace!("Decrypted {} bytes (from {} -> {})", b, bytes.len(), output.len());
    serde_cbor::from_reader(&output[..]).map_err(Error::Deserialise)
}

#[cfg(test)]
mod tests
{
    use crate::*;
    #[test]
    fn test_ser() -> Result<(), eyre::Report>
    {
	let object = "hello world";
	let enc = super::into_text(&String::from(object)).wrap_err(eyre::eyre!("failed to encode"))?;
	println!("Enc: {:?}", enc);
	let dec: String = super::from_text(enc).wrap_err(eyre::eyre!("failed to decode"))?;
	assert_eq!(object,dec);
	Ok(())
    }
}