//! Cryptographic commits for permissioned Spaces repositories.

use atrium_crypto::keypair::Secp256k1Keypair;
use base64::{Engine as _, engine::general_purpose::STANDARD_NO_PAD};
use hkdf::Hkdf;
use hmac::{Hmac, Mac};
use rand::Rng as _;
use serde::{Deserialize, Serialize};
use serde_bytes::ByteBuf;
use sha2::{Digest, Sha256};

use crate::{
    error::{Error, ErrorContext as _, Result},
    space::domain::SpaceRecord,
};

const LANES: usize = 1024;
const STATE_BYTES: usize = LANES * 2;
const DOMAIN: &[u8] = b"atproto-space-v1";

pub(crate) fn sign(
    key: &Secp256k1Keypair,
    space: &str,
    author: &str,
    rev: &str,
    records: &[SpaceRecord],
) -> Result<SignedCommit> {
    let hash = set_hash_digest(records);
    let mut ikm = [0_u8; 32];
    rand::rng().fill(&mut ikm);
    let context = encode_context(space, author, rev, &ikm)?;
    let signature = key
        .sign(&context)
        .context(Error::Internal, "commit signing failed")?;

    let hkdf = Hkdf::<Sha256>::from_prk(&ikm).context(Error::Internal, "commit HKDF failed")?;
    let mut mac_key = [0_u8; 32];
    hkdf.expand(&context, &mut mac_key)
        .context(Error::Internal, "commit HKDF failed")?;
    let mut mac =
        Hmac::<Sha256>::new_from_slice(&mac_key).context(Error::Internal, "commit HMAC failed")?;
    mac.update(&hash);

    Ok(SignedCommit {
        ver: 1,
        hash: LexBytes::new(hash),
        ikm: LexBytes::new(ikm),
        sig: LexBytes::new(signature),
        mac: LexBytes::new(mac.finalize().into_bytes()),
        rev: rev.to_owned(),
    })
}

#[derive(Clone, Debug, Serialize)]
pub(crate) struct SignedCommit<Bytes = LexBytes> {
    pub(crate) ver: u8,
    pub(crate) hash: Bytes,
    pub(crate) ikm: Bytes,
    pub(crate) sig: Bytes,
    pub(crate) mac: Bytes,
    pub(crate) rev: String,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub(crate) struct LexBytes {
    #[serde(rename = "$bytes")]
    pub(crate) value: String,
}

impl LexBytes {
    pub(crate) fn new(bytes: impl AsRef<[u8]>) -> Self {
        Self {
            value: STANDARD_NO_PAD.encode(bytes.as_ref()),
        }
    }

    pub(crate) fn decode(&self) -> Result<Vec<u8>> {
        STANDARD_NO_PAD
            .decode(&self.value)
            .context(Error::Internal, "invalid signed commit byte encoding")
    }
}

impl SignedCommit {
    pub(crate) fn dag_cbor(&self) -> Result<Vec<u8>> {
        serde_ipld_dagcbor::to_vec(&SignedCommit::<ByteBuf> {
            ver: self.ver,
            hash: self.hash.decode()?.into(),
            ikm: self.ikm.decode()?.into(),
            sig: self.sig.decode()?.into(),
            mac: self.mac.decode()?.into(),
            rev: self.rev.clone(),
        })
        .context(Error::Internal, "could not encode signed commit")
    }
}

pub(crate) fn set_hash_digest(records: &[SpaceRecord]) -> [u8; 32] {
    let mut lanes = [0_u16; LANES];
    for record in records {
        let element = format!("{}/{}/{}", record.collection, record.rkey, record.cid);
        let mut expanded = [0_u8; STATE_BYTES];
        let mut reader = blake3::Hasher::new();
        reader.update(element.as_bytes());
        reader.finalize_xof().fill(&mut expanded);
        for (lane, bytes) in lanes.iter_mut().zip(expanded.chunks_exact(2)) {
            *lane = lane.wrapping_add(u16::from_le_bytes([bytes[0], bytes[1]]));
        }
    }
    Sha256::digest(lanes.map(u16::to_le_bytes).as_flattened()).into()
}

fn encode_context(space: &str, author: &str, rev: &str, ikm: &[u8]) -> Result<Vec<u8>> {
    let fields = [space.as_bytes(), author.as_bytes(), rev.as_bytes(), ikm];
    let mut context = Vec::with_capacity(
        DOMAIN.len() + fields.iter().map(|field| field.len() + 2).sum::<usize>(),
    );
    context.extend_from_slice(DOMAIN);
    for field in fields {
        let length = u16::try_from(field.len()).map_err(|_| {
            Error::Internal("commit context field exceeds uint16 length prefix".to_owned())
        })?;
        context.extend_from_slice(&length.to_be_bytes());
        context.extend_from_slice(field);
    }
    Ok(context)
}

#[cfg(test)]
mod tests {
    use atrium_crypto::{keypair::Secp256k1Keypair, verify::verify_signature};
    use base64::{Engine as _, engine::general_purpose::STANDARD_NO_PAD};
    use cid::Cid;
    use hmac::{Hmac, Mac};
    use serde_json::json;
    use sha2::{Digest as _, Sha256};

    use crate::space::domain::SpaceRecord;

    use super::{LexBytes, SignedCommit, encode_context, set_hash_digest, sign};

