cipherstash-client 0.34.1-alpha.1

//! Types for representing EQL payloads, and encryption/decryption functions.

#![cfg(feature = "tokio")]

mod formats;
use formats::{format_sem_term_binary, format_sem_term_ore, format_sem_term_ore_array};

use std::borrow::Cow;
use std::sync::Arc;

use crate::{
    encryption::{self, Encrypted, EncryptedEntry, EncryptedSteVecTerm, IndexTerm, QueryOp},
    zerokms::{self, RecordDecryptError},
};

use crate::{
    encryption::StorageBuilder,
    zerokms::{GenerateKeyPayload, IndexKey},
};

use super::zerokms::EncryptedRecord;
use cipherstash_config::{column::IndexType, ColumnConfig};
use serde::{Deserialize, Serialize};
use thiserror::Error;

use uuid::Uuid;

use crate::encryption::{PlaintextTarget, Queryable, ScopedCipher};

use cipherstash_config::column::Index;

use zerokms_protocol::{Context, UnverifiedContext};

/// Encrypts multiple plaintexts into EQL format.
///
/// This is the main encryption entry point for the EQL system. It takes prepared plaintext values and encrypts them
/// into EQL payloads suitable for database storage or query generation.  The function handles both storage encryption
/// (with all SEM terms) and only SEM term generation for queries.
///
/// # Arguments
///
/// * `cipher` - The scoped cipher for performing cryptographic operations
/// * `plaintexts` - A vector of prepared plaintext values to encrypt
/// * `opts` - Encryption options including keyset ID, lock context, and service token
///
/// # Returns
///
/// A vector of `EqlCiphertext` payloads, one for each input plaintext, in the same order.
///
/// # Errors
///
/// Returns `EqlError` if:
/// - Data key generation fails
/// - Encryption of any plaintext fails
/// - Index generation fails
/// - The ZeroKMS service is unavailable
///
/// # Examples
///
/// ```no_run
/// # use std::sync::Arc;
/// # use cipherstash_client::eql::{encrypt_eql, PreparedPlaintext, EqlEncryptOpts};
/// # use cipherstash_client::encryption::ScopedCipher;
/// # use cipherstash_client::AutoStrategy;
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// # let cipher: Arc<ScopedCipher<AutoStrategy>> = unimplemented!();
/// # let plaintexts: Vec<PreparedPlaintext> = vec![];
/// let opts = EqlEncryptOpts::default();
/// let encrypted = encrypt_eql(cipher, plaintexts, &opts).await?;
/// # Ok(())
/// # }
/// # fn main() {}
/// ```
pub async fn encrypt_eql<'a, C>(
    cipher: Arc<ScopedCipher<C>>,
    plaintexts: Vec<PreparedPlaintext<'a>>,
    opts: &EqlEncryptOpts<'a>,
) -> Result<Vec<EqlCiphertext>, EqlError>
where
    C: Send + Sync + 'static,
    for<'b> &'b C: stack_auth::AuthStrategy,
{
    use std::collections::VecDeque;

    let effective_keyset_id = opts.keyset_id.unwrap_or(cipher.keyset_id());

    let targets: Vec<EncryptionTarget> =
        to_encryption_targets(cipher.index_key(), plaintexts, effective_keyset_id)?;

    let mut data_keys = VecDeque::from(
        cipher
            .generate_data_keys(
                generate_data_key_payloads(opts, &targets),
                opts.service_token.clone(),
                opts.unverified_context.clone(),
            )
            .await?,
    );

    let encrypted = targets
        .into_iter()
        .map(|target| -> Result<EqlCiphertext, EqlError> {
            match target {
                EncryptionTarget::ForStorage(identifier, builder) => {
                    let encrypted = builder.build_for_encryption().encrypt(
                        // PANIC SAFETY: we trust that a successful result from `generate_data_keys` returned the
                        // precise number of keys that we requested, which means `data_keys.remove(0)` cannot fail.
                        data_keys
                            .remove(0)
                            .expect("insufficient data keys to encrypt all plaintexts"),
                    )?;

                    Ok(to_eql_ciphertext(encrypted, &identifier)?)
                }
                EncryptionTarget::ForQuery(identifier, plaintext, index_type, query_op) => {
                    let index = Index::new(index_type.clone());
                    let index_term =
                        (index, plaintext).build_queryable(cipher.clone(), query_op)?;

                    Ok(to_eql_ciphertext_from_sem_term(index_term, &identifier)?)
                }
            }
        })
        .collect::<Result<Vec<_>, _>>()?;

    Ok(encrypted)
}

