lodestone-wire 0.0.2

use std::collections::BTreeMap;

use crate::{BULK_BATCH_ID_DOMAIN, BULK_BATCH_ID_FORMAT_VERSION, Digest, Identity, Resource};
const MAXIMUM_ITEM_ACTION_IDS: usize = 1_024;

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum BulkBatchIdError {
    InvalidCollectionIdentity,
    ZeroCollectionIdentity,
    InvalidIndexKind,
    InvalidIndexName,
    InvalidIndexIncarnation,
    ZeroIndexIncarnation,
    TooManyItemActionIds,
    InvalidItemActionId { position: usize },
    ZeroItemActionId { position: usize },
    DuplicateItemActionId { first: usize, duplicate: usize },
}

impl std::fmt::Display for BulkBatchIdError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(formatter, "invalid Lodestone API bulk batch identity input: {self:?}")
    }
}

impl std::error::Error for BulkBatchIdError {}

/// Derives the sole current Lodestone API batch identity used by document bulk
/// and retained-append bulk requests.
///
/// The BLAKE3-256 input is the versioned domain followed by this canonical
/// binary payload, in order: format version as `u16` big-endian, collection as raw
/// `u128` big-endian, index-name byte length as `u16` big-endian, exact UTF-8
/// index-name bytes, index incarnation as raw `u128` big-endian, item count as
/// `u32` big-endian, then each ordered action identity as 32 raw bytes.
pub fn bulk_batch_id(
    collection: &Identity,
    index: &Resource,
    item_action_ids: &[Digest],
) -> Result<Digest, BulkBatchIdError> {
    let payload = canonical_bulk_batch_payload_v1(collection, index, item_action_ids)?;
    let mut hasher = blake3::Hasher::new();
    hasher.update(BULK_BATCH_ID_DOMAIN);
    hasher.update(&payload);
    Ok(Digest {
        value: hasher.finalize().to_hex().to_string(),
    })
}

fn canonical_bulk_batch_payload_v1(
    collection: &Identity,
    index: &Resource,
    item_action_ids: &[Digest],
) -> Result<Vec<u8>, BulkBatchIdError> {
    let collection = decode_identity(&collection.value).ok_or(BulkBatchIdError::InvalidCollectionIdentity)?;
    if collection == 0 {
        return Err(BulkBatchIdError::ZeroCollectionIdentity);
    }
    if index.kind != "index" {
        return Err(BulkBatchIdError::InvalidIndexKind);
    }
    let index_name = index.name.as_bytes();
    if !(1..=255).contains(&index_name.len()) || index_name.iter().any(|byte| *byte <= 0x1f || *byte == 0x7f) {
        return Err(BulkBatchIdError::InvalidIndexName);
    }
    let incarnation = decode_identity(&index.incarnation.value).ok_or(BulkBatchIdError::InvalidIndexIncarnation)?;
    if incarnation == 0 {
        return Err(BulkBatchIdError::ZeroIndexIncarnation);
    }
    if item_action_ids.len() > MAXIMUM_ITEM_ACTION_IDS {
        return Err(BulkBatchIdError::TooManyItemActionIds);
    }

    let name_length = u16::try_from(index_name.len()).map_err(|_| BulkBatchIdError::InvalidIndexName)?;
    let item_count = u32::try_from(item_action_ids.len()).map_err(|_| BulkBatchIdError::TooManyItemActionIds)?;
    let mut payload = Vec::with_capacity(40 + index_name.len() + item_action_ids.len() * 32);
    payload.extend_from_slice(&BULK_BATCH_ID_FORMAT_VERSION.to_be_bytes());
    payload.extend_from_slice(&collection.to_be_bytes());
    payload.extend_from_slice(&name_length.to_be_bytes());
    payload.extend_from_slice(index_name);
    payload.extend_from_slice(&incarnation.to_be_bytes());
    payload.extend_from_slice(&item_count.to_be_bytes());

    let mut positions = BTreeMap::new();
    for (position, action) in item_action_ids.iter().enumerate() {
        let bytes = decode_digest(&action.value).ok_or(BulkBatchIdError::InvalidItemActionId { position })?;
        if bytes == [0; 32] {
            return Err(BulkBatchIdError::ZeroItemActionId { position });
        }
        if let Some(first) = positions.insert(bytes, position) {
            return Err(BulkBatchIdError::DuplicateItemActionId {
                first,
                duplicate: position,
            });
        }
        payload.extend_from_slice(&bytes);
    }
    Ok(payload)
}

