use super::const_args::*;
use crate::encryption::{EncryptionError, Plaintext};

#[derive(Debug, Clone, PartialEq)]
pub enum ComposablePlaintext {
    Unit(Plaintext),
    ConstArg2(ConstArg2),
    ConstArg3(ConstArg3),
    ConstArg4(ConstArg4),
}

impl ComposablePlaintext {
    pub fn new(plaintext: impl Into<Plaintext>) -> Self {
        Self::Unit(plaintext.into())
    }

    pub fn as_plaintext(self) -> Option<Plaintext> {
        match self {
            ComposablePlaintext::Unit(inner) => Some(inner),
            _ => None,
        }
    }

    pub fn try_compose(self, plaintext: impl Into<Plaintext>) -> Result<Self, EncryptionError> {
        let plaintext: Plaintext = plaintext.into();
        match self {
            Self::Unit(a) => Ok(Self::ConstArg2(ConstArg2::new(a, plaintext))),
            Self::ConstArg2(inner) => Ok(Self::ConstArg3(ConstArg3(plaintext, inner))),
            Self::ConstArg3(inner) => Ok(Self::ConstArg4(ConstArg4(plaintext, inner))),
            Self::ConstArg4(_) => Err(EncryptionError::TooManyArguments),
        }
    }

    /// Returns the head of the cons list and the tail of the cons list (if it exists)
    pub fn pop(self) -> (Self, Option<Self>) {
        match self {
            Self::Unit(head) => (Self::Unit(head), None),
            Self::ConstArg2(ConstArg2(head, tail)) => (Self::Unit(head), Some(Self::Unit(tail))),
            Self::ConstArg3(ConstArg3(head, tail)) => {
                (Self::Unit(head), Some(Self::ConstArg2(tail)))
            }
            Self::ConstArg4(ConstArg4(head, tail)) => {
                (Self::Unit(head), Some(Self::ConstArg3(tail)))
            }
        }
    }
}

impl<T: Into<Plaintext>> From<T> for ComposablePlaintext {
    fn from(plaintext: T) -> Self {
        Self::Unit(plaintext.into())
    }
}

impl<A, B> TryFrom<(A, B)> for ComposablePlaintext
where
    A: Into<Plaintext>,
    B: Into<Plaintext>,
{
    type Error = EncryptionError;

    fn try_from((a, b): (A, B)) -> Result<Self, Self::Error> {
        Self::Unit(a.into()).try_compose(b)
    }
}

impl<A, B, C> TryFrom<(A, B, C)> for ComposablePlaintext
where
    A: Into<Plaintext>,
    B: Into<Plaintext>,
    C: Into<Plaintext>,
{
    type Error = EncryptionError;

    fn try_from((a, b, c): (A, B, C)) -> Result<Self, Self::Error> {
        Self::Unit(a.into()).try_compose(b)?.try_compose(c)
    }
}

impl<A, B, C, D> TryFrom<(A, B, C, D)> for ComposablePlaintext
where
    A: Into<Plaintext>,
    B: Into<Plaintext>,
    C: Into<Plaintext>,
    D: Into<Plaintext>,
{
    type Error = EncryptionError;

    fn try_from((a, b, c, d): (A, B, C, D)) -> Result<Self, Self::Error> {
        Self::Unit(a.into())
            .try_compose(b)?
            .try_compose(c)?
            .try_compose(d)
    }
}

impl TryFrom<ComposablePlaintext> for Plaintext {
    type Error = EncryptionError;

    fn try_from(value: ComposablePlaintext) -> Result<Self, Self::Error> {
        value
            .as_plaintext()
            .ok_or_else(|| EncryptionError::TooManyArguments)
    }
}

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

    #[test]
    fn test_pops_in_reverse_insert_order() -> Result<(), Box<dyn std::error::Error>> {
        let output = ComposablePlaintext::new("First")
            .try_compose("Second")?
            .try_compose("Third")?;

        let (head, output) = output.pop();
        assert_eq!(head, ComposablePlaintext::new("Third"));

        let (head, output) = output.unwrap().pop();
        assert_eq!(head, ComposablePlaintext::new("Second"));

        let (head, _) = output.unwrap().pop();
        assert_eq!(head, ComposablePlaintext::new("First"));

        Ok(())
    }

    #[test]
    fn test_single_plaintext_convers() -> Result<(), Box<dyn std::error::Error>> {
        let output = ComposablePlaintext::new("First");
        let plaintext = output.as_plaintext().expect("expected plaintext");

        assert_eq!(plaintext, Plaintext::from("First"));

        Ok(())
    }
}
