use crate::zerokms::{
    self, DataKeyWithTag, Decryptable, EncryptPayload, EncryptedRecord, Error, GenerateKeyPayload,
    IndexKey, RecordDecryptError, ZeroKMSWithClientKey,
};
use stack_auth::AuthStrategy;
use std::{borrow::Cow, fmt::Debug, sync::Arc};
use uuid::Uuid;
use zeroize::{Zeroize, ZeroizeOnDrop};
use zerokms_protocol::{IdentifiedBy, UnverifiedContext};

use stack_auth::AuthStrategyBounds;

use super::{
    compound_indexer::{Accumulator, ComposableIndex, ComposablePlaintext, CompoundIndex},
    EncryptionError, IndexTerm,
};

const DEFAULT_INDEX_TERM_SIZE: usize = 12;

/// A Scoped Cipher is one which has been initialized for a specific keyset.
/// It can be used *only* to encrypt and decrypt data for that keyset.
///
/// # Lifetime
///
/// **Initialise once per process and hold an `Arc<ScopedCipher<_>>` for the
/// process lifetime.** Recreating a `ScopedCipher` on a periodic schedule
/// (e.g. "refresh every N rows" in a bulk-ingest loop) is an anti-pattern:
///
/// - `ScopedCipher` carries no TTL of its own — it's a thin wrapper around
///   the keyset's index key plus an `Arc` to the underlying ZeroKMS
///   client. The "scope" refers to the keyset binding, not a session
///   lifetime.
/// - Auth tokens used by the underlying client auto-refresh through
///   `stack_auth::AutoRefresh` — preemptively, 90 s before expiry, with
///   in-flight refreshes shared across concurrent callers via a mutex.
///   Recreating `ScopedCipher` doesn't trigger or improve this; it only
///   discards the cached token and forces an immediate re-auth.
/// - Recreating the cipher discards the warm `reqwest::Client` connection
///   pool. The next call after re-init pays a cold-start tax (DNS + TCP +
///   TLS handshake) that can blow past the request timeout under variable
///   network conditions.
///
/// If you're working around symptoms that look like session expiry, check
/// the actual error chain — auth errors surface as
/// [`zerokms::Error::Auth`](crate::zerokms::Error::Auth), not as
/// `ZeroKMS error 'Unexpected error'`. The latter is almost always a
/// transport-level timeout (`SendRequest`) that periodic re-init *makes
/// worse*, not better.
///
/// See the "Tuning for bulk ingest" section on
/// [`ZeroKMSBuilder`](crate::zerokms::ZeroKMSBuilder) for the timeout /
/// pool / batch knobs that *do* help long-running encrypt loops.
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct ScopedCipher<C> {
    #[zeroize(skip)]
    client: Arc<ZeroKMSWithClientKey<C>>,
    #[zeroize(skip)]
    keyset_id: Uuid,
    index_key: IndexKey,
}

impl<C> Debug for ScopedCipher<C> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ScopedCipher")
            .field("keyset_id", &self.keyset_id)
            .finish()
    }
}

impl<C> ScopedCipher<C> {
    /// Initialize a new ScopedCipher for the default keyset_id.
    pub async fn init_default(client: Arc<ZeroKMSWithClientKey<C>>) -> Result<Self, zerokms::Error>
    where
        C: AuthStrategyBounds,
        for<'b> &'b C: AuthStrategy,
    {
        Self::init(client, None).await
    }

    /// Initialize a new ScopedCipher for the given keyset_id.
    /// If the keyset_id is None, the ScopedCipher will be initialized with the default keyset_id for the client.
    pub async fn init(
        client: Arc<ZeroKMSWithClientKey<C>>,
        keyset_id: Option<IdentifiedBy>,
    ) -> Result<Self, zerokms::Error>
    where
        C: AuthStrategyBounds,
        for<'b> &'b C: AuthStrategy,
    {
        client
            .load_keyset_index_key(keyset_id)
            .await
            .map(|(keyset_id, index_key)| Self {
                client,
                keyset_id,
                index_key,
            })
    }