/// Decrypts multiple EQL encrypted payloads back to plaintext.
///
/// This is the main decryption entry point for the EQL system. It takes encrypted EQL
/// payloads (as retrieved from the database) and decrypts them using ZeroKMS, returning
/// the original plaintext values.
///
/// # Arguments
///
/// * `cipher` - The scoped cipher for performing cryptographic operations
/// * `ciphertexts` - An iterator of encrypted EQL payloads to decrypt
/// * `opts` - Decryption options including keyset ID, lock context, and service token
///
/// # Returns
///
/// A vector of decrypted `Plaintext` values, one for each input ciphertext, in the same order.
///
/// # Errors
///
/// Returns `EqlError` if:
/// - Any ciphertext payload is missing its ciphertext field
/// - Decryption fails (e.g., wrong keyset, tampered data)
/// - The ZeroKMS service is unavailable
/// - The decrypted data cannot be parsed as plaintext
///
/// # Examples
///
/// ```no_run
/// # use std::sync::Arc;
/// # use cipherstash_client::eql::{decrypt_eql, EqlCiphertext, EqlDecryptOpts};
/// # use cipherstash_client::encryption::ScopedCipher;
/// # use cipherstash_client::AutoStrategy;
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// # let cipher: Arc<ScopedCipher<AutoStrategy>> = unimplemented!();
/// # let ciphertexts: Vec<EqlCiphertext> = vec![];
/// let opts = EqlDecryptOpts::default();
/// let plaintexts = decrypt_eql(cipher, ciphertexts, &opts).await?;
/// # Ok(())
/// # }
/// # fn main() {}
/// ```
pub async fn decrypt_eql<'a, C>(
    cipher: Arc<ScopedCipher<C>>,
    ciphertexts: impl IntoIterator<Item = EqlCiphertext>,
    opts: &EqlDecryptOpts<'a>,
) -> Result<Vec<encryption::Plaintext>, EqlError>
where
    C: Send + Sync + 'static,
    for<'b> &'b C: stack_auth::AuthStrategy,
{
    use crate::{encryption::DecryptOptions, zerokms::WithContext};

    let decrypt_opts = DecryptOptions {
        keyset_id: opts.keyset_id,
        service_token: opts.service_token.clone(),
        unverified_context: opts.unverified_context.clone(),
    };

    let ciphertexts = ciphertexts
        .into_iter()
        .map(|eql| {
            let ciphertext = eql
                .body
                .ciphertext
                .ok_or_else(|| EqlError::MissingCiphertext(eql.identifier.clone()))?;
            Ok(WithContext {
                record: ciphertext,
                context: opts.lock_context.clone(),
            })
        })
        .collect::<Result<Vec<_>, EqlError>>()?;

    Ok(cipher
        .decrypt(ciphertexts, &decrypt_opts)
        .await
        .map_err(|err| convert_zerokms_error(err, cipher.keyset_id(), opts.keyset_id))?
        .into_iter()
        .map(|decrypted| encryption::Plaintext::from_slice(&decrypted))
        .collect::<Result<Vec<_>, _>>()?)
}

/// Like [`decrypt_eql`] but decryption failure of one or more ciphertexts does not fail the entire operation.
/// Instead, a decryption result (success or failure) is returned for every ciphertext.
///
/// # Per-Item Errors
///
/// The following conditions produce per-item `Err` results without aborting:
/// - Missing ciphertext (query-mode payload): Returns `Err(EqlError::MissingCiphertext)`
/// - Decryption failure: Returns `Err` with the underlying decryption error
///
/// # Operation-Level Errors
///
/// Only a failure to retrieve the access token will result in the entire operation failing.
pub async fn decrypt_eql_fallible<'a, C>(
    cipher: Arc<ScopedCipher<C>>,
    ciphertexts: impl IntoIterator<Item = EqlCiphertext>,
    opts: &EqlDecryptOpts<'a>,
) -> Result<Vec<Result<encryption::Plaintext, EqlError>>, EqlError>
where
    C: Send + Sync + 'static,
    for<'b> &'b C: stack_auth::AuthStrategy,
{
    use crate::{encryption::DecryptOptions, zerokms::WithContext};

    let decrypt_opts = DecryptOptions {
        keyset_id: opts.keyset_id,
        service_token: opts.service_token.clone(),
        unverified_context: opts.unverified_context.clone(),
    };

    // Collect and pre-allocate (matches vitur_client pattern)
    let inputs: Vec<_> = ciphertexts.into_iter().collect();
    let input_count = inputs.len();
    let mut results: Vec<Option<Result<encryption::Plaintext, EqlError>>> =
        (0..input_count).map(|_| None).collect();
    let mut valid_payloads: Vec<(usize, WithContext<EncryptedRecord>)> =
        Vec::with_capacity(input_count);

    // Single pass: partition inputs, pre-fill missing ciphertext errors
    for (index, eql) in inputs.into_iter().enumerate() {
        match eql.body.ciphertext {
            Some(ciphertext) => {
                valid_payloads.push((
                    index,
                    WithContext {
                        record: ciphertext,
                        context: opts.lock_context.clone(),
                    },
                ));
            }
            None => {
                results[index] = Some(Err(EqlError::MissingCiphertext(eql.identifier)));
            }
        }
    }

    // Separate indices and payloads (no clone needed)
    let (indices, payloads): (Vec<usize>, Vec<_>) = valid_payloads.into_iter().unzip();

    // Decrypt valid payloads only
    let decrypt_results = cipher
        .decrypt_fallible(payloads, &decrypt_opts)
        .await
        .map_err(|err| convert_zerokms_error(err, cipher.keyset_id(), opts.keyset_id))?;

    // Fill results at tracked indices
    for (index, decrypt_result) in indices.into_iter().zip(decrypt_results) {
        results[index] = Some(match decrypt_result {
            Ok(bytes) => encryption::Plaintext::from_slice(&bytes).map_err(Into::into),
            Err(err) => Err(EqlError::from(err)),
        });
    }

    // Unwrap - invariant: all slots filled by either pre-filled error or decrypt result
    Ok(results
        .into_iter()
        .map(|r| r.expect("all result slots filled"))
        .collect())
}