fn decode_identity(value: &str) -> Option<u128> {
    if value.len() != 32 || !value.bytes().all(is_lower_hex) {
        return None;
    }
    u128::from_str_radix(value, 16).ok()
}

fn decode_digest(value: &str) -> Option<[u8; 32]> {
    if value.len() != 64 || !value.bytes().all(is_lower_hex) {
        return None;
    }
    let mut decoded = [0_u8; 32];
    for (output, pair) in decoded.iter_mut().zip(value.as_bytes().chunks_exact(2)) {
        *output = (lower_hex_value(pair[0])? << 4) | lower_hex_value(pair[1])?;
    }
    Some(decoded)
}

const fn is_lower_hex(value: u8) -> bool {
    matches!(value, b'0'..=b'9' | b'a'..=b'f')
}

const fn lower_hex_value(value: u8) -> Option<u8> {
    match value {
        b'0'..=b'9' => Some(value - b'0'),
        b'a'..=b'f' => Some(value - b'a' + 10),
        _ => None,
    }
}

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

    const VECTORS: &str = include_str!("../../../clients/native/fixtures/lodestone-bulk-batch-id.json");

    fn identity(value: u128) -> Identity {
        Identity {
            value: format!("{value:032x}"),
        }
    }

    fn digest(byte: u8) -> Digest {
        Digest {
            value: format!("{byte:02x}").repeat(32),
        }
    }

    fn index(name: &str, incarnation: u128) -> Resource {
        Resource {
            kind: "index".to_owned(),
            name: name.to_owned(),
            incarnation: identity(incarnation),
        }
    }

    #[test]
    fn bulk_batch_identity_is_deterministic_ordered_and_scope_bound() -> Result<(), Box<dyn std::error::Error>> {
        let collection = identity(1);
        let resource = index("audit-events", 2);
        let actions = [digest(0x11), digest(0x22)];
        let first = bulk_batch_id(&collection, &resource, &actions)?;
        assert_eq!(
            first.value,
            "25b1694f4d2883cc7c475d51ef6e9fc6ce1d42ef09a9283db30a4acd7ebd0f06"
        );
        let escaped = bulk_batch_id(
            &Identity {
                value: "0123456789abcdef0123456789abcdef".to_owned(),
            },
            &Resource {
                kind: "index".to_owned(),
                name: "quote\"\\雪".to_owned(),
                incarnation: Identity {
                    value: "fedcba9876543210fedcba9876543210".to_owned(),
                },
            },
            &[
                Digest {
                    value: format!("{}01", "00".repeat(31)),
                },
                digest(0xff),
            ],
        )?;
        assert_eq!(
            escaped.value,
            "619b49a49aa5a91437679522dbc272e3344118eea59f6ccd1fda7655b78d5f29"
        );
        assert_eq!(bulk_batch_id(&collection, &resource, &actions)?, first);
        assert_ne!(
            bulk_batch_id(&collection, &resource, &[actions[1].clone(), actions[0].clone()])?,
            first
        );
        assert_ne!(bulk_batch_id(&identity(3), &resource, &actions)?, first);
        assert_ne!(
            bulk_batch_id(&collection, &index("audit-events-2", 2), &actions)?,
            first
        );
        assert_ne!(bulk_batch_id(&collection, &index("audit-events", 3), &actions)?, first);
        Ok(())
    }

    #[test]
    fn bulk_batch_identity_rejects_duplicate_malformed_and_zero_actions() {
        let collection = identity(1);
        let resource = index("audit-events", 2);
        assert_eq!(
            bulk_batch_id(&collection, &resource, &[digest(0x11), digest(0x11)]),
            Err(BulkBatchIdError::DuplicateItemActionId { first: 0, duplicate: 1 })
        );
        assert_eq!(
            bulk_batch_id(&collection, &resource, &[Digest { value: "GG".repeat(32) }]),
            Err(BulkBatchIdError::InvalidItemActionId { position: 0 })
        );
        assert_eq!(
            bulk_batch_id(&collection, &resource, &[digest(0)]),
            Err(BulkBatchIdError::ZeroItemActionId { position: 0 })
        );
    }

    #[test]
    fn bulk_batch_identity_rejects_invalid_scope_and_over_bound_inputs() {
        let resource = index("audit-events", 2);
        assert_eq!(
            bulk_batch_id(&Identity { value: "01".to_owned() }, &resource, &[digest(0x11)],),
            Err(BulkBatchIdError::InvalidCollectionIdentity)
        );
        assert_eq!(
            bulk_batch_id(&identity(0), &resource, &[digest(0x11)]),
            Err(BulkBatchIdError::ZeroCollectionIdentity)
        );
        assert_eq!(
            bulk_batch_id(
                &identity(1),
                &Resource {
                    kind: "alias".to_owned(),
                    name: "audit-events".to_owned(),
                    incarnation: identity(2),
                },
                &[digest(0x11)],
            ),
            Err(BulkBatchIdError::InvalidIndexKind)
        );
        assert_eq!(
            bulk_batch_id(&identity(1), &index("", 2), &[digest(0x11)]),
            Err(BulkBatchIdError::InvalidIndexName)
        );
        assert_eq!(
            bulk_batch_id(&identity(1), &index("audit-events", 0), &[digest(0x11)]),
            Err(BulkBatchIdError::ZeroIndexIncarnation)
        );
        assert_eq!(
            bulk_batch_id(&identity(1), &resource, &vec![digest(0x11); 1_025]),
            Err(BulkBatchIdError::TooManyItemActionIds)
        );
    }

    #[test]
    fn fixed_language_neutral_vectors_match_payload_and_digest() -> Result<(), Box<dyn std::error::Error>> {
        let document: serde_json::Value = serde_json::from_str(VECTORS)?;
        assert_eq!(document["format_version"].as_u64(), Some(1));
        assert_eq!(
            document["batch_id_format_version"].as_u64(),
            Some(u64::from(BULK_BATCH_ID_FORMAT_VERSION))
        );
        assert_eq!(document["algorithm"].as_str(), Some("blake3-256"));
        assert_eq!(
            document["domain_separator_utf8"].as_str().map(str::as_bytes),
            Some(BULK_BATCH_ID_DOMAIN)
        );
        let vectors = document["vectors"].as_array().ok_or("bulk batch vectors are absent")?;
        for vector in vectors {
            let collection = Identity {
                value: vector["collection_identity"]
                    .as_str()
                    .ok_or("bulk batch vector collection is absent")?
                    .to_owned(),
            };
            let index: Resource = serde_json::from_value(vector["index"].clone())?;
            let action_values = vector["item_action_ids"]
                .as_array()
                .ok_or("bulk batch vector actions are absent")?;
            let actions = action_values
                .iter()
                .map(|value| {
                    value
                        .as_str()
                        .map(|value| Digest {
                            value: value.to_owned(),
                        })
                        .ok_or("bulk batch vector action is invalid")
                })
                .collect::<Result<Vec<_>, _>>()?;
            let expected_payload = decode_test_hex(
                vector["canonical_payload_hex"]
                    .as_str()
                    .ok_or("bulk batch vector payload is absent")?,
            )?;
            assert_eq!(
                canonical_bulk_batch_payload_v1(&collection, &index, &actions)?,
                expected_payload
            );
            assert_eq!(
                bulk_batch_id(&collection, &index, &actions)?.value,
                vector["batch_id"]
                    .as_str()
                    .ok_or("bulk batch vector result is absent")?
            );
        }
        Ok(())
    }

    fn decode_test_hex(value: &str) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
        if !value.len().is_multiple_of(2) {
            return Err("bulk batch vector payload has odd hex width".into());
        }
        value
            .as_bytes()
            .chunks_exact(2)
            .map(|pair| Ok(u8::from_str_radix(std::str::from_utf8(pair)?, 16)?))
            .collect()
    }
}