use super::mac_term::MacTerm;
use crate::encryption::{
    json_indexer::{
        prefix_mac::{PrefixMac, UpdatePrefixMac},
        ste_vec::encrypted_term::{EncryptedSteVecTermCompat, EncryptedSteVecTermStandard},
    },
    orderable::orderable_to_u64,
    EncryptionError, Plaintext,
};
use cllw_ore::{
    encrypt_ope_bits, encrypt_ore_bits, orderize_string, CllwOpeEncrypt, CllwOreEncrypt,
    Error as CllwError, Key as CllwKey, OpeCllw8VariableV1, OreCllw8VariableV1,
};
use orderable_bytes::ToOrderableBytes;
use serde_json::Value;
use zeroize::{Zeroize, ZeroizeOnDrop};

/// Tagged-plaintext carrier for STE-vec orderable terms.
///
/// Numeric and string plaintexts share a single variable-width OPE/ORE
/// ciphertext type. Disjointness between the two domains is enforced by a
/// single **type-tag bit prepended to the plaintext bit stream** (Number = 0,
/// String = 1) so that
///
/// 1. all numeric ciphertexts sort below all string ciphertexts at the
///    cryptographic layer, without an extra discriminator byte tacked onto the
///    ciphertext after the fact; and
/// 2. the tag bit itself flows through the CLLW PRF, contributing keystream
///    entropy — not just a free-floating prefix that an attacker can ignore.
///
/// Bit layout: `[tag bit][payload bits]`.
///
/// - `Number(n)` — 1 tag bit (`0`) + 64 payload bits of `n` MSB-first. `n` has
///   already been bit-twiddled by `orderable_to_u64` so plain bit order matches
///   the original signed/float plaintext order.
/// - `String(s)` — 1 tag bit (`1`) + the bytes of `orderize_string(s)` MSB-
///   first, followed by an 8-bit `0x00` terminator. `orderize_string` never
///   emits `0x00`, so the terminator unambiguously marks end-of-string, and
///   because `0x00` sorts below every content byte it keeps prefix order
///   correct (`"" < "a" < "ab"`).
#[derive(Zeroize, ZeroizeOnDrop)]
#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
pub(crate) enum OrderableTerm {
    Number(u64),
    String(String),
}

const TAG_NUMBER: u8 = 0;
const TAG_STRING: u8 = 1;
const STRING_TERMINATOR: u8 = 0x00;

impl OrderableTerm {
    /// Info string mixed into the term-side HMAC before deriving an OPE key.
    const OPE_INFO: &'static str = "OPE";
    /// Info string mixed into the term-side HMAC before deriving an ORE key.
    const ORE_INFO: &'static str = "ORE";

    /// Build the `[tag][payload]` plaintext bit stream described on the type
    /// docs. Each entry is a `0` or `1` byte — the form accepted by
    /// [`cllw_ore::encrypt_ope_bits`] / [`cllw_ore::encrypt_ore_bits`].
    fn tagged_bits(&self) -> Vec<u8> {
        match self {
            Self::Number(n) => {
                let mut bits = Vec::with_capacity(1 + u64::BITS as usize);
                bits.push(TAG_NUMBER);
                for i in (0..u64::BITS).rev() {
                    bits.push(((n >> i) & 1) as u8);
                }
                bits
            }
            Self::String(s) => {
                let bytes = orderize_string(s);
                let mut bits = Vec::with_capacity(1 + (bytes.len() + 1) * u8::BITS as usize);
                bits.push(TAG_STRING);
                for byte in bytes.bytes().chain(std::iter::once(STRING_TERMINATOR)) {
                    for i in (0..u8::BITS).rev() {
                        bits.push((byte >> i) & 1);
                    }
                }
                bits
            }
        }
    }

    /// Derive a term-specific CLLW key by mixing `info` into `macca` and
    /// finalizing to 32 bytes. The same construction is used for both OPE
    /// (Compat) and ORE (Standard) terms — only the info string differs, and
    /// keeping the derivation in one place gives a single audit point for the
    /// cryptographic key-derivation step.
    fn derive_key<M>(macca: &mut M, info: &'static str) -> CllwKey
    where
        M: PrefixMac + UpdatePrefixMac<&'static str>,
    {
        macca.update(info);
        CllwKey::from(macca.finalize_reset::<32>())
    }

    /// Derive an OPE key from `macca` and produce a Compat-mode encrypted
    /// term.
    pub(crate) fn build_compat<M>(
        &self,
        macca: &mut M,
    ) -> Result<EncryptedSteVecTermCompat, EncryptionError>
    where
        M: PrefixMac + UpdatePrefixMac<&'static str>,
    {
        let key = Self::derive_key(macca, Self::OPE_INFO);
        self.encrypt_ope_with_salt(&key, None)
            .map(EncryptedSteVecTermCompat::Ope)
            .map_err(EncryptionError::from)
    }

    /// Derive an ORE key from `macca` and produce a Standard-mode encrypted
    /// term.
    pub(crate) fn build_standard<M>(
        &self,
        macca: &mut M,
    ) -> Result<EncryptedSteVecTermStandard, EncryptionError>
    where
        M: PrefixMac + UpdatePrefixMac<&'static str>,
    {
        let key = Self::derive_key(macca, Self::ORE_INFO);
        self.encrypt_with_salt(&key, None)
            .map(EncryptedSteVecTermStandard::Ore)
            .map_err(EncryptionError::from)
    }
}

impl CllwOpeEncrypt for &OrderableTerm {
    type Output = OpeCllw8VariableV1;