/// The current version of the EQL schema format.
///
/// This version number is included in all encrypted payloads to ensure compatibility
/// between different versions of the encryption system.
pub const EQL_SCHEMA_VERSION: u16 = 2;

/// Identifies a specific database table and column pair.
///
/// This type is used throughout the EQL system to uniquely identify which table and column
/// a piece of data belongs to. It implements `Hash` and `Eq` to allow use as a map key.
///
/// # Fields
///
/// * `table` - The database table name (serialized as "t")
/// * `column` - The column name within the table (serialized as "c")
#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct Identifier {
    /// The database table name (serialized as "t").
    #[serde(rename = "t")]
    pub table: String,

    /// The column name within the table (serialized as "c").
    #[serde(rename = "c")]
    pub column: String,
}

impl Identifier {
    /// Creates a new `Identifier` from table and column names.
    ///
    /// # Arguments
    ///
    /// * `table` - The table name (can be any type convertible to `String`)
    /// * `column` - The column name (can be any type convertible to `String`)
    ///
    /// # Examples
    ///
    /// ```
    /// use cipherstash_client::eql::Identifier;
    ///
    /// let id = Identifier::new("users", "email");
    /// ```
    pub fn new(table: impl Into<String>, column: impl Into<String>) -> Self {
        let table = table.into();
        let column = column.into();

        Self { table, column }
    }

    /// Returns a reference to the table name.
    pub fn table(&self) -> &str {
        &self.table
    }

    /// Returns a reference to the column name.
    pub fn column(&self) -> &str {
        &self.column
    }
}

/// Represents encrypted data in the EQL format.
///
/// This is the top-level structure for encrypted data, containing the identifier, version, and the encrypted body with
/// all associated cryptographic searchable encrypted metadata.
///
/// # Fields
///
/// * `identifier` - The table and column this encrypted data belongs to (serialized as "i")
/// * `version` - The encryption version used (serialized as "v")
/// * `body` - The encrypted ciphertext and indexes (flattened into parent during serialization)
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct EqlCiphertext {
    /// The table and column this encrypted data belongs to (serialized as "i").
    #[serde(rename = "i")]
    pub identifier: Identifier,

    /// The encryption version used (serialized as "v").
    #[serde(rename = "v")]
    pub version: u16,

    /// The encrypted ciphertext and indexes (flattened into parent during serialization).
    #[serde(flatten)]
    pub body: EqlCiphertextBody,
}

/// Contains the encrypted ciphertext and associated cryptographic indexes.
///
/// This structure holds the actual encrypted data along with various encrypted indexes
/// that enable different types of queries on the encrypted data.
///
/// # Fields
///
/// * `ciphertext` - The encrypted record (serialized as "c" in MessagePack Base85 format)
/// * `indexes` - All cryptographic indexes for this encrypted value (flattened into parent)
/// * `is_array_item` - Whether this encrypted value is part of an array (serialized as "a")
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct EqlCiphertextBody {
    /// The encrypted record (serialized as "c" in MessagePack Base85 format).
    #[serde(
        rename = "c",
        default,
        with = "formats::mp_base85",
        skip_serializing_if = "Option::is_none"
    )]
    pub ciphertext: Option<EncryptedRecord>,

    /// All cryptographic searchable encrypted metadata (SEM) for this encrypted value (flattened into parent).
    #[serde(flatten)]
    pub sem: EqlSEM,

    /// Whether this encrypted value is part of an array (serialized as "a").
    #[serde(rename = "a", skip_serializing_if = "Option::is_none")]
    pub is_array_item: Option<bool>,
}

/// Struct that can carry values of all possible cryptographic searchable encrypted metadata (SEM) types.
///
/// This structure contains various SEM types that enable different query operations on encrypted data. Each SEM type
/// supports specific query patterns:
///
/// - **ORE indexes** enable range queries (greater than, less than, etc.)
/// - **Bloom filters** enable approximate membership testing
/// - **HMAC/Blake3** enable exact match queries
/// - **SteVec** enables vector/array containment queries
///
/// # Serialization
///
/// - All fields are optional and null values are not serialized
/// - This allows the database to store only the indexes actually needed for each field
/// - The JSON format matches the EQL specification for database storage
///
/// # Fields
///
/// * `ore_block_u64_8_256` - ORE block index for 64-bit integers (serialized as "ob")
/// * `bloom_filter` - Bloom filter for approximate match queries (serialized as "bf")
/// * `hmac_256` - HMAC-SHA256 hash for exact matches (serialized as "hm")
/// * `selector` - Selector value for deterministic encryption (serialized as "s")
/// * `blake3` - Blake3 hash for exact matches (serialized as "b3")
/// * `ore_cllw_u64_8` - ORE CLLW fixed-width index for 64-bit values (serialized as "ocf")
/// * `ore_cllw_var_8` - ORE CLLW variable-width index (serialized as "ocv")
/// * `ste_vec_index` - Structured encryption vector index for containment queries (serialized as "sv")
#[derive(Clone, Debug, Deserialize, Serialize, Default)]
pub struct EqlSEM {
    /// SEM term for the 64-bit Block ORE scheme (serialized as "ob").
    #[serde(rename = "ob", skip_serializing_if = "Option::is_none")]
    pub ore_block_u64_8_256: Option<Vec<String>>,

