//! Domain-separated keys for the unified state tree.
//!
//! [`NamespaceId`] selects one of the three logical namespaces within the single JMT.
//! [`NamespaceKey`] pairs a namespace with raw key bytes and computes a 32-byte JMT
//! [`KeyHash`](jmt::KeyHash) where byte 0 is the namespace tag and bytes 1..32 are
//! `sha256(raw_key_bytes)[1..32]`.

use alloy_primitives::{Address, B256};
use jmt::KeyHash;
use sha2::{Digest, Sha256};

/// Namespace selector.
///
/// All three namespaces share a single unified Jellyfish Merkle Tree. Separation is
/// achieved by a 1-byte prefix at `leaf_key[0]` rather than separate trees.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(u8)]
pub enum NamespaceId {
    /// Encrypted identity data (KYC attestations, on-chain bindings).
    Identity = 0x01,
    /// Encrypted provider data (blacklists, allowlists).
    Confidential = 0x02,
    /// Policy client secrets (HPKE-encrypted WASM secrets).
    Secrets = 0x03,
}

impl NamespaceId {
    /// The raw byte representation of this namespace.
    pub fn byte(self) -> u8 {
        self as u8
    }

    /// Parse a namespace from its byte representation.
    pub fn from_byte(b: u8) -> Option<Self> {
        match b {
            0x01 => Some(Self::Identity),
            0x02 => Some(Self::Confidential),
            0x03 => Some(Self::Secrets),
            _ => None,
        }
    }
}

/// A key within the unified state tree, scoped to one namespace.
///
/// The JMT [`KeyHash`] is 32 bytes: byte 0 is the [`NamespaceId`] tag, bytes 1..32 are
/// `sha256(raw_key_bytes)[1..32]`. This pins key-space disjointness across namespaces
/// while the domain-separation string inside the SHA-256 preimage (constructed by the
/// caller) pins intra-namespace separation.
///
/// # Invariants
///
/// `raw` **must** be a domain-separated preimage per `docs/PRIVATE_DATA_STORAGE.md` §5
/// (e.g. `"identity_v1" || user_address || ...`). The type system does not enforce the
/// preimage structure defined in spec §5.1–§5.3, so callers that pass bare keys risk
/// cross-schema collisions. Registry adapters (NEWT-1015 / NEWT-1019 / NEWT-1028) wrap
/// these constructors with the correct prefixes downstream.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NamespaceKey {
    namespace: NamespaceId,
    raw: Vec<u8>,
}

impl NamespaceKey {
    /// Create a key in the Identity namespace.
    pub fn identity(key_bytes: impl AsRef<[u8]>) -> Self {
        Self {
            namespace: NamespaceId::Identity,
            raw: key_bytes.as_ref().to_vec(),
        }
    }

    /// Create a key in the Confidential namespace.
    pub fn confidential(key_bytes: impl AsRef<[u8]>) -> Self {
        Self {
            namespace: NamespaceId::Confidential,
            raw: key_bytes.as_ref().to_vec(),
        }
    }

    /// Create a key in the Secrets namespace.
    pub fn secrets(key_bytes: impl AsRef<[u8]>) -> Self {
        Self {
            namespace: NamespaceId::Secrets,
            raw: key_bytes.as_ref().to_vec(),
        }
    }

    /// Create a key with an explicit namespace and raw bytes.
    pub fn new(namespace: NamespaceId, key_bytes: impl AsRef<[u8]>) -> Self {
        Self {
            namespace,
            raw: key_bytes.as_ref().to_vec(),
        }
    }

    /// The namespace this key belongs to.
    pub fn namespace(&self) -> NamespaceId {
        self.namespace
    }

    /// The raw key bytes (domain-separated preimage).
    pub fn raw(&self) -> &[u8] {
        &self.raw
    }

