use super::builder::StorageBuilder;
use super::indexer::{IndexerInit, Indexes, IndexesForQuery, QueryOp};
use super::text::TokenFilter;
use super::QueryBuilder;
use super::{errors::EncryptionError, plaintext::Plaintext, IndexTerm};
use crate::zerokms::IndexKey;
use cipherstash_config::column;
use hmac::{Hmac, Mac};
use sha2::Sha256;

type HmacSha256 = Hmac<Sha256>;

impl IndexerInit for UniqueIndexer {
    type Args = UniqueIndexerOptions;
    type Error = EncryptionError;

    fn try_init<A>(args: A) -> Result<Self, Self::Error>
    where
        Self::Args: TryFrom<A, Error = Self::Error>,
    {
        let args = UniqueIndexerOptions::try_from(args)?;
        Ok(Self::new(args))
    }
}

impl<'k> Indexes<'k, Plaintext> for UniqueIndexer {
    fn index(
        &self,
        mut builder: StorageBuilder<'k, Plaintext>,
    ) -> Result<StorageBuilder<'k, Plaintext>, EncryptionError> {
        let index_term = self.encrypt(builder.plaintext(), builder.index_key())?;
        builder.add_index_term(index_term);

        Ok(builder)
    }
}

impl<C> IndexesForQuery<Plaintext, C> for UniqueIndexer {
    fn query_index(
        &self,
        builder: QueryBuilder<Plaintext, C>,
        _op: QueryOp,
    ) -> Result<IndexTerm, EncryptionError> {
        let index_term = self.encrypt(builder.plaintext(), builder.index_key())?;
        Ok(index_term)
    }
}

pub struct UniqueIndexerOptions {
    pub token_filters: Vec<TokenFilter>,
}

impl TryFrom<&column::IndexType> for UniqueIndexerOptions {
    type Error = EncryptionError;

    fn try_from(value: &column::IndexType) -> Result<Self, Self::Error> {
        match value {
            column::IndexType::Unique { token_filters, .. } => Ok(Self {
                token_filters: token_filters.iter().copied().map(From::from).collect(),
            }),
            _ => Err(EncryptionError::IndexingError(
                "UniqueIndexerOptions can only be created from a Unique Index configuration"
                    .to_string(),
            )),
        }
    }
}

impl Default for UniqueIndexerOptions {
    fn default() -> Self {
        Self {
            token_filters: vec![TokenFilter::Downcase],
        }
    }
}

impl Default for UniqueIndexer {
    fn default() -> Self {
        Self::new(Default::default())
    }
}

// Note that Debug is safe because no key material or sensitive data is exposed
#[derive(Debug)]
pub struct UniqueIndexer {
    token_filters: Vec<TokenFilter>,
}

impl UniqueIndexer {
    pub fn new(opts: UniqueIndexerOptions) -> Self {
        Self {
            token_filters: opts.token_filters,
        }
    }

    pub(super) fn create_hmac(&self, index_key: &IndexKey) -> Result<HmacSha256, EncryptionError> {
        Ok(HmacSha256::new_from_slice(index_key.key())?)
    }

    pub(super) fn encrypt_into_hmac(
        &self,
        mac: &mut HmacSha256,
        plaintext: &Plaintext,
    ) -> Result<(), EncryptionError> {
        let plaintext_bytes = match plaintext {
            Plaintext::Text(None)
            | Plaintext::BigInt(None)
            | Plaintext::Boolean(None)
            | Plaintext::Decimal(None)
            | Plaintext::Float(None)
            | Plaintext::Int(None)
            | Plaintext::NaiveDate(None)
            | Plaintext::SmallInt(None)
            | Plaintext::Timestamp(None) => return Ok(()),

            Plaintext::Text(Some(utf_str)) => {
                let filtered_string = self
                    .token_filters
                    .iter()
                    .try_fold(utf_str.to_string(), |s, filter| filter.process_single(s));
                Plaintext::from(filtered_string).to_vec()
            }

            x => x.to_vec(),
        };

        mac.update(&plaintext_bytes);

        Ok(())
    }

    pub fn encrypt(
        &self,
        plaintext: &Plaintext,
        index_key: &IndexKey,
    ) -> Result<IndexTerm, EncryptionError> {
        if plaintext.is_null() {
            Ok(IndexTerm::Null)
        } else {
            let mut mac = self.create_hmac(index_key)?;
            self.encrypt_into_hmac(&mut mac, plaintext)?;
            Ok(IndexTerm::Binary(mac.finalize().into_bytes().to_vec()))
        }
    }
}

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

    #[test]
    fn test_should_case_insensitive_compare() {
        let index_key = [1; 32].into();
        let indexer = UniqueIndexer::new(Default::default());

        let first = indexer
            .encrypt(&"hello WORLD".into(), &index_key)
            .expect("Failed to encrypt");
        let second = indexer
            .encrypt(&"HELLO world".into(), &index_key)
            .expect("Failed to encrypt");

        assert_eq!(first, second);
    }
}