    /// Bloom filter for approximate match queries (serialized as "bf").
    #[serde(rename = "bf", skip_serializing_if = "Option::is_none")]
    pub bloom_filter: Option<Vec<u16>>,

    /// HMAC-SHA256 hash for exact matches (serialized as "hm").
    #[serde(rename = "hm", skip_serializing_if = "Option::is_none")]
    pub hmac_256: Option<String>,

    /// Selector to field-selection in an SteVec (for Json queries).
    #[serde(rename = "s", skip_serializing_if = "Option::is_none")]
    pub selector: Option<String>,

    /// Blake3 hash for exact matches in a structured encryption vector AKA `SteVec` (serialized as "b3").
    #[serde(rename = "b3", skip_serializing_if = "Option::is_none")]
    pub blake3: Option<String>,

    /// ORE CLLW fixed-width scheme for 64-bit values in a structured encryption vector AKA `SteVec` (serialized as
    /// "ocf").
    #[serde(rename = "ocf", skip_serializing_if = "Option::is_none")]
    pub ore_cllw_u64_8: Option<String>,

    /// ORE CLLW variable-width scheme used for string comparison in an `SteVec` (serialized as "ocv")
    #[serde(rename = "ocv", skip_serializing_if = "Option::is_none")]
    pub ore_cllw_var_8: Option<String>,

    /// Structured encryption vector (STE) scheme for supporting containment queries and a subset of JSON-style
    /// operations (serialized as "sv").
    #[serde(rename = "sv", skip_serializing_if = "Option::is_none")]
    pub ste_vec: Option<Vec<EqlCiphertextBody>>,
}

impl EqlSEM {
    /// Updates `self` from an [`IndexTerm`].
    fn update(&mut self, term: IndexTerm) {
        match term {
            IndexTerm::Binary(bytes) => {
                self.hmac_256 = Some(format_sem_term_binary(&bytes));
            }
            IndexTerm::BinaryVec(_) => unimplemented!("this is not used by EQL"),
            IndexTerm::BitMap(bloom_filter) => {
                self.bloom_filter = Some(bloom_filter);
            }
            IndexTerm::OreFull(bytes) => {
                self.ore_block_u64_8_256 = Some(format_sem_term_ore(&bytes));
            }
            IndexTerm::OreArray(bytes) => {
                self.ore_block_u64_8_256 = Some(format_sem_term_ore_array(&bytes));
            }
            IndexTerm::OreLeft(bytes) => {
                self.ore_block_u64_8_256 = Some(format_sem_term_ore(&bytes));
            }
            IndexTerm::SteVecSelector(selector) => {
                self.selector = Some(hex::encode(selector.as_bytes()));
            }
            IndexTerm::SteVecTerm(ste_vec_term) => match ste_vec_term {
                EncryptedSteVecTerm::Mac(bytes) => self.blake3 = Some(hex::encode(bytes)),
                EncryptedSteVecTerm::OreFixed(ore) => self.ore_cllw_u64_8 = Some(hex::encode(ore)),
                EncryptedSteVecTerm::OreVariable(ore) => {
                    self.ore_cllw_var_8 = Some(hex::encode(&ore))
                }
            },
            IndexTerm::SteQueryVec(_) => {}
            IndexTerm::Null => {}
        };
    }
}

/// Converts an `IndexTerm` into an `EqlCiphertext` structure.
///
/// This function is specifically designed to handle `SteVecSelector` index terms, creating an encrypted payload with
/// only the selector index populated.
///
/// # Arguments
///
/// * `index_term` - The index term to convert (must be `IndexTerm::SteVecSelector`)
/// * `identifier` - The table and column identifier for this encrypted data
///
/// # Returns
///
/// Returns an `EqlCiphertext` structure with the selector index populated, or an
/// `EqlError::InvalidIndexTerm` if the index term is not a `SteVecSelector`.
///
/// # Errors
///
/// Returns `EqlError::InvalidIndexTerm` if the provided index term is not a `SteVecSelector`.
fn to_eql_ciphertext_from_sem_term(
    index_term: IndexTerm,
    identifier: &Identifier,
) -> Result<EqlCiphertext, EqlError> {
    let mut sem = EqlSEM::default();
    sem.update(index_term);

    Ok(EqlCiphertext {
        identifier: identifier.to_owned(),
        version: EQL_SCHEMA_VERSION,
        body: EqlCiphertextBody {
            ciphertext: None,
            sem,
            is_array_item: None,
        },
    })
}

