use super::{accumulator::Accumulator, composable_plaintext::ComposablePlaintext, ComposableIndex};
use crate::{
    encryption::{
        text::TokenFilter,
        unique_indexer::{UniqueIndexer, UniqueIndexerOptions},
        EncryptionError, Plaintext,
    },
    zerokms::IndexKey,
};
use hmac::Mac;

#[derive(Debug)]
pub struct ExactIndex {
    token_filters: Vec<TokenFilter>,
}

impl ExactIndex {
    // TODO: Create an options type for this
    pub fn new(token_filters: Vec<TokenFilter>) -> Self {
        Self { token_filters }
    }
}

impl Default for ExactIndex {
    /// Creates a new `ExactIndex` with no token filters.
    fn default() -> Self {
        Self {
            token_filters: Vec::new(),
        }
    }
}

impl ComposableIndex for ExactIndex {
    fn compose_index(
        &self,
        key: &IndexKey,
        plaintext: ComposablePlaintext,
        accumulator: Accumulator,
    ) -> Result<Accumulator, EncryptionError> {
        let indexer = UniqueIndexer::new(UniqueIndexerOptions {
            token_filters: self.token_filters.clone(),
        });
        let plaintext: Plaintext = plaintext.try_into()?;

        match accumulator {
            Accumulator::Term(term) => {
                let mut mac = indexer.create_hmac(key)?;
                indexer.encrypt_into_hmac(&mut mac, &plaintext)?;
                mac.update(term.as_ref());
                Ok(Accumulator::Term(mac.finalize().into_bytes().to_vec()))
            }
            Accumulator::Terms(terms) => terms
                .into_iter()
                .map(|term| {
                    let mut mac = indexer.create_hmac(key)?;
                    indexer.encrypt_into_hmac(&mut mac, &plaintext)?;
                    mac.update(term.as_ref());
                    Ok(mac.finalize().into_bytes().to_vec())
                })
                .collect::<Result<Vec<_>, EncryptionError>>()
                .map(Accumulator::Terms),
        }
    }
}
