use super::EncryptionError;

pub mod accumulator;
pub mod composable_plaintext;
pub mod const_args;
pub mod exact_index;
pub mod prefix_indexer;

use crate::zerokms::IndexKey;
pub use accumulator::Accumulator;
pub use composable_plaintext::ComposablePlaintext;
pub use exact_index::ExactIndex;
pub use prefix_indexer::PrefixIndex;

/// Trait to represent any index that can be composed with other indexes
/// in a Structured Encryption scheme.
pub trait ComposableIndex: std::fmt::Debug {
    fn compose_index(
        &self,
        key: &IndexKey,
        plaintext: ComposablePlaintext,
        accumulator: Accumulator,
    ) -> Result<Accumulator, EncryptionError>;

    fn compose_query(
        &self,
        key: &IndexKey,
        plaintext: ComposablePlaintext,
        accumulator: Accumulator,
    ) -> Result<Accumulator, EncryptionError> {
        self.compose_index(key, plaintext, accumulator)
    }
}

impl ComposableIndex for Box<dyn ComposableIndex + Send> {
    fn compose_index(
        &self,
        key: &IndexKey,
        plaintext: ComposablePlaintext,
        accumulator: Accumulator,
    ) -> Result<Accumulator, EncryptionError> {
        (**self).compose_index(key, plaintext, accumulator)
    }

    fn compose_query(
        &self,
        key: &IndexKey,
        plaintext: ComposablePlaintext,
        accumulator: Accumulator,
    ) -> Result<Accumulator, EncryptionError> {
        (**self).compose_query(key, plaintext, accumulator)
    }
}

#[derive(Debug)]
pub struct CompoundIndexOfTwo<A, B> {
    indexes: (A, B),
}

impl<A, B> CompoundIndexOfTwo<A, B> {
    pub fn new(a: A, b: B) -> Self {
        Self { indexes: (a, b) }
    }
}

impl<A: ComposableIndex, B: ComposableIndex> ComposableIndex for CompoundIndexOfTwo<A, B> {
    fn compose_index(
        &self,
        key: &IndexKey,
        inputs: ComposablePlaintext,
        accumulator: Accumulator,
    ) -> Result<Accumulator, EncryptionError> {
        let (head, tail) = inputs.pop();
        self.indexes.0.compose_index(
            key,
            head,
            self.indexes.1.compose_index(
                key,
                tail.ok_or(EncryptionError::TooFewArguments)?,
                accumulator,
            )?,
        )
    }

    fn compose_query(
        &self,
        key: &IndexKey,
        inputs: ComposablePlaintext,
        accumulator: Accumulator,
    ) -> Result<Accumulator, EncryptionError> {
        let (head, tail) = inputs.pop();
        self.indexes.0.compose_query(
            key,
            head,
            self.indexes.1.compose_query(
                key,
                tail.ok_or(EncryptionError::TooFewArguments)?,
                accumulator,
            )?,
        )
    }
}

#[derive(Debug)]
pub struct CompoundIndex<I: ComposableIndex>(I);

impl<I: ComposableIndex + Send> CompoundIndex<I> {
    pub fn new(index: I) -> Self {
        Self(index)
    }

    pub fn and<J: ComposableIndex + Send>(
        self,
        other: J,
    ) -> CompoundIndex<CompoundIndexOfTwo<J, I>> {
        CompoundIndex(CompoundIndexOfTwo::new(other, self.0))
    }
}

impl<I: ComposableIndex + Send> ComposableIndex for CompoundIndex<I> {
    fn compose_index(
        &self,
        key: &IndexKey,
        plaintext: ComposablePlaintext,
        accumulator: Accumulator,
    ) -> Result<Accumulator, EncryptionError> {
        self.0.compose_index(key, plaintext, accumulator)
    }

    fn compose_query(
        &self,
        key: &IndexKey,
        plaintext: ComposablePlaintext,
        accumulator: Accumulator,
    ) -> Result<Accumulator, EncryptionError> {
        self.0.compose_query(key, plaintext, accumulator)
    }
}

#[cfg(test)]
mod tests {
    use crate::encryption::Plaintext;

    use super::*;