/// Converts an `Encrypted` value into an `EqlCiphertext` structure.
///
/// This is the primary conversion function for transforming encrypted data and its associated SEM into the EQL format
/// suitable for database storage. It handles both regular encrypted records and structured encryption vectors (SteVec).
///
/// # Arguments
///
/// * `encrypted` - The encrypted data with its cryptographic SEM
/// * `identifier` - The table and column identifier for this encrypted data
///
/// # Returns
///
/// Returns an `EqlCiphertext` structure containing the ciphertext and all applicable SEM terms, formatted according to
/// the EQL specification.
///
/// # Errors
///
/// Returns an `EqlError` if the encrypted data cannot be properly converted to EQL format.
fn to_eql_ciphertext(
    encrypted: Encrypted,
    identifier: &Identifier,
) -> Result<EqlCiphertext, EqlError> {
    let mut sem = EqlSEM::default();

    match encrypted {
        Encrypted::Record(ciphertext, terms) => {
            for term in terms {
                sem.update(term);
            }

            Ok(EqlCiphertext {
                identifier: identifier.to_owned(),
                version: EQL_SCHEMA_VERSION,
                body: EqlCiphertextBody {
                    ciphertext: Some(ciphertext),
                    sem,
                    is_array_item: None,
                },
            })
        }
        Encrypted::SteVec(ste_vec) => {
            let ciphertext = ste_vec.root_ciphertext()?.clone();

            let ste_vec_sem: Vec<EqlCiphertextBody> = ste_vec
                .into_iter()
                .map(
                    |EncryptedEntry {
                         tokenized_selector,
                         term,
                         record,
                         parent_is_array,
                     }| {
                        let sem = match term {
                            EncryptedSteVecTerm::Mac(bytes) => EqlSEM {
                                selector: Some(hex::encode(tokenized_selector.as_bytes())),
                                blake3: Some(hex::encode(bytes)),
                                ..Default::default()
                            },
                            EncryptedSteVecTerm::OreFixed(ore) => EqlSEM {
                                selector: Some(hex::encode(tokenized_selector.as_bytes())),
                                ore_cllw_u64_8: Some(hex::encode(ore)),
                                ..Default::default()
                            },
                            EncryptedSteVecTerm::OreVariable(ore) => EqlSEM {
                                selector: Some(hex::encode(tokenized_selector.as_bytes())),
                                ore_cllw_var_8: Some(hex::encode(&ore)),
                                ..Default::default()
                            },
                        };

                        EqlCiphertextBody {
                            ciphertext: Some(record),
                            sem,
                            is_array_item: Some(parent_is_array),
                        }
                    },
                )
                .collect();

            sem.ste_vec = Some(ste_vec_sem);

            Ok(EqlCiphertext {
                identifier: identifier.to_owned(),
                version: EQL_SCHEMA_VERSION,
                body: EqlCiphertextBody {
                    ciphertext: Some(ciphertext),
                    sem,
                    is_array_item: None,
                },
            })
        }
    }
}