    #[test]
    fn commit_json_and_cbor_keep_the_same_fields_and_native_byte_encoding() {
        let mut commit = SignedCommit {
            ver: 1,
            hash: LexBytes::new([1]),
            ikm: LexBytes::new([2]),
            sig: LexBytes::new([3]),
            mac: LexBytes::new([4]),
            rev: "rev".to_owned(),
        };
        assert_eq!(
            serde_json::to_value(&commit).unwrap(),
            json!({
                "ver": 1, "hash": {"$bytes": "AQ"}, "ikm": {"$bytes": "Ag"},
                "sig": {"$bytes": "Aw"}, "mac": {"$bytes": "BA"}, "rev": "rev",
            })
        );
        // Canonical map-key order: ikm, mac, rev, sig, ver, hash.
        assert_eq!(commit.dag_cbor().unwrap(),
            b"\xa6\x63ikm\x41\x02\x63mac\x41\x04\x63rev\x63rev\x63sig\x41\x03\x63ver\x01\x64hash\x41\x01");

        commit.hash.value = "invalid!".to_owned();
        assert_eq!(
            commit.dag_cbor().unwrap_err().to_string(),
            commit.hash.decode().unwrap_err().to_string()
        );
    }

    fn record(collection: &str, rkey: &str, cid: &str) -> SpaceRecord {
        SpaceRecord {
            space: "at://did:example:owner/space/com.example.forum/default".to_owned(),
            did: "did:example:alice".to_owned(),
            collection: collection.to_owned(),
            rkey: rkey.to_owned(),
            cid: cid.parse::<Cid>().unwrap(),
            value: json!({"$type": collection}),
        }
    }

    #[test]
    fn set_hash_is_order_independent_and_path_sensitive() {
        let a = record(
            "com.example.post",
            "one",
            "bafyreidefdycgbfy3oglcb6ism3eqhyp5llsrpzxjsuac2gsy4mtrtx244",
        );
        let b = record(
            "com.example.post",
            "two",
            "bafyreidpw4cbv6gr4ukh33z23pvvrpr3wi4gnpmi4doamlsl3sa4rgri2a",
        );
        assert_eq!(
            set_hash_digest(&[a.clone(), b.clone()]),
            set_hash_digest(&[b, a.clone()])
        );
        let moved = record("com.example.post", "moved", &a.cid.to_string());
        assert_ne!(set_hash_digest(&[a]), set_hash_digest(&[moved]));
    }

    #[test]
    fn set_hash_retains_little_endian_lanes_and_wrapping_addition() {
        let record = record(
            "com.example.post",
            "é-one",
            "bafyreidefdycgbfy3oglcb6ism3eqhyp5llsrpzxjsuac2gsy4mtrtx244",
        );
        let element = format!("{}/{}/{}", record.collection, record.rkey, record.cid);
        let mut expanded = [0_u8; super::STATE_BYTES];
        blake3::Hasher::new()
            .update(element.as_bytes())
            .finalize_xof()
            .fill(&mut expanded);
        let lanes: Vec<_> = expanded
            .chunks_exact(2)
            .map(|bytes| u16::from_le_bytes([bytes[0], bytes[1]]))
            .collect();
        assert!(lanes.iter().any(|&lane| lane > u16::MAX / 2));
        assert!(lanes.iter().any(|&lane| lane != lane.swap_bytes()));

        for copies in [0_u16, 1, 2, 3, 257] {
            let mut state = [0_u8; super::STATE_BYTES];
            for (lane, bytes) in lanes.iter().zip(state.chunks_exact_mut(2)) {
                bytes.copy_from_slice(&lane.wrapping_mul(copies).to_le_bytes());
            }
            let expected: [u8; 32] = Sha256::digest(state).into();
            assert_eq!(
                set_hash_digest(&vec![record.clone(); usize::from(copies)]),
                expected,
                "{copies} copies"
            );
        }
    }

    #[test]
    fn each_signed_view_uses_fresh_binding_material() {
        let signer = Secp256k1Keypair::import(&[7; 32]).unwrap();
        let first = sign(
            &signer,
            "at://did:example:o/space/com.example.x/a",
            "did:example:a",
            "3kbcq3p7ad400",
            &[],
        )
        .unwrap();
        let second = sign(
            &signer,
            "at://did:example:o/space/com.example.x/a",
            "did:example:a",
            "3kbcq3p7ad400",
            &[],
        )
        .unwrap();
        assert_eq!(first.hash, second.hash);
        assert_ne!(first.ikm, second.ikm);
        assert_ne!(first.mac, second.mac);
        assert_ne!(first.sig, second.sig);
    }

    #[test]
    fn commit_signature_and_mac_verify_against_the_published_key() {
        let signer = Secp256k1Keypair::import(&[7; 32]).unwrap();
        let space = "at://did:example:o/space/com.example.x/a";
        let author = "did:example:a";
        let commit = sign(&signer, space, author, "3kbcq3p7ad400", &[]).unwrap();
        let ikm = STANDARD_NO_PAD.decode(&commit.ikm.value).unwrap();
        let context = encode_context(space, author, &commit.rev, &ikm).unwrap();
        verify_signature(
            &atrium_crypto::keypair::Did::did(&signer),
            &context,
            &STANDARD_NO_PAD.decode(&commit.sig.value).unwrap(),
        )
        .unwrap();

        let hkdf = hkdf::Hkdf::<Sha256>::from_prk(&ikm).unwrap();
        let mut key = [0_u8; 32];
        hkdf.expand(&context, &mut key).unwrap();
        let mut mac = Hmac::<Sha256>::new_from_slice(&key).unwrap();
        mac.update(&STANDARD_NO_PAD.decode(&commit.hash.value).unwrap());
        mac.verify_slice(&STANDARD_NO_PAD.decode(&commit.mac.value).unwrap())
            .unwrap();
    }
}