    #[test]
    fn test_single_exact() -> Result<(), Box<dyn std::error::Error>> {
        let index_key = IndexKey::from([1; 32]);
        let index = ExactIndex::new(vec![]);

        let result = index.compose_index(
            &index_key,
            "foo@bar.com".into(),
            Accumulator::from_salt("user#email"),
        )?;

        assert_eq!(result.terms().len(), 1);

        Ok(())
    }

    #[test]
    fn test_single_exact_in_a_compound() -> Result<(), Box<dyn std::error::Error>> {
        let index_key = IndexKey::from([1; 32]);

        let index = CompoundIndex::new(ExactIndex::new(vec![]));
        let result = index.compose_index(
            &index_key,
            "foo@bar.com".into(),
            Accumulator::from_salt("user#email"),
        )?;

        assert_eq!(result.terms().len(), 1);

        Ok(())
    }

    #[test]
    fn test_two_exact_indexes() -> Result<(), Box<dyn std::error::Error>> {
        let index_key = IndexKey::from([1; 32]);
        let index = CompoundIndex::new(ExactIndex::new(vec![])).and(ExactIndex::new(vec![]));

        let result = index.compose_index(
            &index_key,
            ("foo@bar.com", "Person").try_into()?,
            Accumulator::from_salt("user#email/name"),
        )?;
        assert_eq!(result.terms().len(), 1);

        Ok(())
    }

    #[test]
    fn test_one_exact_one_prefix() -> Result<(), Box<dyn std::error::Error>> {
        let index_key = IndexKey::from([1; 32]);
        let index = CompoundIndex::new(ExactIndex::new(vec![])).and(PrefixIndex::new(vec![]));

        let result = index.compose_index(
            &index_key,
            ("foo@bar.com", "Person").try_into()?,
            Accumulator::from_salt("user#email/name"),
        )?;
        assert_eq!(result.terms().len(), 4);

        Ok(())
    }

    #[test]
    fn test_three_exact_indexes() -> Result<(), Box<dyn std::error::Error>> {
        let index_key = IndexKey::from([1; 32]);
        let index = CompoundIndex::new(ExactIndex::new(vec![]))
            .and(ExactIndex::new(vec![]))
            .and(ExactIndex::new(vec![]));

        let result = index.compose_index(
            &index_key,
            ("foo@bar.com", "Person", true).try_into()?,
            Accumulator::from_salt("user#email/name/active"),
        )?;

        assert_eq!(result.terms().len(), 1);

        Ok(())
    }

    #[test]
    fn test_one_exact_one_prefix_one_exact() -> Result<(), Box<dyn std::error::Error>> {
        let index_key = IndexKey::from([1; 32]);
        let index = CompoundIndex::new(ExactIndex::new(vec![]))
            .and(PrefixIndex::new(vec![]))
            .and(ExactIndex::new(vec![]));

        let result = index.compose_index(
            &index_key,
            ("foo@bar.com", "Person", 100i32).try_into()?,
            Accumulator::from_salt("user#email/name/login-count"),
        )?;
        assert_eq!(result.terms().len(), 4);

        Ok(())
    }

    #[test]
    fn test_one_exact_one_prefix_one_prefix() -> Result<(), Box<dyn std::error::Error>> {
        let index_key = IndexKey::from([1; 32]);
        let index = CompoundIndex::new(ExactIndex::new(vec![]))
            .and(PrefixIndex::new(vec![]))
            .and(PrefixIndex::new(vec![]));

        let result = index.compose_index(
            &index_key,
            ("foo@bar.com", "Daniel", "Draper").try_into()?,
            Accumulator::from_salt("user#email/first-name/last-name"),
        )?;

        let acc = Accumulator::from_salt("user#email/first-name/last-name");

        let acc = ExactIndex::new(vec![]).compose_index(
            &index_key,
            Plaintext::from("foo@bar.com").into(),
            acc,
        )?;

        let acc = PrefixIndex::new(vec![]).compose_index(
            &index_key,
            Plaintext::from("Daniel").into(),
            acc,
        )?;

        let acc = PrefixIndex::new(vec![]).compose_index(
            &index_key,
            Plaintext::from("Draper").into(),
            acc,
        )?;

        let terms = result.terms();
        let acc_terms = acc.terms();

        assert_eq!(terms.len(), 16);
        assert_eq!(acc_terms.len(), 16);

        assert_eq!(terms, acc_terms);

        Ok(())
    }
}