/// Errors that can occur during EQL encryption, decryption, and index operations.
///
/// This enum covers all error cases in the EQL system, including serialization failures,
/// configuration mismatches, keyset management issues, and cryptographic operation errors.
#[derive(Error, Debug)]
pub enum EqlError {
    /// Failed to serialize ciphertext to JSON format.
    #[error(transparent)]
    CiphertextCouldNotBeSerialised(#[from] serde_json::Error),

    /// An encrypted column value could not be parsed from the database format.
    #[error("Encrypted column could not be parsed")]
    ColumnCouldNotBeParsed,

    /// An encrypted column value was NULL when a value was expected.
    #[error("Encrypted column is null")]
    ColumnIsNull,

    /// Failed to deserialize an encrypted column from the database.
    #[error("Column '{column}' in table '{table}' could not be deserialised")]
    ColumnCouldNotBeDeserialised { table: String, column: String },

    /// Failed to encrypt a column value.
    #[error("Column '{column}' in table '{table}' could not be encrypted")]
    ColumnCouldNotBeEncrypted { table: String, column: String },

    /// The encryption configuration for a column doesn't match the encrypted data.
    #[error("Column configuration for column '{column}' in table '{table}' does not match the encrypted column")]
    ColumnConfigurationMismatch { table: String, column: String },

    /// Decryption failed using the specified keyset.
    #[error("Could not decrypt data using keyset '{keyset_id}'")]
    CouldNotDecryptDataForKeyset { keyset_id: String },

    /// An index term was not of the expected type for the operation.
    #[error("InvalidIndexTerm")]
    InvalidIndexTerm,

    /// The EQL payload is missing its ciphertext field.
    ///
    /// This typically occurs when attempting to decrypt a query-mode payload,
    /// which only contains searchable encrypted metadata (SEM) terms.
    #[error("EQL payload for column '{}' in table '{}' is missing ciphertext", _0.column(), _0.table())]
    MissingCiphertext(Identifier),

    /// The keyset ID provided via `SET CIPHERSTASH.KEYSET_ID` is not a valid UUID.
    #[error("KeysetId `{id}` could not be parsed using `SET CIPHERSTASH.KEYSET_ID`. KeysetId should be a valid UUID")]
    KeysetIdCouldNotBeParsed { id: String },

    /// Failed to set the keyset ID using `SET CIPHERSTASH.KEYSET_ID`.
    #[error("Keyset Id could not be set using `SET CIPHERSTASH.KEYSET_ID`")]
    KeysetIdCouldNotBeSet,

    /// Failed to set the keyset name using `SET CIPHERSTASH.KEYSET_NAME`.
    #[error("Keyset Name could not be set using `SET CIPHERSTASH.KEYSET_NAME`")]
    KeysetNameCouldNotBeSet,

    /// Missing encryption configuration for a specific column type.
    ///
    /// This should be unreachable in practice.
    #[error("Missing encrypt configuration for column type `{plaintext_type}`")]
    MissingEncryptConfiguration { plaintext_type: &'static str },

    /// Failed to encode decrypted plaintext as the expected data type.
    #[error("Decrypted column could not be encoded as the expected type")]
    PlaintextCouldNotBeEncoded,

    /// An error occurred in the encryption pipeline.
    #[error(transparent)]
    Pipeline(#[from] encryption::EncryptionError),

    /// Failed to decode plaintext from the expected type format.
    #[error(transparent)]
    PlaintextCouldNotBeDecoded(#[from] encryption::TypeParseError),

    /// No keyset identifier was provided when one was required.
    #[error("Missing keyset identifer")]
    MissingKeysetIdentifier,

    /// Attempted to set a keyset when a default keyset is already configured.
    #[error("Cannot SET CIPHERSTASH.KEYSET if a default keyset has been configured")]
    UnexpectedSetKeyset,

    /// The specified column has no encryption configuration.
    #[error("Column '{column}' in table '{table}' has no Encrypt configuration")]
    UnknownColumn { table: String, column: String },

    /// The keyset name or ID is not found in the configured credentials.
    #[error("Unknown keyset name or id '{keyset}'. Check the configured credentials")]
    UnknownKeysetIdentifier { keyset: String },

    /// The specified table has no encryption configuration.
    #[error("Table '{table}' has no Encrypt configuration")]
    UnknownTable { table: String },

    /// An unknown or unsupported index term type was encountered for the column.
    #[error("Unknown Index Term for column '{}' in table '{}'", _0.column(), _0.table())]
    UnknownIndexTerm(Identifier),

    /// A ZeroKMS error
    #[error("ZeroKMS error '{}'", _0)]
    ZeroKMS(#[from] zerokms::Error),

    /// A record decryption error
    #[error("Record decryption error '{}'", _0)]
    RecordDecrypt(#[from] RecordDecryptError),
}

/// Specifies what to encrypt when encrypting [`encryption::Plaintext`] values.
///
/// `EqlOperation` determines whether to perform full encryption for storage or just generate a specific encrypted
/// searchable encrypted metadata (SEM) term.
#[derive(Debug)]
pub enum EqlOperation<'a> {
    /// Encrypt both the plaintext and its associated encrypted search terms.
    ///
    /// Use this eql_op for persisting encrypted data while keeping it searchable.
    Store,

    /// Generate an encrypted search term for this specific index type.
    ///
    /// The tuple contains the index type and query operation to use.
    Query(&'a IndexType, QueryOp),
}

/// A prepared plaintext value ready for EQL encryption.
///
/// `PreparedPlaintext` bundles together all the information needed to encrypt a plaintext value into an EQL payload:
/// the value itself, its location identifier (table/column), the encryption eql_op (storage or query), and the column's
/// encryption configuration.
///
/// # Fields
///
/// * `identifier` - The table and column this plaintext belongs to
/// * `plaintext` - The actual plaintext value to encrypt
/// * `eql_op` - Whether to encrypt for storage or generate a query index
/// * `column_config` - The encryption configuration for this column
///
/// # Examples
///
/// ```no_run
/// use cipherstash_client::eql::{PreparedPlaintext, Identifier, EqlOperation};
/// use cipherstash_client::encryption::Plaintext;
/// use std::borrow::Cow;
/// # use cipherstash_config::ColumnConfig;
///
/// let identifier = Identifier::new("users", "email");
/// let plaintext = Plaintext::Utf8Str(Some("user@example.com".into()));
/// # let column_config: ColumnConfig = unimplemented!();
///
/// let prepared = PreparedPlaintext::new(
///     Cow::Owned(column_config),
///     identifier,
///     plaintext,
///     EqlOperation::Store,
/// );
/// ```
pub struct PreparedPlaintext<'a> {
    /// The table and column this plaintext belongs to.
    identifier: Identifier,

    /// The actual plaintext value to encrypt.
    plaintext: encryption::Plaintext,

    /// Whether to encrypt for storage or generate a cryptographic searchable encrypted metadata (SEM) term.
    eql_op: EqlOperation<'a>,

    /// The encryption configuration for this column.
    column_config: Cow<'a, ColumnConfig>,
}

impl<'a> PreparedPlaintext<'a> {
    pub fn new(
        column_config: Cow<'a, ColumnConfig>,
        identifier: Identifier,
        plaintext: encryption::Plaintext,
        eql_op: EqlOperation<'a>,
    ) -> Self {
        Self {
            identifier,
            plaintext,
            eql_op,
            column_config,
        }
    }
}

/// Internal representation of an encryption target during EQL encryption.
///
/// This enum distinguishes between encrypting data for storage (which generates both ciphertext and _all_ configured
/// SEM terms) and generating a single query encrypted SEM term for a query.
enum EncryptionTarget<'a> {
    /// Encrypt plaintext for storage with all configured SEM term types.
    ///
    /// Contains the identifier and a configured storage builder.
    ForStorage(Identifier, StorageBuilder<'a, encryption::Plaintext>),

    /// Generate a single SEM term for a specific SEM term type.
    ///
    /// Contains the identifier, plaintext, SEM type, and query operation.
    ForQuery(Identifier, encryption::Plaintext, &'a IndexType, QueryOp),
}

/// Generates data key payloads for encryption targets that require storage.
///
/// This function filters the encryption targets to identify those that need data key
/// generation (i.e., `ForStorage` targets) and creates the corresponding key generation
/// payloads for the ZeroKMS service.
///
/// # Arguments
///
/// * `opts` - Encryption options containing the lock context
/// * `targets` - The encryption targets to generate keys for
///
/// # Returns
///
/// A vector of `GenerateKeyPayload` objects, one for each storage target. Query-only
/// targets are excluded as they don't require new data keys.
fn generate_data_key_payloads<'a>(
    opts: &EqlEncryptOpts<'a>,
    targets: &'a Vec<EncryptionTarget<'a>>,
) -> Vec<GenerateKeyPayload<'a>> {
    targets
        .iter()
        .filter_map(|target| match target {
            EncryptionTarget::ForStorage(_, builder) => Some(GenerateKeyPayload::new(
                builder.descriptor(),
                opts.lock_context.clone(),
            )),
            EncryptionTarget::ForQuery(_, _, _, _) => None,
        })
        .collect()
}

/// Converts prepared plaintexts into encryption targets.
///
/// This function transforms high-level `PreparedPlaintext` values into low-level
/// `EncryptionTarget` representations, preparing them for encryption. It distinguishes
/// between storage encryption (which uses `StorageBuilder`) and query index generation.
///
/// # Arguments
///
/// * `index_key` - The index key for generating searchable indexes
/// * `plaintexts` - The prepared plaintext values to convert
/// * `effective_keyset_id` - The keyset ID to use for encryption
///
/// # Returns
///
/// A vector of `EncryptionTarget` values ready for encryption.
///
/// # Errors
///
/// Returns `EncryptionError` if the plaintext values cannot be properly configured
/// for encryption (e.g., invalid column configuration).
fn to_encryption_targets<'a>(
    index_key: &'a IndexKey,
    plaintexts: Vec<PreparedPlaintext<'a>>,
    effective_keyset_id: Uuid,
) -> Result<Vec<EncryptionTarget<'a>>, encryption::EncryptionError> {
    plaintexts
        .into_iter()
        .map(
            move |prepared_plaintext| -> Result<EncryptionTarget, encryption::EncryptionError> {
                use crate::encryption::Encryptable;

                let PreparedPlaintext {
                    identifier,
                    plaintext,
                    eql_op,
                    column_config,
                } = prepared_plaintext;

                match eql_op {
                    EqlOperation::Store => Ok(EncryptionTarget::ForStorage(
                        identifier,
                        PlaintextTarget::new(plaintext, (*column_config).clone())
                            .build_encryptable(index_key, effective_keyset_id)?,
                    )),
                    EqlOperation::Query(index_type, query_op) => Ok(EncryptionTarget::ForQuery(
                        identifier, plaintext, index_type, query_op,
                    )),
                }
            },
        )
        .collect::<Result<Vec<_>, _>>()
}

/// Options for EQL decryption operations.
///
/// This structure provides configuration options for the [`decrypt_eql`] function,
/// allowing control over which keyset to use, the lock context for key unwrapping,
/// and authentication tokens.
///
/// # Fields
///
/// * `keyset_id` - Optional keyset UUID; if `None`, uses the cipher's default keyset
/// * `lock_context` - Context information for unwrapping encrypted keys
/// * `service_token` - Optional service authentication token
/// * `unverified_context` - Optional unverified context for additional metadata
///
/// # Examples
///
/// ```
/// use cipherstash_client::eql::EqlDecryptOpts;
/// use std::borrow::Cow;
///
/// let opts = EqlDecryptOpts {
///     keyset_id: None, // Use default keyset
///     lock_context: Cow::Borrowed(&[]),
///     service_token: None,
///     unverified_context: None,
/// };
/// ```
#[derive(Debug, Default)]
pub struct EqlDecryptOpts<'a> {
    /// Keyset to use for all decryption operations.
    ///
    /// When set to `None`, decryption will use the keyset associated with the ZeroKMS client.
    pub keyset_id: Option<Uuid>,

    /// The lock context to use for decryption operations.
    pub lock_context: Cow<'a, [Context]>,

    /// Optional service authentication token for ZeroKMS requests.
    pub service_token: Option<Cow<'a, crate::credentials::ServiceToken>>,

    /// Optional unverified context for additional metadata.
    pub unverified_context: Option<Cow<'a, UnverifiedContext>>,
}

/// Options for EQL encryption operations.
///
/// This structure provides configuration options for the [`encrypt_eql`] function,
/// controlling which keyset to use, the lock context for key wrapping, authentication
/// tokens, and which indexes to generate.
///
/// # Fields
///
/// * `keyset_id` - Optional keyset UUID; if `None`, uses the cipher's default keyset
/// * `lock_context` - Context information for wrapping data encryption keys
/// * `service_token` - Optional service authentication token
/// * `unverified_context` - Optional unverified context for additional metadata
/// * `index_types` - Optional filter to generate only specific index types instead of all configured indexes
///
/// # Examples
///
/// ```
/// use cipherstash_client::eql::EqlEncryptOpts;
/// use std::borrow::Cow;
///
/// // Use default options (all indexes, default keyset)
/// let opts = EqlEncryptOpts::default();
///
/// // Or customize
/// # use cipherstash_config::column::IndexType;
/// let opts = EqlEncryptOpts {
///     keyset_id: None,
///     lock_context: Cow::Owned(vec![]),
///     service_token: None,
///     unverified_context: None,
///     index_types: Some(Cow::Owned(vec![IndexType::Ore])), // Generate only ORE indexes
/// };
/// ```
#[derive(Debug, Default)]
pub struct EqlEncryptOpts<'a> {
    /// Keyset to use for all encryption operations.
    ///
    /// When set to `None`, encryption will use the keyset associated with the ZeroKMS client.
    pub keyset_id: Option<Uuid>,