    /// Compute the 32-byte JMT [`KeyHash`]: `[namespace_byte, sha256(raw)[1..32]]`.
    pub fn key_hash(&self) -> KeyHash {
        let digest = Sha256::digest(&self.raw);
        let mut result = [0u8; 32];
        result[0] = self.namespace.byte();
        result[1..].copy_from_slice(&digest[1..32]);
        KeyHash(result)
    }
}

/// Domain-separation prefix for identity leaf keys (see `docs/PRIVATE_DATA_STORAGE.md` § 5.1).
const IDENTITY_LEAF_SEP: &[u8] = b"identity/leaf_v1";
/// Domain-separation prefix for identity linkage keys (see `docs/PRIVATE_DATA_STORAGE.md` § 5.1).
const IDENTITY_LINK_SEP: &[u8] = b"identity/link_v1";
/// Domain-separation prefix for confidential data keys (see `docs/PRIVATE_DATA_STORAGE.md` § 5.2).
const CONFIDENTIAL_DATA_SEP: &[u8] = b"confidential/data_v1";
/// Domain-separation prefix for confidential grant keys (see `docs/PRIVATE_DATA_STORAGE.md` § 5.2).
const CONFIDENTIAL_GRANT_SEP: &[u8] = b"confidential/grant_v1";
/// Domain-separation prefix for secrets leaf keys (see `docs/PRIVATE_DATA_STORAGE.md` § 5.3).
const SECRETS_LEAF_SEP: &[u8] = b"secrets/leaf_v1";

/// Length of an Ethereum address in bytes.
const ADDRESS_LEN: usize = 20;
/// Length of a B256 (bytes32) in bytes.
const B256_LEN: usize = 32;

/// Decode a `confidential_grant` raw preimage back into its `(provider,
/// policy_client, domain_id)` triple, or return `None` if the bytes do not
/// match the canonical layout.
///
/// Used by the operator's `authorize_read` (PDS §9.2 Layer 2) to recognize
/// that an inbound `leaf_key` targets a grant-status leaf and to extract the
/// triple without exposing prefix bytes outside this crate.
pub fn extract_confidential_grant(raw: &[u8]) -> Option<(Address, Address, B256)> {
    let prefix = CONFIDENTIAL_GRANT_SEP;
    let expected_len = prefix.len() + ADDRESS_LEN + ADDRESS_LEN + B256_LEN;
    if raw.len() != expected_len {
        return None;
    }
    if !raw.starts_with(prefix) {
        return None;
    }
    let mut cursor = prefix.len();
    let provider = Address::from_slice(&raw[cursor..cursor + ADDRESS_LEN]);
    cursor += ADDRESS_LEN;
    let policy_client = Address::from_slice(&raw[cursor..cursor + ADDRESS_LEN]);
    cursor += ADDRESS_LEN;
    let domain_id = B256::from_slice(&raw[cursor..cursor + B256_LEN]);
    Some((provider, policy_client, domain_id))
}

/// Decode an `identity_linkage` raw preimage back into `(user, policy_client)`,
/// or return `None` if the bytes do not match the canonical layout.
///
/// Tuple order matches the [`NamespaceKey::identity_linkage`] preimage layout
/// (`b"identity/link_v1" || user[20] || policy_client[20]`), so this function
/// is the inverse of that builder. Used by the operator's `authorize_read` to
/// recognize that an inbound Identity-namespace `leaf_key` targets a linkage
/// record and to enforce `signer == policy_client` per PDS §9.2 Layer 2.
pub fn extract_identity_linkage(raw: &[u8]) -> Option<(Address, Address)> {
    let prefix = IDENTITY_LINK_SEP;
    let expected_len = prefix.len() + ADDRESS_LEN + ADDRESS_LEN;
    if raw.len() != expected_len {
        return None;
    }
    if !raw.starts_with(prefix) {
        return None;
    }
    let mut cursor = prefix.len();
    let user = Address::from_slice(&raw[cursor..cursor + ADDRESS_LEN]);
    cursor += ADDRESS_LEN;
    let policy_client = Address::from_slice(&raw[cursor..cursor + ADDRESS_LEN]);
    Some((user, policy_client))
}