    /// This value is used for term index keys and "encrypted" partition / sort keys
    pub fn mac<const N: usize>(&self, value: &str, prefix: Option<&str>) -> [u8; N] {
        // TODO: Make sure this uses an unambiguous encoding scheme
        // see https://linear.app/cipherstash/issue/BUG-95/wip-ensure-unambiguous-coding-for-mac-prefixes
        let mut hasher = blake3::Hasher::new_keyed(self.index_key.key());
        if let Some(prefix) = prefix {
            hasher.update(prefix.as_bytes());
        }
        hasher.update(value.as_bytes());
        let mut result = [0; N];
        hasher.finalize_xof().fill(&mut result);
        hasher.zeroize();
        result
    }

    pub fn compound_index<I>(
        &self,
        index: I,
        plaintext: ComposablePlaintext,
        info: String,
    ) -> Result<IndexTerm, EncryptionError>
    where
        I: ComposableIndex + Send,
    {
        let index = CompoundIndex::new(index);
        let accumulator = Accumulator::from_salt(info);

        let term = index
            .compose_index(&self.index_key, plaintext, accumulator)?
            .truncate(DEFAULT_INDEX_TERM_SIZE)?;

        Ok(term.into())
    }

    pub fn compound_query<I>(
        &self,
        index: I,
        input: ComposablePlaintext,
        info: String,
    ) -> Result<IndexTerm, EncryptionError>
    where
        I: ComposableIndex + Send,
    {
        let index = CompoundIndex::new(index);
        let accumulator = Accumulator::from_salt(info);

        let term = index
            .compose_query(&self.index_key, input, accumulator)?
            .exactly_one()?
            .truncate(DEFAULT_INDEX_TERM_SIZE)?;

        Ok(term.try_into()?)
    }

    /// Encrypt a stream of [`EncryptPayload`] and return them as an [`EncryptedRecord`].
    /// This function wraps the [`ZeroKMSWithClientKey::encrypt`] function but with the keyset_id set.
    pub async fn encrypt(
        &self,
        payloads: impl IntoIterator<Item = EncryptPayload<'_>>,
    ) -> Result<Vec<EncryptedRecord>, Error>
    where
        C: AuthStrategyBounds,
        for<'b> &'b C: AuthStrategy,
    {
        self.client.encrypt(payloads, Some(self.keyset_id)).await
    }

    /// Decrypt a stream of encrypted values (of type `D` where `D: Decryptable`) and return the raw decrypted binary
    /// blob.
    pub async fn decrypt<'a, D>(
        &self,
        payloads: impl IntoIterator<Item = D>,
        opts: &DecryptOptions<'a>,
    ) -> Result<Vec<Vec<u8>>, Error>
    where
        D: Decryptable,
        C: AuthStrategyBounds,
        for<'b> &'b C: AuthStrategy,
    {
        self.client
            .decrypt(
                payloads,
                opts.keyset_id.or(Some(self.keyset_id)),
                opts.unverified_context.as_deref(),
            )
            .await
    }

    pub async fn decrypt_fallible<'a, D>(
        &self,
        payloads: impl IntoIterator<Item = D>,
        opts: &DecryptOptions<'a>,
    ) -> Result<Vec<Result<Vec<u8>, RecordDecryptError>>, Error>
    where
        D: Decryptable,
        C: AuthStrategyBounds,
        for<'b> &'b C: AuthStrategy,
    {
        let unverified_context = opts.unverified_context.as_deref().map(Cow::Borrowed);
        self.client
            .decrypt_fallible(payloads, unverified_context)
            .await
    }

    // TODO: Dan to make this pub(crate) again
    pub fn index_key(&self) -> &IndexKey {
        &self.index_key
    }

    pub fn keyset_id(&self) -> Uuid {
        self.keyset_id
    }

    /// Generate data keys for a stream of [`GenerateKeyPayload`] and return them as a [`DataKeyWithTag`].
    /// Scoped to the keyset_id of this [`ScopedCipher`].
    #[allow(dead_code)]
    pub(crate) async fn generate_data_keys<'a>(
        &self,
        payloads: impl IntoIterator<Item = GenerateKeyPayload<'_>>,
        unverified_context: Option<Cow<'a, UnverifiedContext>>,
    ) -> Result<Vec<DataKeyWithTag>, Error>
    where
        C: AuthStrategyBounds,
        for<'b> &'b C: AuthStrategy,
    {
        self.client
            .generate_data_keys(payloads, Some(self.keyset_id), unverified_context)
            .await
    }
}

#[derive(Debug, Default)]
pub struct DecryptOptions<'a> {
    pub keyset_id: Option<Uuid>,
    pub unverified_context: Option<Cow<'a, UnverifiedContext>>,
}