    /// The lock context to use for encryption operations.
    pub lock_context: Cow<'a, [Context]>,

    /// Optional service authentication token for ZeroKMS requests.
    pub service_token: Option<Cow<'a, crate::credentials::ServiceToken>>,

    /// Optional unverified context for additional metadata.
    pub unverified_context: Option<Cow<'a, UnverifiedContext>>,

    /// Optional filter to generate only specific index types.
    ///
    /// When present, generate _only_ the specified index types instead of all index types
    /// specified in the config.
    pub index_types: Option<Cow<'a, [IndexType]>>,
}

/// Converts ZeroKMS errors into more user-friendly EQL encryption errors.
///
/// This function inspects ZeroKMS errors and provides more specific error messages,
/// particularly for common cases like missing or incorrect keysets.
///
/// # Arguments
///
/// * `err` - The ZeroKMS error to convert
/// * `keyset_id` - The optional keyset ID that was used in the operation
///
/// # Returns
///
/// An `EqlError` with a more specific error message when possible, or the
/// original ZeroKMS error wrapped in `EqlError::ZeroKMS`.
fn convert_zerokms_error(
    err: zerokms::Error,
    cipher_keyset_id: Uuid,
    keyset_id_override: Option<Uuid>,
) -> EqlError {
    match err {
        zerokms::Error::Decrypt(_) => {
            let error_msg = err.to_string();
            if error_msg.contains("Failed to retrieve key") {
                EqlError::CouldNotDecryptDataForKeyset {
                    keyset_id: keyset_id_override
                        .map(|id| id.to_string())
                        .unwrap_or(cipher_keyset_id.to_string()),
                }
            } else {
                EqlError::ZeroKMS(err)
            }
        }
        _ => EqlError::ZeroKMS(err),
    }
}

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

    #[test]
    fn test_eql_ciphertext_missing_ciphertext_field() {
        // Create an EqlCiphertext with no ciphertext (simulating query-mode payload)
        let identifier = Identifier::new("test_table", "test_column");
        let eql = EqlCiphertext {
            identifier: identifier.clone(),
            version: EQL_SCHEMA_VERSION,
            body: EqlCiphertextBody {
                ciphertext: None, // Missing ciphertext
                sem: EqlSEM::default(),
                is_array_item: None,
            },
        };

        // Verify that attempting to extract ciphertext returns MissingCiphertext error
        let result = eql
            .body
            .ciphertext
            .ok_or_else(|| EqlError::MissingCiphertext(identifier));
        assert!(matches!(result, Err(EqlError::MissingCiphertext(_))));
    }

    #[test]
    fn test_mp_base85_deserialize_invalid_input_returns_error() {
        use serde::de::value::{Error as ValueError, StrDeserializer};
        use serde::de::IntoDeserializer;

        // Invalid base85 string that should fail to decode
        let invalid: StrDeserializer<ValueError> = "not-valid-base85!!!".into_deserializer();

        // This tests that invalid base85 input returns an error, not Ok(None)
        let result: Result<Option<EncryptedRecord>, ValueError> =
            formats::mp_base85::deserialize(invalid);

        assert!(
            result.is_err(),
            "Invalid base85 input should return error, not Ok(None)"
        );
    }

    #[test]
    fn test_fallible_contract_missing_ciphertext_is_per_item_error() {
        // This test documents the fallible contract:
        // - Missing ciphertext should produce Err in result vector
        // - Should NOT abort the entire operation
        //
        // The actual async decrypt_eql_fallible requires integration test setup,
        // but we verify the error type and contract expectation here.

        let identifier = Identifier::new("test_table", "test_column");
        let eql = EqlCiphertext {
            identifier: identifier.clone(),
            version: EQL_SCHEMA_VERSION,
            body: EqlCiphertextBody {
                ciphertext: None, // Missing ciphertext (query-mode payload)
                sem: EqlSEM::default(),
                is_array_item: None,
            },
        };

        // Per the fallible contract, missing ciphertext should be a per-item error:
        // Ok(vec![Err(MissingCiphertext)]) - NOT - Err(MissingCiphertext)
        let per_item_error: Result<encryption::Plaintext, EqlError> =
            Err(EqlError::MissingCiphertext(eql.identifier.clone()));

        assert!(matches!(
            per_item_error,
            Err(EqlError::MissingCiphertext(_))
        ));
    }
}