/// Typed, domain-aware constructors that mirror the namespace schemas in
/// `docs/PRIVATE_DATA_STORAGE.md` § 5. Each constructor builds a fixed-width
/// preimage from typed primitives (`Address`, `B256`, `u32` BE) so that adjacent
/// fields cannot fuse — a `&str`-based API would let `domain_id="foo" || field_tag="bar"`
/// collide with `domain_id="foob" || field_tag="ar"`.
impl NamespaceKey {
    /// Build an identity-namespace leaf key for one `(user, domain_id, field_tag)` tuple.
    ///
    /// Preimage layout (72 bytes): `b"identity/leaf_v1" || user[20] || domain_id[32]
    /// || field_tag.to_be_bytes()[4]`. The resulting key hash is
    /// `[0x01, sha256(preimage)[1..32]]`.
    pub fn identity_leaf(user: Address, domain_id: B256, field_tag: u32) -> Self {
        const LEN: usize = IDENTITY_LEAF_SEP.len() + 20 + 32 + 4;
        let mut raw = [0u8; LEN];
        let (sep, rest) = raw.split_at_mut(IDENTITY_LEAF_SEP.len());
        sep.copy_from_slice(IDENTITY_LEAF_SEP);
        let (user_slot, rest) = rest.split_at_mut(20);
        user_slot.copy_from_slice(user.as_slice());
        let (domain_slot, tag_slot) = rest.split_at_mut(32);
        domain_slot.copy_from_slice(domain_id.as_slice());
        tag_slot.copy_from_slice(&field_tag.to_be_bytes());
        Self::new(NamespaceId::Identity, raw)
    }

    /// Build an identity-namespace linkage key for one `(user, policy_client)` pair.
    ///
    /// Preimage layout (56 bytes): `b"identity/link_v1" || user[20] || policy_client[20]`.
    pub fn identity_linkage(user: Address, policy_client: Address) -> Self {
        const LEN: usize = IDENTITY_LINK_SEP.len() + 20 + 20;
        let mut raw = [0u8; LEN];
        let (sep, rest) = raw.split_at_mut(IDENTITY_LINK_SEP.len());
        sep.copy_from_slice(IDENTITY_LINK_SEP);
        let (user_slot, client_slot) = rest.split_at_mut(20);
        user_slot.copy_from_slice(user.as_slice());
        client_slot.copy_from_slice(policy_client.as_slice());
        Self::new(NamespaceId::Identity, raw)
    }

    /// Build a confidential-namespace data key for one `(provider, domain_id, data_ref_id)` tuple.
    ///
    /// Preimage layout (104 bytes): `b"confidential/data_v1" || provider[20] ||
    /// domain_id[32] || data_ref_id[32]`.
    pub fn confidential_data(provider: Address, domain_id: B256, data_ref_id: B256) -> Self {
        const LEN: usize = CONFIDENTIAL_DATA_SEP.len() + 20 + 32 + 32;
        let mut raw = [0u8; LEN];
        let (sep, rest) = raw.split_at_mut(CONFIDENTIAL_DATA_SEP.len());
        sep.copy_from_slice(CONFIDENTIAL_DATA_SEP);
        let (provider_slot, rest) = rest.split_at_mut(20);
        provider_slot.copy_from_slice(provider.as_slice());
        let (domain_slot, ref_slot) = rest.split_at_mut(32);
        domain_slot.copy_from_slice(domain_id.as_slice());
        ref_slot.copy_from_slice(data_ref_id.as_slice());
        Self::new(NamespaceId::Confidential, raw)
    }