    fn encrypt_ope_with_salt(
        self,
        key: &CllwKey,
        salt: Option<&[u8]>,
    ) -> Result<Self::Output, CllwError> {
        let bits = self.tagged_bits();
        let n = bits.len();
        Ok(OpeCllw8VariableV1::from_bytes(encrypt_ope_bits(
            bits, n, key, salt,
        )))
    }
}

impl CllwOreEncrypt for &OrderableTerm {
    type Output = OreCllw8VariableV1;

    fn encrypt_with_salt(
        self,
        key: &CllwKey,
        salt: Option<&[u8]>,
    ) -> Result<Self::Output, CllwError> {
        let bits = self.tagged_bits();
        Ok(OreCllw8VariableV1::from(encrypt_ore_bits(bits, key, salt)))
    }
}

impl TryFrom<&Plaintext> for OrderableTerm {
    type Error = EncryptionError;

    fn try_from(value: &Plaintext) -> Result<Self, Self::Error> {
        match value {
            Plaintext::Int(Some(n)) => Ok(Self::Number(*n as u64)),
            Plaintext::BigUInt(Some(n)) => Ok(Self::Number(*n)),
            Plaintext::BigInt(Some(n)) => Ok(Self::Number(orderable_to_u64(
                n.to_orderable_bytes().as_ref(),
            ))),
            Plaintext::Float(Some(f)) => Ok(Self::Number(orderable_to_u64(
                f.to_orderable_bytes().as_ref(),
            ))),
            Plaintext::SmallInt(Some(n)) => Ok(Self::Number(orderable_to_u64(
                n.to_orderable_bytes().as_ref(),
            ))),
            Plaintext::Text(Some(s)) => Ok(Self::String(s.to_string())),
            Plaintext::Int(None)
            | Plaintext::BigUInt(None)
            | Plaintext::BigInt(None)
            | Plaintext::Float(None)
            | Plaintext::SmallInt(None)
            | Plaintext::Text(None)
            | Plaintext::Json(None) => Err(null_value()),
            Plaintext::Boolean(x) => Err(invalid_type(x)),
            Plaintext::Decimal(x) => Err(invalid_type(x)),
            Plaintext::Json(Some(_)) => Err(EncryptionError::InvalidValue(
                "SteVecTerm only supports scalar values (strings or numbers). \
                 For JSON containment queries, use QueryOp::Default instead."
                    .to_string(),
            )),
            Plaintext::NaiveDate(x) => Err(invalid_type(x)),
            Plaintext::Timestamp(x) => Err(invalid_type(x)),
        }
    }
}

fn invalid_type<T>(_x: T) -> EncryptionError {
    EncryptionError::InvalidValue(format!(
        "Cannot convert value of type `{}` to orderable term",
        std::any::type_name::<T>()
    ))
}

fn null_value() -> EncryptionError {
    EncryptionError::InvalidValue("Cannot convert null value to orderable term".to_string())
}

#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
pub enum StePlaintextTerm {
    Mac(MacTerm),
    Orderable(OrderableTerm),
}

impl From<&Value> for StePlaintextTerm {
    fn from(value: &Value) -> Self {
        match value {
            Value::Bool(x) => StePlaintextTerm::Mac(MacTerm::Bool(*x)),
            Value::Null => StePlaintextTerm::Mac(MacTerm::Null),
            Value::Number(x) => {
                let number: f64 = x.as_f64().expect("conversion to f64 will succeed unless the serde_json feature 'arbitrary_precision' is enabled");
                StePlaintextTerm::Orderable(OrderableTerm::Number(orderable_to_u64(
                    number.to_orderable_bytes().as_ref(),
                )))
            }
            Value::String(x) => StePlaintextTerm::Orderable(OrderableTerm::String(x.to_string())),
            Value::Array(_) => StePlaintextTerm::Mac(MacTerm::Array),
            Value::Object(_) => StePlaintextTerm::Mac(MacTerm::Map),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cllw_ore::{CllwOpeEncrypt, CllwOreEncrypt, Key as CllwKey};
    use serde_json::{Number, Value};

    fn key() -> CllwKey {
        CllwKey::from([7u8; 32])
    }

    #[test]
    fn from_value_number() {
        assert_eq!(
            StePlaintextTerm::from(&Value::Number(Number::from_f64(10.0).unwrap())),
            StePlaintextTerm::Orderable(OrderableTerm::Number(orderable_to_u64(
                10.0_f64.to_orderable_bytes().as_ref()
            )))
        );
    }

    #[test]
    fn from_value_other_variants() {
        assert_eq!(
            StePlaintextTerm::from(&Value::Bool(true)),
            StePlaintextTerm::Mac(MacTerm::Bool(true))
        );
        assert_eq!(
            StePlaintextTerm::from(&Value::Null),
            StePlaintextTerm::Mac(MacTerm::Null)
        );
        assert_eq!(
            StePlaintextTerm::from(&Value::String("hello".into())),
            StePlaintextTerm::Orderable(OrderableTerm::String("hello".into()))
        );
        assert_eq!(
            StePlaintextTerm::from(&Value::Array(vec![])),
            StePlaintextTerm::Mac(MacTerm::Array)
        );
        assert_eq!(
            StePlaintextTerm::from(&Value::Object(Default::default())),
            StePlaintextTerm::Mac(MacTerm::Map)
        );
    }

    #[test]
    fn try_from_plaintext_text() {
        let p = Plaintext::Text(Some("test".into()));
        let t: OrderableTerm = (&p).try_into().unwrap();
        assert_eq!(t, OrderableTerm::String("test".into()));
    }

    #[test]
    fn try_from_plaintext_int() {
        let p = Plaintext::Int(Some(42));
        let t: OrderableTerm = (&p).try_into().unwrap();
        assert_eq!(t, OrderableTerm::Number(42));
    }

    #[test]
    fn try_from_plaintext_null_returns_error() {
        let p = Plaintext::Int(None);
        let err: EncryptionError = (&p).try_into().map(|_: OrderableTerm| ()).unwrap_err();
        assert!(err.to_string().contains("null value"));
    }

    #[test]
    fn try_from_plaintext_jsonb_returns_helpful_error() {
        let p = Plaintext::Json(Some(serde_json::json!({"role": "admin"})));
        let err: EncryptionError = (&p).try_into().map(|_: OrderableTerm| ()).unwrap_err();
        let msg = err.to_string();
        assert!(
            msg.contains("SteVecTerm only supports scalar values"),
            "{msg}"
        );
        assert!(msg.contains("QueryOp::Default"), "{msg}");
    }

    #[test]
    fn try_from_plaintext_jsonb_none_returns_null_error() {
        let p = Plaintext::Json(None);
        let err: EncryptionError = (&p).try_into().map(|_: OrderableTerm| ()).unwrap_err();
        assert!(err.to_string().contains("null value"));
    }

    #[test]
    fn try_from_plaintext_bool_returns_invalid_type() {
        let p = Plaintext::Boolean(Some(true));
        let err: EncryptionError = (&p).try_into().map(|_: OrderableTerm| ()).unwrap_err();
        assert!(err.to_string().contains("Cannot convert value of type"));
    }

    /// Numeric tag (0) sorts before string tag (1), so every Number ciphertext
    /// lex-compares below every String ciphertext under OPE — the property
    /// that lets us collapse `ocf`/`ocv` into a single `oc` field.
    #[test]
    fn ope_numbers_sort_below_strings() {
        let k = key();
        // Maximum-ish Number vs minimum-ish String. The numeric tag bit (0)
        // dominates over any payload, so the comparison is decided by tag.
        let big_number = (&OrderableTerm::Number(u64::MAX))
            .encrypt_ope_with_salt(&k, None)
            .unwrap();
        let empty_string = (&OrderableTerm::String(String::new()))
            .encrypt_ope_with_salt(&k, None)
            .unwrap();
        assert!(
            big_number.as_ref() < empty_string.as_ref(),
            "Number(MAX) ciphertext must sort below String(\"\") ciphertext"
        );
    }

    /// Same disjointness property at the ORE layer, using cllw-ore's ORE
    /// comparison (`OreCllw8VariableV1::Ord`).
    #[test]
    fn ore_numbers_sort_below_strings() {
        let k = key();
        let big_number = (&OrderableTerm::Number(u64::MAX))
            .encrypt_with_salt(&k, None)
            .unwrap();
        let empty_string = (&OrderableTerm::String(String::new()))
            .encrypt_with_salt(&k, None)
            .unwrap();
        assert!(
            big_number < empty_string,
            "Number(MAX) ORE ciphertext must sort below String(\"\") ORE ciphertext"
        );
    }

    /// OPE within the numeric domain preserves ordering.
    #[test]
    fn ope_numeric_ordering_preserved() {
        let k = key();
        let xs = [1u64, 2, 100, 1_000, u64::MAX / 2, u64::MAX];
        let cts: Vec<_> = xs
            .iter()
            .map(|n| {
                (&OrderableTerm::Number(*n))
                    .encrypt_ope_with_salt(&k, None)
                    .unwrap()
            })
            .collect();
        for w in cts.windows(2) {
            assert!(w[0].as_ref() < w[1].as_ref());
        }
    }

    /// OPE within the string domain preserves ordering.
    #[test]
    fn ope_string_ordering_preserved() {
        let k = key();
        let xs = ["a", "ab", "b", "ba", "z"];
        let cts: Vec<_> = xs
            .iter()
            .map(|s| {
                (&OrderableTerm::String((*s).into()))
                    .encrypt_ope_with_salt(&k, None)
                    .unwrap()
            })
            .collect();
        for w in cts.windows(2) {
            assert!(w[0].as_ref() < w[1].as_ref());
        }
    }
}