    /// Build a confidential-namespace grant key for one `(provider, policy_client, domain_id)`
    /// tuple. Grants are domain-scoped, not record-scoped: one grant authorises a policy
    /// client to read every data record the provider stores under that domain.
    ///
    /// Preimage layout (93 bytes): `b"confidential/grant_v1" || provider[20] ||
    /// policy_client[20] || domain_id[32]`.
    pub fn confidential_grant(provider: Address, policy_client: Address, domain_id: B256) -> Self {
        const LEN: usize = CONFIDENTIAL_GRANT_SEP.len() + 20 + 20 + 32;
        let mut raw = [0u8; LEN];
        let (sep, rest) = raw.split_at_mut(CONFIDENTIAL_GRANT_SEP.len());
        sep.copy_from_slice(CONFIDENTIAL_GRANT_SEP);
        let (provider_slot, rest) = rest.split_at_mut(20);
        provider_slot.copy_from_slice(provider.as_slice());
        let (client_slot, domain_slot) = rest.split_at_mut(20);
        client_slot.copy_from_slice(policy_client.as_slice());
        domain_slot.copy_from_slice(domain_id.as_slice());
        Self::new(NamespaceId::Confidential, raw)
    }

    /// Build a secrets-namespace leaf key for one `(policy_client, secret_id)` pair.
    ///
    /// Preimage layout (67 bytes): `b"secrets/leaf_v1" || policy_client[20] || secret_id[32]`.
    pub fn secrets_leaf(policy_client: Address, secret_id: B256) -> Self {
        const LEN: usize = SECRETS_LEAF_SEP.len() + 20 + 32;
        let mut raw = [0u8; LEN];
        let (sep, rest) = raw.split_at_mut(SECRETS_LEAF_SEP.len());
        sep.copy_from_slice(SECRETS_LEAF_SEP);
        let (client_slot, secret_slot) = rest.split_at_mut(20);
        client_slot.copy_from_slice(policy_client.as_slice());
        secret_slot.copy_from_slice(secret_id.as_slice());
        Self::new(NamespaceId::Secrets, raw)
    }
}

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

    #[test]
    fn same_bytes_different_hash_across_namespaces() {
        let bytes = [0xab; 32];
        let identity = NamespaceKey::identity(bytes);
        let confidential = NamespaceKey::confidential(bytes);
        let secrets = NamespaceKey::secrets(bytes);

        assert_ne!(identity.key_hash(), confidential.key_hash());
        assert_ne!(identity.key_hash(), secrets.key_hash());
        assert_ne!(confidential.key_hash(), secrets.key_hash());
    }

    #[test]
    fn same_namespace_same_bytes_deterministic() {
        let a = NamespaceKey::identity([1, 2, 3]);
        let b = NamespaceKey::identity([1, 2, 3]);
        assert_eq!(a.key_hash(), b.key_hash());
    }

    #[test]
    fn different_bytes_different_hashes() {
        let a = NamespaceKey::identity([1, 2, 3]);
        let b = NamespaceKey::identity([4, 5, 6]);
        assert_ne!(a.key_hash(), b.key_hash());
    }

    #[test]
    fn namespace_byte_at_position_zero() {
        let key = NamespaceKey::identity([0xff; 16]);
        assert_eq!(key.key_hash().0[0], 0x01);

        let key = NamespaceKey::confidential([0xff; 16]);
        assert_eq!(key.key_hash().0[0], 0x02);

        let key = NamespaceKey::secrets([0xff; 16]);
        assert_eq!(key.key_hash().0[0], 0x03);
    }

    #[test]
    fn from_byte_roundtrip() {
        for ns in [NamespaceId::Identity, NamespaceId::Confidential, NamespaceId::Secrets] {
            assert_eq!(NamespaceId::from_byte(ns.byte()), Some(ns));
        }
        assert_eq!(NamespaceId::from_byte(0x00), None);
        assert_eq!(NamespaceId::from_byte(0x04), None);
    }

    // ---------- Domain-aware constructor tests ----------

    /// Spec § 5.1 — identity leaf preimage: `"identity/leaf_v1" || user || domain || field_tag_be`.
    #[test]
    fn identity_leaf_matches_spec_preimage() {
        let user = Address::from([0xaa; 20]);
        let domain_id = B256::from([0xbb; 32]);
        let field_tag: u32 = 0x01020304;

        let mut preimage = Vec::new();
        preimage.extend_from_slice(b"identity/leaf_v1");
        preimage.extend_from_slice(user.as_slice());
        preimage.extend_from_slice(domain_id.as_slice());
        preimage.extend_from_slice(&field_tag.to_be_bytes());

        let mut expected = [0u8; 32];
        expected[0] = 0x01;
        expected[1..].copy_from_slice(&Sha256::digest(&preimage)[1..32]);

        let key = NamespaceKey::identity_leaf(user, domain_id, field_tag);
        assert_eq!(key.key_hash().0[0], 0x01, "identity namespace byte");
        assert_eq!(key.key_hash().0, expected);
    }

    /// Spec § 5.1 — identity linkage preimage: `"identity/link_v1" || user || policy_client`.
    #[test]
    fn identity_linkage_matches_spec_preimage() {
        let user = Address::from([0xaa; 20]);
        let policy_client = Address::from([0xcc; 20]);

        let mut preimage = Vec::new();
        preimage.extend_from_slice(b"identity/link_v1");
        preimage.extend_from_slice(user.as_slice());
        preimage.extend_from_slice(policy_client.as_slice());

        let mut expected = [0u8; 32];
        expected[0] = 0x01;
        expected[1..].copy_from_slice(&Sha256::digest(&preimage)[1..32]);

        let key = NamespaceKey::identity_linkage(user, policy_client);
        assert_eq!(key.key_hash().0[0], 0x01, "identity namespace byte");
        assert_eq!(key.key_hash().0, expected);
    }

    /// Spec § 5.2 — confidential data preimage:
    /// `"confidential/data_v1" || provider || domain || data_ref_id`.
    #[test]
    fn confidential_data_matches_spec_preimage() {
        let provider = Address::from([0x01; 20]);
        let domain_id = B256::from([0x02; 32]);
        let data_ref_id = B256::from([0x03; 32]);

        let mut preimage = Vec::new();
        preimage.extend_from_slice(b"confidential/data_v1");
        preimage.extend_from_slice(provider.as_slice());
        preimage.extend_from_slice(domain_id.as_slice());
        preimage.extend_from_slice(data_ref_id.as_slice());

        let mut expected = [0u8; 32];
        expected[0] = 0x02;
        expected[1..].copy_from_slice(&Sha256::digest(&preimage)[1..32]);

        let key = NamespaceKey::confidential_data(provider, domain_id, data_ref_id);
        assert_eq!(key.key_hash().0[0], 0x02, "confidential namespace byte");
        assert_eq!(key.key_hash().0, expected);
    }

    /// Spec § 5.2 — confidential grant preimage:
    /// `"confidential/grant_v1" || provider || policy_client || domain_id`.
    /// Grants are domain-scoped, not data-record-scoped.
    #[test]
    fn confidential_grant_matches_spec_preimage() {
        let provider = Address::from([0x01; 20]);
        let policy_client = Address::from([0x02; 20]);
        let domain_id = B256::from([0x03; 32]);

        let mut preimage = Vec::new();
        preimage.extend_from_slice(b"confidential/grant_v1");
        preimage.extend_from_slice(provider.as_slice());
        preimage.extend_from_slice(policy_client.as_slice());
        preimage.extend_from_slice(domain_id.as_slice());

        let mut expected = [0u8; 32];
        expected[0] = 0x02;
        expected[1..].copy_from_slice(&Sha256::digest(&preimage)[1..32]);

        let key = NamespaceKey::confidential_grant(provider, policy_client, domain_id);
        assert_eq!(key.key_hash().0[0], 0x02, "confidential namespace byte");
        assert_eq!(key.key_hash().0, expected);
    }

    /// Spec § 5.3 — secrets leaf preimage: `"secrets/leaf_v1" || policy_client || secret_id`.
    #[test]
    fn secrets_leaf_matches_spec_preimage() {
        let policy_client = Address::from([0xdd; 20]);
        let secret_id = B256::from([0xee; 32]);

        let mut preimage = Vec::new();
        preimage.extend_from_slice(b"secrets/leaf_v1");
        preimage.extend_from_slice(policy_client.as_slice());
        preimage.extend_from_slice(secret_id.as_slice());

        let mut expected = [0u8; 32];
        expected[0] = 0x03;
        expected[1..].copy_from_slice(&Sha256::digest(&preimage)[1..32]);

        let key = NamespaceKey::secrets_leaf(policy_client, secret_id);
        assert_eq!(key.key_hash().0[0], 0x03, "secrets namespace byte");
        assert_eq!(key.key_hash().0, expected);
    }

    /// Domain separation across builders within the identity namespace: leaf and
    /// linkage share `leaf_key[0] == 0x01` but their domain-sep prefixes differ,
    /// so the resulting key hashes must differ.
    #[test]
    fn identity_leaf_and_linkage_disjoint() {
        let addr = Address::from([0xaa; 20]);

        let leaf = NamespaceKey::identity_leaf(addr, B256::ZERO, 0);
        let link = NamespaceKey::identity_linkage(addr, addr);

        assert_eq!(leaf.key_hash().0[0], 0x01);
        assert_eq!(link.key_hash().0[0], 0x01);
        assert_ne!(leaf.key_hash(), link.key_hash());
    }

    /// Domain separation across builders within the confidential namespace: data
    /// and grant share `leaf_key[0] == 0x02` but the domain-sep prefixes differ.
    #[test]
    fn confidential_data_and_grant_disjoint() {
        let provider = Address::from([0x01; 20]);
        let other = Address::from([0x02; 20]);
        let id = B256::from([0x03; 32]);

        let data = NamespaceKey::confidential_data(provider, id, B256::from([0x02; 32]));
        let grant = NamespaceKey::confidential_grant(provider, other, id);

        assert_eq!(data.key_hash().0[0], 0x02);
        assert_eq!(grant.key_hash().0[0], 0x02);
        assert_ne!(data.key_hash(), grant.key_hash());
    }

    /// Argument-order separation: swapping `provider` and `policy_client` in
    /// `confidential_grant` must produce a different key. Defends against the
    /// "two providers grant each other" symmetric-input class of bug.
    #[test]
    fn confidential_grant_arg_order_matters() {
        let a = Address::from([0x01; 20]);
        let b = Address::from([0x02; 20]);
        let domain = B256::from([0xff; 32]);

        let g1 = NamespaceKey::confidential_grant(a, b, domain);
        let g2 = NamespaceKey::confidential_grant(b, a, domain);

        assert_ne!(g1.key_hash(), g2.key_hash());
    }

    /// Version-suffix separation: a synthetic `"identity/leaf_v2"` preimage with identical
    /// arguments produces a different key from the `_v1` constructor. Bumping the suffix
    /// is a hard fork of the key space, not a backwards-compatible change.
    #[test]
    fn identity_leaf_version_suffix_pinned() {
        let user = Address::from([0xaa; 20]);
        let domain_id = B256::from([0xbb; 32]);
        let field_tag: u32 = 0;

        let v1_key = NamespaceKey::identity_leaf(user, domain_id, field_tag);

        let mut v2_preimage = Vec::new();
        v2_preimage.extend_from_slice(b"identity/leaf_v2");
        v2_preimage.extend_from_slice(user.as_slice());
        v2_preimage.extend_from_slice(domain_id.as_slice());
        v2_preimage.extend_from_slice(&field_tag.to_be_bytes());

        let mut v2_hash = [0u8; 32];
        v2_hash[0] = 0x01;
        v2_hash[1..].copy_from_slice(&Sha256::digest(&v2_preimage)[1..32]);

        assert_ne!(v1_key.key_hash().0, v2_hash);
    }
}
